Table of Contents
Fetching ...

Threshold Differential Attention for Sink-Free, Ultra-Sparse, and Non-Dispersive Language Modeling

Xingyue Huang, Xueying Ding, Mingxuan Ju, Yozen Liu, Neil Shah, Tong Zhao

TL;DR

This work identifies core flaws in Softmax attention for long sequences—attention sinks and dispersion—and introduces Threshold Differential Attention (TDA), a sink-free, ultra-sparse mechanism that combines length-aware thresholding with an inhibitory differential view. Theoretical analysis under sub-Gaussian noise shows that per-row spurious survivors remain $O(1)$ and consensus spurious matches vanish as context grows, while empirical results report >99% exact zeros and elimination of attention sinks. TRA provides a strong baseline by bounding noise via a length-dependent gate, and TDA further suppresses spurious matches by combining two independent views into a signed, differential map. Across standard and long-context benchmarks, TDA achieves competitive accuracy with significantly improved sparsity, and the authors provide a fused Triton kernel that delivers memory and runtime benefits, making long-context Transformers more practical. Overall, TDA offers a principled, efficient path toward robust long-context modeling without projection-based sparsification.

Abstract

Softmax attention struggles with long contexts due to structural limitations: the strict sum-to-one constraint forces attention sinks on irrelevant tokens, and probability mass disperses as sequence lengths increase. We tackle these problems with Threshold Differential Attention (TDA), a sink-free attention mechanism that achieves ultra-sparsity and improved robustness at longer sequence lengths without the computational overhead of projection methods or the performance degradation caused by noise accumulation of standard rectified attention. TDA applies row-wise extreme-value thresholding with a length-dependent gate, retaining only exceedances. Inspired by the differential transformer, TDA also subtracts an inhibitory view to enhance expressivity. Theoretically, we prove that TDA controls the expected number of spurious survivors per row to $O(1)$ and that consensus spurious matches across independent views vanish as context grows. Empirically, TDA produces $>99\%$ exact zeros and eliminates attention sinks while maintaining competitive performance on standard and long-context benchmarks.

Threshold Differential Attention for Sink-Free, Ultra-Sparse, and Non-Dispersive Language Modeling

TL;DR

This work identifies core flaws in Softmax attention for long sequences—attention sinks and dispersion—and introduces Threshold Differential Attention (TDA), a sink-free, ultra-sparse mechanism that combines length-aware thresholding with an inhibitory differential view. Theoretical analysis under sub-Gaussian noise shows that per-row spurious survivors remain and consensus spurious matches vanish as context grows, while empirical results report >99% exact zeros and elimination of attention sinks. TRA provides a strong baseline by bounding noise via a length-dependent gate, and TDA further suppresses spurious matches by combining two independent views into a signed, differential map. Across standard and long-context benchmarks, TDA achieves competitive accuracy with significantly improved sparsity, and the authors provide a fused Triton kernel that delivers memory and runtime benefits, making long-context Transformers more practical. Overall, TDA offers a principled, efficient path toward robust long-context modeling without projection-based sparsification.

Abstract

Softmax attention struggles with long contexts due to structural limitations: the strict sum-to-one constraint forces attention sinks on irrelevant tokens, and probability mass disperses as sequence lengths increase. We tackle these problems with Threshold Differential Attention (TDA), a sink-free attention mechanism that achieves ultra-sparsity and improved robustness at longer sequence lengths without the computational overhead of projection methods or the performance degradation caused by noise accumulation of standard rectified attention. TDA applies row-wise extreme-value thresholding with a length-dependent gate, retaining only exceedances. Inspired by the differential transformer, TDA also subtracts an inhibitory view to enhance expressivity. Theoretically, we prove that TDA controls the expected number of spurious survivors per row to and that consensus spurious matches across independent views vanish as context grows. Empirically, TDA produces exact zeros and eliminates attention sinks while maintaining competitive performance on standard and long-context benchmarks.
Paper Structure (58 sections, 8 theorems, 41 equations, 4 figures, 6 tables)

This paper contains 58 sections, 8 theorems, 41 equations, 4 figures, 6 tables.

Key Result

Theorem 4.3

Under Assumption ass:tra-subg, fix any $\kappa>0$, then for all $i\ge 1$, $\mathbb{E}[S_i] \;\le\; \kappa.$

Figures (4)

  • Figure 1: Mechanistic diagnostics for TDA. We visualize (a) the sparsity of attention weights across layers, (b) the maximum absolute attention values, and (c) the first token attention sink ratio $\mathrm{gSinkRatio}(1)$ as context length increases. Attention 1 indicates an excitatory view, while Attention 2 indicates an inhibitory view.
  • Figure 2: Inhibition in TDA. Per-token attention maps for a representative head (Layer 6, Head 0) on the sequence "The quick brown fox jumps over the lazy dog". The differential attention maps show $\Delta {\bm{A}} = {\bm{A}}^{(1)} - \lambda {\bm{A}}^{(2)}$ where $\lambda=1$; negative values indicate inhibition, while positive values indicate excitation.
  • Figure 3: Passkey retrieval results. We report correct runs over 100 trials for each context length, with randomly positioned passkey per trial.
  • Figure 4: Runtime and memory of fused Triton TRA. We compare dense Softmax (PyTorch), naive TRA (PyTorch), fused SDPA, and our fused Triton TRA kernel under FP32.

Theorems & Definitions (12)

  • Theorem 4.3: TRA keeps $O(1)$ spurious survivors per row
  • Theorem 4.4: TRA is Non-Dispersive
  • Theorem 4.6: Consensus spurious survivors vanish
  • Corollary 4.7: TDA is Non-Dispersive
  • Theorem A.4: TRA keeps $O(1)$ spurious survivors per row (restated \ref{['thm:tra-survivors']})
  • proof
  • Theorem A.5: TRA is Non-Dispersive (restated \ref{['thm:tra-non-dispersion']})
  • proof
  • Theorem A.6: Consensus spurious survivors vanish in TDA (restated \ref{['thm:tda-consensus']})
  • proof
  • ...and 2 more