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

AI summary

Overview

Research area: Machine learning / large language model architecture — specifically efficient attention mechanisms for long-context language modeling.

Technical level: Advanced. The paper combines an architectural proposal with extreme-value theory, sub-Gaussian concentration arguments, and entropy-based definitions of attention dispersion, alongside standard language-modeling benchmarks.

Scope in one sentence: The paper proposes Threshold Differential Attention (TDA), a non-softmax attention mechanism that uses a length-dependent extreme-value threshold plus an inhibitory differential view to produce sink-free, ultra-sparse (>99% exact zeros) attention that stays non-dispersive as context grows.

What This Paper Is About

Softmax attention enforces a strict sum-to-one constraint on attention weights, which the authors argue causes two structural pathologies: attention sinks (probability mass forced onto irrelevant tokens, often the first position) and attention dispersion (mass diluting toward uniformity as sequence length grows). The paper's goal is an attention mechanism that removes the sum-to-one constraint, produces exact zeros, remains selective at long contexts, and does so without the sorting or iterative-projection overhead of prior sparse methods.

Key Contributions

1. TDA, a drop-in non-softmax attention mechanism. TDA applies a length-dependent extreme-value threshold to rectified scores and subtracts an inhibitory thresholded view (inspired by the differential transformer), yielding signed, ultra-sparse, sink-free attention. The single-view variant is called Threshold Rectified Attention (TRA).

2. Theoretical guarantees under sub-Gaussian assumptions. The expected number of spurious survivors per row stays O(1), and consensus spurious matches across the two independent views vanish as context grows (E[C_i] ≤ κ²/(i+1)). TRA and TDA are both proved non-dispersive.

3. Empirical validation. TDA reaches >99% exact-zero sparsity while remaining competitive on standard and long-context QA, and eliminates attention sinks across layers and heads.

