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

• URL: http://arxiv.org/abs/2605.08123v2
• Date: Wed, 20 May 2026 01:20:05 GMT
• Title: Block-Wise Differentiable Sinkhorn Attention: Tail-Refinement Gradients with a Gap-Aware Dustbin Bridge
• Abstract summary: We study long-context balanced optimal transport attention on TPU hardware through a stopped-base, fixed-depth tail-refinement surrogate.<n>We provide a local surrogate-bias bound, an a posteriori bias certificate, and a projective contraction certificate for strictly positive active blocks.<n>On synthetic masked problems, the optimized matches exact autodiff of the same centered surrogate to within $10-5$--$10-10$.
• Score: 0.0
• License: http://creativecommons.org/licenses/by/4.0/
• 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 \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 $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.

Related papers

• Online Non-Monotone DR-Submodular Maximization Matching the Offline $0.401$ Factor [47.30196146041476]
  We study asymmetric non-negative, non-monotone DR-submodular functions over compact convex subsets of the $d$-oracle unit.<n>Online algorithm attains factor 0.401$ with sub approximate regret when feedback is conditionally unbiased and bounded.
  arXiv  Detail & Related papers  (2026-09-02T06:00:10Z)
• BiKAN: Restoring Collapsed Basis of Binary Kolmogorov--Arnold Networks [6.860988566886594]
  Binarizing a Kolmogorov--Arnold Network (KAN) changes function space available to each layer.<n>Our proposed BiKAN addresses this critical issue by augmenting each binary KAN with selected degree-2 Walsh characters.<n>At W1A1, BiKAN reaches $99.48%$, $84.38%$, and $55.81%$ on M, 401AR-10, and 72AR-100, respectively.
  arXiv  Detail & Related papers  (2026-08-02T20:49:59Z)
• ForgettingOT: Certified Speculative Batching from Sinkhorn's Projective Forgetting [0.0]
  We show that the eigenmode of the fixed-point Jacobian $J_tstar=QP=Pstar P$ controls the strict correction tail.<n>ForgettingOT turns this nonlinear Perron--Frobenius fact into a certified executor for streams of related Sinkhorn problems.
  arXiv  Detail & Related papers  (2026-07-27T17:59:36Z)
• Scaling Federated Linear Contextual Bandits via Sketching [49.12000877146222]
  This paper proposes Federated Sketch Contextual Linear Bandits (FSCLB)<n>Experiments on both synthetic and real-world datasets show that FSCLB significantly reduces computational and communication costs by over 90 %.
  arXiv  Detail & Related papers  (2026-05-01T08:22:06Z)
• Bond-dimension scaling of a local-refinement advantage over hyperoptimized tensor-network contraction on Sycamore like topologies [0.0]
  We identify a missing local-refinement stage in the cotengra tensor-network contraction pipeline.<n>We show that its impact grows monotonically with bond dimension on the emphconnectivity graph of Sycamore-like topologies.
  arXiv  Detail & Related papers  (2026-04-28T11:59:31Z)
• Dimensional Criticality at Grokking Across MLPs and Transformers [2.652953665748039]
  Abrupt transitions between distinct dynamical regimes are a hallmark of complex systems.<n>We introduce textbfTDU--OFC (Thresholded Diffusion Update--Olami-Feder-Christensen), an offline avalanche probe.<n>Across Transformers trained on modular addition and generalizations trained on XOR, we discover a localized crossing of the diffusion baseline $D=1$.
  arXiv  Detail & Related papers  (2026-04-06T13:43:20Z)
• IsoQuant: Hardware-Aligned SO(4) Isoclinic Rotations for LLM KV Cache Compression [0.4496256885343706]
  We propose a blockwise rotation framework based on quaternion algebra and the isoclinic decomposition of $SO(4)$.<n>IsoQuant achieves mean kernel-level speedups of $4.5times$--$4.7times$ over RotorQuant with peak speedups above $6times$.
  arXiv  Detail & Related papers  (2026-03-30T13:37:45Z)
• Large Language Models as Optimization Controllers: Adaptive Continuation for SIMP Topology Optimization [4.365822392824561]
  We present a framework in which a large language model acts as an online controller for SIMP optimization.<n>A hard grayness gate prevents premature binarization, and a meta-optimization loop uses a second pass to tune the agent's call frequency.<n>The agent achieves the final compliance on every benchmark: $-5.7%$ to $-1% relative to the fixed baseline, with all solutions fully binary.
  arXiv  Detail & Related papers  (2026-03-26T07:14:31Z)
• Witness wedges in fidelity-deviation plane: separating teleportation advantage and Bell-inequality violation [0.0]
  We develop a unified framework to analyze quantum teleportation through the joint geometry of two complementary figures of merit.<n>In particular, any measured pair $(F D)$ can be ported into a visibility estimate for isotropic resources.
  arXiv  Detail & Related papers  (2025-11-26T05:43:31Z)
• Proving the Limited Scalability of Centralized Distributed Optimization via a New Lower Bound Construction [57.93371273485736]
  We consider a centralized distributed learning setup where all workers jointly find an unbiased bound LDeltaepsilon2,$ better poly-logarithmically in $n$, even in the homogeneous (i.i.d.) case, where all workers access the same distribution.
  arXiv  Detail & Related papers  (2025-06-30T13:27:39Z)
• Improved Algorithm for Adversarial Linear Mixture MDPs with Bandit Feedback and Unknown Transition [71.33787410075577]
  We study reinforcement learning with linear function approximation, unknown transition, and adversarial losses. We propose a new algorithm that attains an $widetildeO(dsqrtHS3K + sqrtHSAK)$ regret with high probability.
  arXiv  Detail & Related papers  (2024-03-07T15:03:50Z)
• Horizon-Free and Variance-Dependent Reinforcement Learning for Latent Markov Decision Processes [62.90204655228324]
  We study regret minimization for reinforcement learning (RL) in Latent Markov Decision Processes (LMDPs) with context in hindsight. We design a novel model-based algorithmic framework which can be instantiated with both a model-optimistic and a value-optimistic solver.
  arXiv  Detail & Related papers  (2022-10-20T21:32:01Z)
• Variance-Aware Confidence Set: Variance-Dependent Bound for Linear Bandits and Horizon-Free Bound for Linear Mixture MDP [76.94328400919836]
  We show how to construct variance-aware confidence sets for linear bandits and linear mixture Decision Process (MDP) For linear bandits, we obtain an $widetildeO(mathrmpoly(d)sqrt1 + sum_i=1Ksigma_i2) regret bound, where $d is the feature dimension. For linear mixture MDP, we obtain an $widetildeO(mathrmpoly(d)sqrtK)$ regret bound, where
  arXiv  Detail & Related papers  (2021-01-29T18:57:52Z)
• Linear Time Sinkhorn Divergences using Positive Features [51.50788603386766]
  Solving optimal transport with an entropic regularization requires computing a $ntimes n$ kernel matrix that is repeatedly applied to a vector. We propose to use instead ground costs of the form $c(x,y)=-logdotpvarphi(x)varphi(y)$ where $varphi$ is a map from the ground space onto the positive orthant $RRr_+$, with $rll n$.
  arXiv  Detail & Related papers  (2020-06-12T10:21:40Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.