FlashSinkhorn 2: Block-Sparse Entropic Optimal Transport

1 INTRODUCTION

• At 10^5–10^8 points per measure, dense-kernel storage exceeds GPU memory beyond about 10^5 points, while streaming solvers still perform O(nm) work per update.The cited applications include generative modeling, cell dynamics, computational biology, medical imaging, and cosmology.Source: §1 · p. 1
• FS2 solves squared-Euclidean EOT in d ≤3 by lifting a centroid solution, checking sampled rows and columns, and using sparse point updates if the lift remains inaccurate.Further centroid updates can replace most point-level updates before the sparse fine stage removes centroid error.Source: §1 · p. 2
• A block-mass threshold bounds each omitted tile’s contribution to every row and column, while Morton-ordered blocks enable fused tensor-core execution and reduce fine-stage work to O(nm/B^2).The fine stage stores block-pair indices rather than kernel entries.Source: §1 · p. 2
• On one A100, FS2 solved discrete EOT between two 1.34 × 10^8-particle measures to marginal residual below 0.01 in under 2.5 hours, which the authors call the largest such solve within hours.Source: §1 · p. 2

2.1 ENTROPIC OT AND FLASHSINKHORN

• Entropic OT minimizes transport cost plus an ε-weighted KL regularizer, and log-domain Sinkhorn alternates updates of source and target dual potentials.FlashSinkhorn computes the updates as tiled reductions, using tensor-core matrix products while keeping O((n + m)d) working memory rather than materializing the dense score matrix.Source: §2.1 · p. 2

2.2 MULTISCALE SINKHORN AND RELATED APPROACHES

• At blur √ε, the plan’s kernel factor falls to 1/e; when the blur spans a few nearest-neighbour spacings, most pairs carry negligible mass in locally uniform clouds, motivating screening in d ≤3.The number of points in a blur-sized ball grows as the d-th power of blur over spacing.Source: §2.2 · p. 3
• Related multiscale methods combine coarsening with truncation, but GeomLoss lifts centroid potentials without checking them and its centroid-cost mask has no bound on dropped mass.Other reviewed approaches include optimization variants, low-rank methods, domain decomposition, and hierarchical refinement, which can produce a different output from an EOT plan at prescribed residual.Source: §2.2 · p. 3, §2.3 · p. 3

3 FLASHSINKHORN 2: ALGORITHM AND ANALYSIS

• Morton sorting makes spatially compact cells and consecutive-point blocks contiguous, supporting coalesced GPU access; the solver selects block and cell sizes to fit summary state within the L2 budget.The automatic policy chooses B from 64, 128, or 256 and keeps at most twice as many cells as blocks.Source: §3 · pp. 4–5

3.1 COARSE SINKHORN UPDATES

• The coarse stage anneals on cell records, lifts potentials to points, and accepts a lift only when sampled row and column residual estimates plus three standard errors fall below τ.Exact residual computation costs O(nm); sampled row sums cost O(md), and the reported residuals are independently evaluated at returned potentials.Source: §3.1 · p. 5
• If a lift fails, additional centroid updates shrink coarse marginal error until stalling or 2048 updates; the remaining centroid-replacement error requires fine point updates.The lift-error bound separates coarse marginal error, which vanishes at coarse convergence, from centroid error, which further coarse updates cannot remove.Source: §3.1 · p. 5

3.2 SPARSE SINKHORN UPDATES

• Fine updates screen block pairs using block-level potential maxima and boxes, then verify candidate pairs by evaluating their exact mass over B^2 point pairs.Admitted tiles are stored in CSR form; blocks with no admitted tile use a dense update, while other blocks stay sparse.Source: §3.2 · p. 6
• The sparse half-steps run on the same mask, and the omitted-mass bound controls both marginal residual and the difference between sparse and full updates.The mask also bounds the error added to the barycentric projection used for gradients; sparse updates become cheaper than dense ones when the admitted tile count falls well below all block pairs.Source: §3.2 · p. 6
• A block-mass threshold bounds each omitted tile’s contribution, and the automatic δ_rel = τ/(4 max(NB, MB)) policy limits total omitted mass to τ/4 across rows or columns.Because the mask is certified at the potentials used to build it, FS2 rebuilds it at every fine update.Source: §3.2 · p. 6

