Block-Wise Differentiable Sinkhorn Attention: Tail-Refinement Gradients with a Gap-Aware Dustbin Bridge
Abstract
We study long-context balanced entropic optimal transport (OT) attention on TPU hardware through a stopped-base, fixed-depth tail-refinement surrogate. After a stopped -step Sinkhorn solve, we unroll a short refinement tail and differentiate that surrogate exactly. For the reported TPU path, the backward pass contains four staircase plan factors. We prove an exact one-reference-tile schedule: the score cotangent is a single reference plan tile times an explicit modifier field built from vector cotangents and dual differences. This yields block-wise cost , input storage, and additional HBM usage for fixed head dimension and band width on the balanced fixed-support path. We also formalize the current \texttt{dustbin\_block} path as the same unit-target surrogate on an augmented support, so the adjoint schedule lifts to the single-active-dustbin path used in our TPU runs; this bridge is algebraic and does not claim a general KL-unbalanced or arbitrary-capacity gap model. We provide a local surrogate-bias bound, an a posteriori bias certificate, and a projective contraction certificate for strictly positive active blocks. On synthetic masked problems, the optimized kernel matches exact autodiff of the same centered surrogate to within --. On TPU v6e-8, a four-configuration Pfam screen completes end-to-end, and a promoted balanced run sustains roughly examples per second through a three-hour budget, reaching step . Held-out Pfam test shards improve reconstruction from to and sparse CE from to relative to step , with CE logged diagnostically rather than optimized directly; target-barycenter alignment metrics do not materially improve, and a deterministic diagonal reference remains stronger on those metrics.
Keywords
Cite
@article{arxiv.2605.08123,
title = {Block-Wise Differentiable Sinkhorn Attention: Tail-Refinement Gradients with a Gap-Aware Dustbin Bridge},
author = {Dylan Forde},
journal= {arXiv preprint arXiv:2605.08123},
year = {2026}
}