Block-Wise Differentiable Sinkhorn Attention: Tail-Refinement Gradients with a Gap-Aware Dustbin Bridge
cs.LG, cs.CL
Submitted: 2026-04-28
Updated: 2026-09-07
License: http://creativecommons.org/licenses/by/4.0/
The gist: We study long-context balanced entropic optimal transport (OT) attention on TPU hardware through a stopped-base, fixed-depth tail-refinement surrogate.
Terminology
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 T-step Sinkhorn solve, we unroll a short refinement tail and differentiate that surrogate exactly. For the reported R=2 TPU path, the backward pass contains four staircase plan factors. We prove an exact one-reference-tile schedule: the R=2 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 O((T+R)LW), O(Ld) input storage, and O(L) additional HBM usage for fixed head dimension d and band width W on the balanced fixed-support path. We also formalize the current 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 10-5 -- 10-10. On TPU v6e-8, a four-configuration Pfam screen completes end-to-end, and a promoted balanced R=2 run sustains roughly 8.5 examples per second through a three-hour budget, reaching step 1437. Held-out Pfam test shards improve reconstruction from 5.57 to 2.05 and sparse CE from 5.53 to 5.30 relative to step 0, 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.
Sources
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: Local-Consistency Optimisation for Bayesian Neural Networks
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks