Research map / Learning theory and algorithms / Optimal transport

Block-Wise Differentiable Sinkhorn Attention: Tail-Refinement Gradients with a Gap-Aware Dustbin Bridge

NeurIPS 2026 · Learning theory and algorithms › Optimal transport · approach: Sinkhorn and entropic OT algorithms · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Derives an exact fixed-depth (R=2) backward for a stopped-base Sinkhorn attention surrogate, showing four staircase plan factors lie in one row/column scaling orbit so the backward can use one reference tile, and validates it on synthetic tests, TPU v6e scaling, and a small Pfam screen.

Problem

Differentiating long-context entropic optimal-transport (Sinkhorn) attention is memory-hungry: exact autodiff through many Sinkhorn steps is costly and implicit differentiation needs unstable linear solves; the exact R=2 adjoint naively needs four distinct plan tiles.

Approach

Treat the differentiated object as a stopped-base, fixed-depth tail-refinement surrogate; prove that all staircase plans lie in one row/column scaling orbit, so the exact score cotangent can be computed from one resident reference plan tile plus O(B) vector modifiers, with O(L) auxiliary state.

Key results

• Exact R=2 one-reference backward matches autodiff through the centered surrogate with max absolute gradient deviations of about 1e-5 to 1e-9 (relative l2 errors ~5-6%) on synthetic L=512-2048 problems.
• Complete arbitrary-cotangent Q/K/V VJP verified at per-device context L=262,144 on all-eight-device TPU v6e data-parallel; at that length one-reference takes 6020 ms vs direct-four/recompute 5415 ms.
• In a 100-update, three-seed Pfam reconstruction screen, trained FlashAttention had lower reconstruction MSE than Sinkhorn in all observed pairs (mean Sinkhorn-minus-FlashAttention difference +0.0018).

Surprising findings

• Trained FlashAttention beats the Sinkhorn attention on the short reconstruction screen in every observed seed pair.
• Direct-four/recompute has a lower VJP median than the one-reference schedule in the TPU follow-up, so the one-reference representation is not a latency win.

Limitations

• Exactness is for the stopped-base surrogate, not full backprop through the stopped base solve; the surrogate-bias bound depends on contraction assumptions.
• Biological alignment results do not show superiority over diagonal reference; target-barycenter metrics do not improve, and the direct CE objective fails its selector.
• Dustbin path is a single-active-token unit-target bridge with widened band cost O(L*max(W,L)), and TPU memory numbers are static compiler estimates rather than measured dynamic HBM.

Experimental setup

Datasets: Pfam ESM-2 embeddings

Benchmarks: Pfam held-out test shard (100 examples) synthetic exact-autodiff comparison

Metrics: max absolute gradient deviation relative l2 error reconstruction MSE sparse CE barycenter MAE within-five accuracy VJP latency (ms) compiler MiB post-compilation CPU speedup

Baselines: FlashAttention (trained) direct four-plan adjoint exact autodiff through centered surrogate full BPTT through stopped base diagonal reference alignment

Builds on: Sinkhorn-Knopp scaling FlashAttention FlashSinkhorn Birkhoff-Hopf contraction

Closest papers by meaning

• Phase Space Attention: A Hairer Lift Resolves the Single-Layer Induction Obstruction (NeurIPS 2026 · similarity 0.84)
• PISA: Piecewise Sparse Attention Is Wiser for Efficient Diffusion Transformers (NeurIPS 2026 · similarity 0.84)
• On Partial Optimal Transport: Revising the Infeasibility of Sinkhorn and Efficient Gradient Methods (AAAI 2024 · similarity 0.83)
• Log-Averaged Mirror Prox for Fast, Large-Scale Optimal Transport in Linear Space (NeurIPS 2026 · similarity 0.83)
• Faster Video Diffusion with Trainable Sparse Attention (NeurIPS 2025 · similarity 0.83)
• Fast T2T: Optimization Consistency Speeds Up Diffusion-Based Training-to-Testing Solving for Combinatorial Optimization (NeurIPS 2024 · similarity 0.83)