Scaling Balanced Optimal Transport with the HELLO Solver

The HELLO (Hierarchical Edge Localization for Large-scale Optimal Transport) solver enables computing exact balanced optimal transport on million-point datasets by restricting the optimization to a sparse subgraph. Standard discrete optimal transport relies on linear programming, where the number of variables grows quadratically with the number of samples. This produces a dense matrix of every possible connection, quickly exhausting available memory for even moderate datasets. HELLO identifies and stores only the subset of edges likely to carry mass, which reduces memory complexity to linear. This allows researchers to scale to 1.28 million samples in 8,192 dimensions on a single H100 GPU with a 41.6 GiB memory footprint, delivering an order-of-magnitude reduction in runtime compared to strong baselines.

The Mechanism of Edge Localization

Optimal transport seeks a mapping that minimizes the total cost of moving mass from a source distribution to a target distribution. The HELLO solver treats this as a graph problem by building a recursive hierarchy of the input data points, typically using a tree structure like a cover tree or hierarchical clustering. The solver initializes dual potentials—the shadow prices associated with each point—at the coarsest level of the hierarchy by solving a small-scale instance of the problem using standard linear programming.

Once the coarse problem is solved, the system propagates these potentials down the hierarchy to guide the selection of relevant edges. The solver evaluates the condition for complementary slackness, which states that for an optimal transport plan, the cost of an edge between two points must be less than or equal to the sum of the dual potentials of its endpoints. If the cost exceeds this sum, moving mass across that edge is inefficient, and the solver discards it. If the cost is less than or equal to the sum, the edge is kept as a candidate for the active graph. This approach ensures the solver focuses its memory on the high-probability connections while ignoring the vast majority of near-zero-mass pairs that dominate dense formulations.

Because initial potential estimates are approximations, the solver must handle errors in its pruning through a process of violator insertion. After optimizing the sparse subgraph, the solver checks for violators—edges that technically satisfy the optimality condition but were initially pruned due to coarse estimation. Instead of a brute-force scan, the solver queries the hierarchy to identify regions where the dual potential constraints are violated. By traversing the tree and locally refining these regions, the algorithm identifies missing edges without ever materializing the full cost matrix. Iterative updates to the dual variables and expansion of the sparse graph ensure the solver converges to an exact solution.

Practical Implications for Builders

Use HELLO to replace entropy-regularized approximations like Sinkhorn divergence with exact balanced optimal transport on datasets larger than 100k samples. Many developers currently use regularizers to bypass the memory wall, but this creates numerical artifacts where mass is spread across many small, physically meaningless connections. HELLO allows you to achieve the true, sparse optimal transport solution without these approximations. It serves as an engine for specialized variations such as Gromov-Wasserstein distance, unbalanced optimal transport, and OT-based Flow Matching.

Consider an alignment task matching 500,000 embeddings from two language model checkpoints. A standard solver would attempt to allocate space for 250 billion edge weights, exceeding the capacity of most hardware. HELLO constructs a hierarchy to identify that most point pairs are distant enough that their transport mass is effectively zero, concentrating memory only on high-probability connections. This reduces computation to a single-GPU task. Because the solver relies on a hierarchy rather than a flat dense matrix, it handles general pairwise costs effectively. The main unknown for implementers remains how the choice of hierarchy construction parameters—such as the branching factor of the tree—impacts convergence speed on highly irregular data distributions.