4.1 BENCHMARKS

• FS2 meets τ on all 32 synthetic problems, while GeomLoss meets it on 10; where both succeed, FS2 is 74–645× faster.On clustered problems, GeomLoss misses τ at every tested size and regularization; the reported speed range is for the ten problems where both solvers meet τ.Source: Setup · pp. 7–8, Figure 2 · p. 7
• On clustered n=m=220 at ε=0.001, FS2 uses 2196 coarse updates and 15 fine updates; dense fine execution is 8.00× slower.Without coarse updates, the solve misses τ at the fine-update cap; stopping coarse updates after 404 instead requires 1115 fine updates and is 51.6× slower.Source: Setup · p. 8, Table 2 · p. 8
• FS2’s mask raises peak allocation to 8–10× either baseline, with 24 bytes of index storage per screened or admitted tile.At n=223, using B=256 instead of 128 lowers peak memory from 16.2 to 6.9 GiB for slightly more time.Source: Table 3 · p. 8, Setup · p. 9
• Across all 32 inputs, FS2’s surrogate-gradient steps lower the grouped objective at every checkpoint; at n=221, it is 25–61× faster end to end on heterogeneous data, with tied final scores.On clustered data, FS2 reaches a lower final grouped objective at both tested ε values and is faster at ε=0.001.Source: Figure 2 · p. 7, Setup · p. 9

4.2 QUIJOTE PARTICLE TRANSPORT

• At τ=0.01, FS2 solves the full Quijote box on one A100 in under 2.5 hours with full marginal residual below 0.01.From 224 to 227 particles, runtime grows 9.2-fold for an 8-fold increase in n, while peak memory stays within the 80 GiB card.Source: Problem and data · pp. 9–10, Figure 3 · p. 10
• On 4096 particles sampled from the 227 solution at τ=0.01, barycentric-projection displacements have RMSE 35.7% below a zero-displacement predictor.Particles whose reference trajectories cross a periodic face have larger errors than the rest.Source: Problem and data · p. 10, Figure 3 · p. 10

5 CONCLUSION AND LIMITATIONS

• FS2 reaches the target residual on all 32 synthetic problems and solves a 1.34 × 10^8-particle EOT problem on one A100 to below 0.01 residual in under 2.5 hours.The reported blur equals one lattice spacing.Source: §5 · p. 10
• The method targets squared-Euclidean cost in d ≤3, where blur-sized neighborhoods contain few points and screening pays off.Source: §5 · p. 10
• Reaching 10^10 cosmology points requires multiple GPUs because one measure’s FP32 coordinates alone occupy 112 GiB, exceeding an A100’s 80 GiB.The box screen also tests every block pair, so this scale requires a sub-quadratic pair search.Source: §5 · p. 10

AI USE STATEMENT

• Generative AI tools assisted with code, experiments, figures, and editing; the authors checked every reported number against experiment records and accept responsibility for the content.Source: AI USE STATEMENT · p. 10

A.1 PROOF OF PROPOSITION 1

• At coarse convergence, the centroid-replacement term can remain positive: collapsing two-point clouds at −h and h into one cell leaves nonzero marginal residual.In the example, the optimal coarse potentials are zero, yet the lifted potentials produce residual cosh(2h^2/ε) − 1.Source: §A.1 · p. 14
• On 32 grid problems, the bound rejects every lift despite 20 lifted returns having residuals at most 0.0041; fine updates are needed in some cases.Extra coarse updates lower lifted residuals on clustered problems at ε ≥0.002, but residuals stay flat on specified heterogeneous and clustered cases at ε = 0.001.Source: §A.1 · p. 14
• Rounding aggregated coordinates once to TF32 changes the potentials by about 1.6 × 10^-4 and the lift residual by 0.5%.Source: §A.1 · p. 14

A.2 COORDINATE GRADIENTS FROM COARSE POTENTIALS (PROPOSITIONS 3 AND 4)