4. A fused Triton kernel implementation (released at https://github.com/snap-research/TDA.git) whose TRA kernel is competitive with FlashAttention-2 under BF16, with consistent speedups at long contexts.

Main Findings

• Two pathologies attributed to softmax. The strict sum-to-one constraint is said to cause (i) attention sinks on irrelevant tokens and (ii) attention dispersion, where probability mass is progressively diluted as sequence length grows. The paper defines a mechanism as dispersive if lim E[H(a_i)] / log i = 1, and non-dispersive if this ratio approaches 0.

• TRA bounds spurious survivors at O(1). Theorem 4.3 states that under Assumption 4.1 (mean-zero, σ²/d sub-Gaussian noise per row) and with β ≥ σ, E[S_i] ≤ κ for all i ≥ 1. The threshold used is τ_i = β · sqrt(2 log((i+1)/κ) / d) with κ > 0.

• TDA drives consensus noise to zero. Theorem 4.6 states that under Assumptions 4.1 and 4.5 (independence of the two views' noise), E[C_i] ≤ κ²/(i+1), so lim E[C_i] = 0 as i → ∞. Corollary 4.7 states TDA is non-dispersive.

• TRA and TDA differ from all listed baselines on four features. Table 1 compares exact-zero output, signed weights, no sum-to-one normalization, and length awareness: TDA is the only method in the table with all four; TRA has all but signed weights; Softmax has none of the four.

• Sparsity profile is depth-dependent. Marking an entry inactive when its magnitude is exactly 0, the paper reports that early and late layers are highly sparse while middle layers are substantially more active, with the near-zero rate dropping to about 50%. Attention 1 is the excitatory view and Attention 2 the inhibitory view.

• Differential subtraction reduces peak attention. The layerwise maximum absolute attention of the differential map ΔA = A⁽¹⁾ − λA⁽²⁾ is substantially smaller than either individual view, indicating that inhibition cancels large common-mode exceedances.

• Sinks are removed. Under Softmax, the first-token sink ratio gSinkRatio(1) increases sharply with sequence length; under TDA it stays near or below the uniform baseline as context length grows.

• Inhibition is semantic as well as global. In a modified GPT-2 (RoPE, TDA replacing softmax) run on "The quick brown fox jumps over the lazy dog," the high-frequency preposition "the" is broadly inhibited across many queries, while content tokens such as "quick" and "brown" show query-dependent inhibition.

• Language modeling results. Pretrained on FineWebEdu-10B (10B tokens, first 100M reserved for validation) using a GPT-2-162M variant with RoPE, TDA achieves the lowest validation loss at 3.1190 and 99% sparsity (Table 2). Softmax is 3.1196 at 0% sparsity; Diff ReLA is 3.1294 at 96%; TRA is 3.1320 at 92%; ReLA is 3.1657 at 94%; Entmax is 3.1941 at 43%; Gated Softmax, SSMax, LSSA, and Diff Softmax are all 0%. On individual tasks TDA scores HellaSwag 0.337/0.415 (Acc/Acc-Norm), ARC-Easy 0.524/0.488, ARC-Challenge 0.220/0.239, OpenBookQA 0.216/0.320, PIQA 0.628/0.626, and Winogrande 0.489 (Acc only), versus Softmax at 0.345/0.409, 0.526/0.487, 0.223/0.245, 0.180/0.304, 0.641/0.621, and 0.490.

• Long-context results. On SCROLLS (Table 3), TDA scores 11.46 ROUGE-1 on QMSum, 9.13 on SummScreenFD, 5.24 on GovReport, and 11.41 F1 on Qasper. It is second-best on QMSum and Qasper, where Entmax leads at 11.52 and 11.54. TRA achieves the best GovReport score at 5.61 and is second on SummScreenFD at 9.47 (Entmax 10.16). The paper notes Entmax is substantially more expensive at long context.

• Passkey retrieval. Testing target context lengths from 500 to 4000 tokens in increments of 500 with 100 trials per length, both Softmax and TDA degrade as irrelevant text grows, but TDA is consistently more robust. At 4000 tokens, TDA achieves 15% correct versus 6% for Softmax.

• Multi-needle retrieval. TDA outperforms Softmax in both the 2-needle and 4-needle settings (Table 4). At 4000 tokens with 2 needles, TDA reaches 82.0 vs. Softmax 0.0; with 4 needles, TDA reaches 20.0 vs. Softmax 0.0. At 500 tokens with 4 needles, TDA is 62.0 vs. Softmax 18.0.

• Kernel efficiency. Under BF16 (Table 5), the TRA kernel is slower at 512 and 1024 tokens (speedups 1.09× and 0.89×) due to fixed kernel overhead, roughly matches FlashAttention-2 around 4096 tokens (0.98×), and speeds up at longer contexts: 1.13× at 8192, 1.25× at 16384, 1.29× at 32768, and 1.24× at 65536 tokens.

• Hyperparameter sensitivity. For the power p (Table 6a), p=1 gives validation loss 3.2068 and average accuracy 0.3945; p=2 gives 3.1190 and 0.4023; p=3 gives 3.1408 and 0.4020; p=5 gives 3.1412 and 0.3922. For the threshold scaling β (Table 6b), 1.0 gives 3.1190 and 0.4023, 0.8 gives 3.1140 and 0.4015, and 0.5 gives 3.1288 and 0.4018. The paper states β=1.0 achieves the best validation loss and highest average accuracy while also describing performance as fairly robust across the tested range. Average accuracy is computed over HellaSwag, ARC-Easy, ARC-Challenge, OpenBookQA, PIQA, and Winogrande.

Methodology in Plain English

The authors start from rectified attention, which replaces softmax with a ReLU-style activation. This removes the sum-to-one constraint and naturally yields exact zeros, but a fixed rectifier threshold stops filtering effectively as context grows, because random (query, key) dot products produce larger extremes by chance as more keys become visible. They call this noise accumulation and credit it with ReLA's long-context weakness.

Their fix in TRA is to make the threshold grow with the visible context size. Normalizing queries and keys to unit length, each row i gets a threshold τ_i = β · sqrt(2 log((i+1)/κ) / d), where β is a learnable scalar, κ controls the expected number of spurious survivors, and d is the per-head dimension. The scale follows the standard sub-Gaussian maximum bound, under which a threshold of order sqrt(2 log i / d) keeps tail exceedances under control. Attention weights are then (s_ij − τ_i)_+^p, and the output is an RMSNorm over the weighted sum of values. Because the threshold is a row-wise order statistic rather than a learned projection, no sorting or iterative projection is required.

TDA extends TRA by running two independent sets of query and key projections, computing a thresholded map for each, and subtracting them: Δa_ij = a_ij⁽¹⁾ − λ·a_ij⁽²⁾ with λ in (0,1) learned. The second view acts as an inhibitory signal trained to capture non-selective exceedances that carry no useful content, so noise that survives in both views is what remains. Because the two views are treated as independent, the probability that the same noise key exceeds the threshold in both falls off as 1/(i+1), while a single view retains only O(1) spurious survivors.

Evaluation covers pretraining from scratch on FineWebEdu-10B, zero-shot multiple-choice QA, four SCROLLS long-context tasks, passkey and multi-needle retrieval stress tests, and BF16 kernel timing against FlashAttention-2. All experiments use κ=1, β=1, and p=2, run on 8 NVIDIA A100-80GB GPUs, with NTK-aware RoPE scaling and 500 additional training steps when extending context.

Why This Matters

Impact on research. The paper reframes the attention-sink and dispersion problems as consequences of the sum-to-one constraint rather than incidental artifacts, and offers a construction that satisfies all four properties it identifies (exact zeros, signed weights, no normalization, length awareness) — a combination no baseline in its Table 1 achieves. It also supplies matching non-dispersiveness and noise-control proofs, giving subsequent work a formal

Authors’ 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.