• The coordinate-gradient construction holds target centroids, masses, and coarse potentials fixed, so gradients do not pass through those quantities.At the original coordinates, differentiating the lift gives ∇zi eL(Z) = 2ai(zi − eTi(Z)), the source-coordinate gradient formula.Source: §A.2 · p. 14
• The surrogate-gradient discrepancy includes both point-to-centroid quantization and differences in conditional weights, including coarse-solve error.Source: §A.2 · pp. 14–15
• For a β-Lipschitz gradient, a direction with positive alignment c decreases the objective when 0 < ηopt < 2∥G∥F c/β.Source: §A.2 · p. 15

• When a tile’s upper bound mIJ is at most ρIJ, its omitted contribution is bounded by δrel ai for each row and δrel bj for each column.Source: §A.3 · p. 15
• The screen is invariant to shifting f by κ and g by −κ, which leaves the plan unchanged.Source: §A.3 · p. 16
• With non-uniform weights, FS2 meets τ in all four cases and the per-point bound of Theorem 1 holds everywhere.Source: §A.3 · p. 16
• The mass-dependent rule admits 0.3–0.5% fewer block pairs under one weighting and 2.9–5.3% fewer under a rank−1/2 power law than a scalar rule, with the same fine updates.Source: §A.3 · p. 16

A.4 PROOF OF PROPOSITION 2

• The middle projection-error bound can be much smaller than the diameter bound when omitted mass lies near the masked projection, although the diameter bound is sharp.Source: §A.4 · p. 16
• At returned potentials, no row exceeds either masking bound, while the looser mask raises the largest measured error 50–70-fold.The FP32 sparse-gradient kernel differs from the FP64 masked gradient by at most 1.0 × 10^-4 in relative ℓ2 on clustered data and 1.0 × 10^-2 on heterogeneous data.Source: §A.4 · p. 16

A.5 IO MODEL (PROPOSITION 5)

• For sparse traversal, each neighbour-record miss moves Θ(B(d+1)) scalars, and the two miss counts lie between active neighbour blocks and admitted pairs.This IO count assumes source-block queries and accumulators stay resident, costs are computed online, indices are read once, and caches start cold each half-step.Source: §A.5 · p. 17
• Sparse traversal also moves Θ((n + m)(d + 1)) scalars for source-block queries and outputs, plus Θ(NB + |A|) for traversal indices.Source: §A.5 · p. 17

B.1 BLOCK SIZE AND COARSE CELLS

• The block-size policy selects B from {64, 128, 256} so block summaries fit within 0.8 times L2 cache; empirical cell limits favor finer partitions as memory allows.The summary memory estimate is 32(NB + MB) bytes for d=3, and the cell-count budget can increase memory traffic.Source: §B.1 · p. 17

B.2 FORWARD SOLVER AND STOPPING RULES

• Coarse solves anneal on Morton-prefix cells, lift potentials, and continue in batches when sampled checks reject the lift; coarse stopping uses mean errors and does not bound the maximum cell error.Checks use 1,024 sampled indices per side, may expand to 8,192, and can accept potentials whose full residual exceeds τ because samples are reused.Source: §B.2 · p. 18

B.3 MASK CONSTRUCTION AND SPARSE HALF-STEPS

• Fine updates rebuild a shared sparse mask, screen block pairs in two passes, and stream admitted tiles through CSR-based log-sum-exp updates.Mask construction scans the full block grid; each tile test costs about 6d + 3 operations, while the cost matrix is never formed.Source: §B.3 · pp. 18–19

B.4 OPERATION COUNT

• The operation bound accounts for coarse updates, lifts, checks, dense repairs, and sparse fine work; memory stores candidate and admitted indices only when their allocation limits fit.No score matrix is stored, but the selector’s coarse-count bounds remain quadratic in n at fixed B.Source: §B.4 · pp. 19–20

C.1 SETUP

• The benchmark uses three-dimensional uniform-weight inputs and evaluates full residuals against τ = 0.005, with warmed complete-call timings and separate peak-memory measurements.The grid uses clustered and heterogeneous clouds; FS2 uses one fixed calibration without per-problem tuning.Source: §C.1 · p. 20

C.2 GEOMLOSS MULTISCALE WORKFLOW

• GeomLoss switches from centroid resolutions to points according to its blur schedule, lifts coarse potentials, and builds a mask at the last coarse level when fine updates follow.It reuses the resulting KeOps index ranges for fine updates; if the schedule does not reach an eligible blur, it lifts and takes no fine update.Source: §C.2 · pp. 20–21
• At FS2’s clustered-data potentials, GeomLoss’s mask omits more than half the plan mass in 964 of 2^20 rows and nearly all in 12, while FS2 omits at most 6.7 × 10^-6 per row and keeps half as many point pairs.At GeomLoss’s own potentials its mask omits 9.8 × 10^-6 of mass, only 0.02% of its residual 0.054; on heterogeneous data the maximum omitted row mass is 7 × 10^-18 for GeomLoss and 2.9 × 10^-5 for FS2.Source: §C.2 · p. 21

C.3 SPEED

• The speed grid covers 32 inputs, and the target-ε continuation comparisons are specific to two n = m = 221 problems.FS2 uses a fixed kernel calibration with no per-problem tuning.Source: §C.3 · p. 21
• On heterogeneous ε = 0.01, all six GeomLoss settings meet τ; default cluster scale takes 57.8 s, 212 times FS2, while matched cells take 0.39 s, 1.4 times FS2.The two scaling-0.9 runs differ only in cluster scale, which the passage identifies as the source of the default-setting gap on this problem.Source: §C.3 · p. 21
• On clustered ε = 0.001, GeomLoss needs 118 extra target-ε entries and 421 s to meet τ, 27.3 times FS2’s 15.4 s; the grid setting never meets τ at scored K.The comparison gives GeomLoss an oracle stop without charging for residual tests, whereas FS2’s times include sampled checks.Source: §C.3 · p. 21

C.4 RESIDUAL CERTIFICATE AND SAMPLED CHECK

• All 15 evaluated returns are certified; on the grid, Ω is at most 1% of τ and R_A agrees with recorded full residuals within 1.2 × 10^-6, while certification takes minutes rather than about 38 hours.The certificate combines the FP32 residual of the admitted plan with an FP64 bound on omitted tile mass; the two 227 solutions have Ω below 10^-7.Source: §C.4 · p. 21, Table 8 · p. 23
• None of 40 reported full-residual returns exceeds τ, and none of 140 seeded re-solves of seven near-threshold problems exceeds τ, bounding false acceptance at 2.6% (95% confidence).The 2.6% bound is reported for the 140 re-solves; it is 6.0% over the 60 solves on three problems whose stopping update moved with the seed.Source: §C.4 · p. 22
• The production check accepts none of 6,000 perturbed states, although the row-only check accepts up to 95.8%; these measurements do not cover n ≥ 2^22, other ε, non-uniform weights, or the Quijote blur.The perturbations affect a random 0.01%, 0.1%, or 1% of rows, and the production check also tests columns.Source: §C.4 · p. 22

C.5 ABLATIONS

• Removing the coarse stage slows successful solves by 154–1643×, while shorter continuation can miss τ or require 51.6–60.0× the FS2 time.On clustered ε = 0.001, the no-coarse control stops at the 128-update fine cap with R = 0.02135; raising the cap to 8192 after shorter continuation requires 1115–1305 fine updates.Source: §C.5 · p. 24, Table 9 · p. 24, Table 10 · p. 24
• On clustered ε = 0.001, sparse execution cuts time from 33.327 to 4.165 s (8.00×); at ε = 0.01, neither geometry takes a fine update.On heterogeneous ε = 0.001, the mask admits 60.6% of tiles and sparse execution gains 1.33×.Source: §C.5 · p. 24
• Tile verification cuts clustered point-pair work from 3.049×10^12 to 2.087×10^12 and time from 4.335 to 4.167 s, but costs time on heterogeneous data.On heterogeneous data, verification removes 3.95% of the work and takes 1.824 s versus 1.576 s without it; threshold controls still meet the target.Source: §C.5 · pp. 24–25, Table 10 · p. 24

C.6 PROFILING

• On the same B = 128 mask, FS2’s fine update takes 504.392 ms versus KeOps’ 791.073 ms, with 55.7% fewer warp instructions per pair.The profiled comparison uses the same sorted points, incoming potentials, weights, and transposed support; the dense FlashSinkhorn reduction takes 3.6 s, 14× longer than FS2’s masked reduction.Source: §C.6 · p. 25, Table 11 · p. 25

C.7 MEMORY

• At n = 223, increasing block size from 128 to 256 cuts admitted tiles from 2.25×10^8 to 6.0×10^7 and peak memory from 16.17 to 6.88 GiB, for 7% more time.Both settings return R = 0.004694; at n = 222 and 223 on heterogeneous data, FS2 returns lifted coarse potentials without a fine update or mask.Source: §C.7 · p. 25

C.8 OPTIMIZATION

• The optimization evaluator’s grouped score ranks trajectories rather than reporting the full-point OT value, and GeomLoss reaches the time limit on six problems.The score combines entropic cost between group centroids with within-group squared-distance averages; the twenty-step runs include compilation and autotuning and stop at 7200 s.Source: §C.8 · p. 26, Table 12 · p. 26
• At the four n = m = 221 initializations, the surrogate direction aligns more closely with the reference on clustered inputs, while heterogeneous relative errors are 0.023 and 0.012.On clustered data, cosine similarity is 0.92 and 0.96 for the surrogate versus 0.49 and 0.40 for GeomLoss; these compare to a numerical reference, not the exact gradient.Source: §C.8 · p. 26

D.1 PROBLEM AND DATA

• The cosmological inputs match redshift-zero positions to Lagrangian lattice sites; the 224 instance is a central 256^3 subcube and the 227 instance the full 512^3 box.At ε = 3.814697265625, blur √ε = 1.953125 Mpc/h equals lattice spacing; TF32-rounded coordinates move each particle by at most 0.125 Mpc/h, 6% of the blur.Source: §D.1 · p. 26

D.2 CONFIGURATION AND TIMING

• At 227 and τ = 0.01, verification admits 24,451,822 tiles holding 1.60×10^12 point pairs; their 24-byte-per-tile index takes 2.9 GiB versus 11.7 TiB for a truncated pair kernel.The box screen passes 107,397,067 of the 238 block pairs; the selected cell level and check interval are configured separately for the two instance sizes.Source: §D.2 · pp. 26–27
• At 227, τ = 0.01, the full-residual evaluator takes about 38 hours on an A100, and validation is excluded from the timed calls.The 227 configuration uses t = 45 because t = 46 makes the sampled check non-finite; using the longer coarse continuation takes 2030.4 s and 119 ordinary updates rather than 1351.6 s and 80.Source: §D.2 · p. 27

D.3 EVALUATOR PRECISION

• FP64 correction gives row/column residuals 0.04653/0.04661, 0.04883/0.04889, 0.00939/0.00942, and 0.00882/0.00869; each remains below its τ.Near box corners, FP32 spacing is 1.6% of ε, so sampled row errors are recomputed in FP64 from coordinate differences and corrected by the mean paired difference.Source: §D.3 · p. 27

D.4 SIZE SWEEP

• For the τ = 0.05 size sweep, runs pass the solver’s sampled check and are validated on 65,536 sampled outputs per side, without enumerating the full residual.The sweep uses nested cuboids from 256^3 to 512^3 lattice sites with lattice spacing and ε fixed.Source: §D.4 · p. 27

D.5 DISPLACEMENT RECONSTRUCTION

• On 4096 sampled particles, displacement component correlations are 0.797, 0.789, and 0.802, and unwrapped position RMSE is 122.268 Mpc/h.The minimum-image wrapping is used only for scoring, while the solver’s cost remains nonperiodic; projection takes 9.303 s, excluding 43.474 s of loading and preparation.Source: §D.5 · p. 28, Table 14 · p. 28