We call query--key relevance absolute when its values lie on a fixed bounded scale, depend on neither competing keys nor sequence length, require no sequence-length-dependent calibration, and can all be zero. To realize this notion, we introduce screening, whose explicit threshold transforms bounded query--key similarities into relevance values, enabling exact rejection, empty selection, and direct inspection on a common scale. In a controlled comparison of 12 attention mechanisms on a matched Transformer backbone, only screening maintains both low long-context perplexity and robust retrieval beyond the training context; notably, it does so without inference-time scaling. Building on screening, we introduce Multiscreen, a language-model architecture composed of parallel gated screening tiles. Multiscreen retains these long-context gains while achieving greater parameter efficiency, stronger general zero-shot downstream performance, lower training cost at larger scales, and lower model-side time to first token than Transformer baselines. We further develop a normalization design that keeps Multiscreen training stable even at a learning rate of 1 and show that an adapted version likewise stabilizes Transformer at the same learning rate.
Figures & tables
Figure 1: Conceptual comparison between softmax attention and screening. Each cell corresponds to a scalar value associated with an individual key. Colors indicate scalar values (blue: negative, red: positive; intensity indicates magnitude), while hatched cells denote zero values. Softmax attention maps unbounded attention scores to nonnegative attention weights that sum to one, redistributing mass across all keys. In contrast, screening maps bounded query–key similarities to distance-unaware relevance via Trim, which assigns exactly zero relevance to similarities below a threshold, and then applies Softmask to produce distance-aware relevance.
Figure 2: (a) Multiscreen comprises NL residual layers with NH parallel gated screening tiles per layer; WE is the row-wise unit-normalized embedding matrix shared by the input and language-modeling head, and esE and esF are learned input and logit scales. Final hidden states are unit-normalized before the head. (b) Each gated screening tile applies query, key, value, and gate projections, screening, gating, and an output projection with learned residual-output scale esO . (c) The screening unit applies QKV unit-length normalization, MiPE, Trim, Softmask, and TanhNorm. In the diagrams, “@” denotes matrix multiplication and “/RSS” denotes row-wise unit-length normalization.
Figure 3: Comparison between RoPE and MiPE. RoPE applies rotations across all dimensions with fixed frequencies. MiPE applies rotation only to two dimensions and adapts the rotation angle based on the learned screening window w , becoming inactive when w≥wth .
Figure 4: Position-dependent perplexity of the 1.3B models. Colored curves show Transformer with different YaRN scaling factors, whereas the black curve shows Multiscreen, which requires no inference-time scaling. The dashed line marks the pretraining context length ( 212 ).
Figure 5: (a) Example ABCDigits prompt. (b) ABCDigits exact-match accuracy across model sizes, context lengths, and target depths. Each cell averages 1,000 instances; results are further averaged across three independently trained models, except at 1.3B, where one model is used. Transformer uses context-matched YaRN scaling, while Multiscreen uses our default inference rule. The dashed line marks the pretraining context length ( 212 ).
Figure 6: Controlled attention-mechanism comparison on the 353M Transformer backbone with all other settings fixed. The top row shows position-dependent perplexity; the bottom row shows ABCDigits retrieval accuracy. Each variant uses one trained model. In the top row, red denotes no inference-time scaling; for methods requiring scaling, red corresponds to factor 1 and the other colors to the indicated factors. For ABCDigits, methods requiring scaling use a factor equal to the evaluation-to-training context-length ratio. For ReLU2 , scaling factors 4 and above exceed the plotted perplexity range. Dashed lines mark the pretraining context length ( 212 ).
Figure 7: Learning-rate stability of Transformer and Multiscreen. (a) Validation loss across learning rates, with arrows at the upper plot boundary indicating divergence. Each point uses one training run, except Multiscreen at learning rate 2−4 , which reuses the three runs from the scaling experiments and shows their mean and standard deviation. (b) Training outcomes at learning rate 1 after 236 tokens ( 214 steps), where “Stable” indicates completion of this budget without divergence. PLNorm denotes pre-logit normalization. Complete results are provided in Appendix I .
Figure 8: Scaling and computational efficiency. (a) Validation loss and (b) training cost in GPU hours across model sizes, with each model trained on 238 tokens. (c) Model-side TTFT across model sizes at context length 217 . (d) Model-side TTFT across context lengths for the 1.3B models. Error bars indicate standard deviation across independently trained models where multiple runs are available.
Appendix figures & tables25 assets
Supplementary material from the paper’s appendix.
Appendix
Accuracy (%) ↑
Benchmark
Transformer
Multiscreen
HellaSwag
55.58
59.02
PIQA
72.91
73.99
ARC-Easy
52.82
55.60
ARC-Challenge
29.35
29.69
WinoGrande
58.41
59.43
Appendix
Table 1: Zero-shot downstream accuracy (%) of the 1.3B models, with the better result for each benchmark shown in bold.
Hyperparameter
Symbol
Value / scaling
Number of layers
NL
Ψ
Number of heads
NH
Ψ
Embedding dimension
dE
Ψ2
Query/key dimension
dK
16
Value dimension
dV
64
MiPE threshold
wth
256
Appendix
Table 2: Architectural hyperparameters of Multiscreen used in our experiments.
Quantity
4M
28M
286M
1.3B
4B
Supraparameter ( Ψ )
8
16
32
48
64
Total params
4,134,146
27,546,626
286,347,266
1,304,884,226
3,963,961,346
Non-embedding params
917,698
14,680,834
234,884,098
1,189,092,098
3,758,108,674
Appendix
Table 3: Architecture configurations and parameter counts for the Multiscreen models used in the scaling experiments.
Quantity
8M
45M
353M
1.3B
Number of layers ( NL )
6
6
24
24
Number of heads ( NH )
4
8
16
16
Embedding dimension ( dE )
128
512
1,024
2,048
Total params
7,613,312
44,609,024
353,453,056
1,310,935,040
Non-embedding params
1,180,416
18,877,440
301,989,888
1,208,008,704
Appendix
Table 4: Architecture configurations and parameter counts for the Transformer baseline models used in the scaling experiments.
Parameter
Shape
Initialization
WQ
(dE,dK)
σ=0.1/dK
WK
(dE,dK)
σ=0.1/dK
WV
(dE,dV)
σ=0.1/dV
WG
(dE,dV)
σ=0.1
WO
(dV,dE)
σ=0.1/dE
WE
(∣V∣,dE)
σ=0.1/dE
Appendix
Table 5: Parameter shapes and initialization of Multiscreen. All matrix initializations are zero-mean Gaussian, and σ denotes their standard deviation.
Figure 9: Scaling behavior under alternative definitions of model size. Left: scaling behavior of Transformer and Multiscreen with respect to non-embedding parameters. Right: scaling behavior of Multiscreen with respect to the supraparameter Ψ . Error bars indicate one standard deviation across independently trained models where multiple runs are available.
Figure 10: Validation-loss trajectories of Multiscreen models trained under the fixed 238 -token budget. Left: validation loss as a function of training tokens. Right: validation-loss reduction over the second half of training. For the 4M, 28M, and 286M models, curves and shaded regions in the left panel show the mean and one standard deviation across three independently trained models, and error bars in the right panel indicate one standard deviation across the same runs. The 1.3B and 4B results each use a single trained model.
Figure 11: Training loss trajectories from the same runs as in fig. 7 , shown for representative learning rates. Left: 45M Transformer without weight decay or gradient clipping; center: the Transformer control with weight decay 0.1 and gradient clipping at 1.0 ; right: 28M Multiscreen without either. Curves are smoothed using a moving average over 256 training steps. Divergent Transformer runs may lie entirely outside the displayed loss range.
Figure 12: Global gradient-norm dynamics during training for the 1.3B Transformer and Multiscreen models used in the scaling experiments. For visualization, values above 1 are clipped and shown with × markers.
Configuration
Outcome
Transformer
Diverged
Transformer + SIO
Diverged
Transformer + PLNorm
Diverged
Transformer + SIO − PLNorm
Diverged
Transformer + QKNorm
Diverged
Transformer + QKNorm + SIO
Stable
Appendix
Table 6: Training outcomes at learning rate 1 under controlled modifications to Transformer and Multiscreen. “Stable” indicates completion of the 214 -step diagnostic run without divergence. “SIO − PLNorm” denotes the SIO components other than PLNorm.
Figure 13: Validation-loss trajectories for five selected configurations that remain stable at learning rate 1 . PLNorm denotes pre-logit normalization, and each curve shows one 214 -step training run.
Figure 14: Global gradient norms for five selected configurations that remain stable at learning rate 1 , together with the standard 45M Transformer at learning rate 10−3 as a low-learning-rate reference. PLNorm denotes pre-logit normalization, and each panel shows one 214 -step training run. For visualization, values above 2 are clipped and shown with × markers.
Figure 15: Distance-aware relevance maps across layers and heads. Each map shows the distance-aware relevance αijd , with rows indexing query positions and columns indexing key positions. Dark gray regions indicate causally inaccessible positions or positions outside the learned screening window. Each tile is annotated with its layer and head indices, the learned screening window w , the acceptance width r , and the fraction of nonzero relevance values Pr(αijd>0) among causally accessible positions within the window, providing a summary of its sparsity and coverage.
Figure 16: Distribution of learned screening parameters in the 4B Multiscreen model. Each point corresponds to one screening unit, with horizontal position indicating the learned screening window w and vertical position indicating the learned acceptance width r . The horizontal axis is shown on a logarithmic scale. Marginal histograms show the corresponding distributions of w and r . The dotted vertical line marks the MiPE threshold ( wth=28 ), at or above which MiPE becomes inactive, and the dashed vertical line marks the pretraining context length ( 212 ), beyond which the default inference rule sets w=∞ . Raw learned windows span several orders of magnitude, with many local or medium-range units and a smaller number exceeding the pretraining context length.
Figure 17: Long-context perplexity comparison for 353M Transformer and 286M Multiscreen models. Curves show mean perplexity over three independently trained models, and shaded regions indicate one standard deviation. The horizontal axis indicates the context length at which position-dependent perplexity is evaluated, and the vertical axis is perplexity. The dashed vertical line marks the pretraining context length ( 212 ). Colored curves show Transformer with different YaRN scaling factors, while the black curve shows Multiscreen, which has no RoPE component and therefore requires no YaRN scaling. Multiscreen maintains stable perplexity substantially beyond the training context, while Transformer exhibits a tradeoff between delayed breakdown and increased perplexity at shorter contexts.
Figure 18: Passkey retrieval accuracy across model sizes for Transformer and Multiscreen. The setup follows that of ABCDigits, with identical context lengths, target depths, and evaluation protocol. Top row: Transformer baselines. Bottom row: Multiscreen models using our default inference rule. For Transformer, the YaRN scaling factor is set to the evaluation context length divided by the training context length. Multiscreen is more robust to increasing context length than Transformer, with the 286M and 1.3B Multiscreen models achieving 100% accuracy across the grid.
Variant
Pos. enc.
Context
Scaling
# Params (M)
Softmax
RoPE
Full
YaRN
353.45
Softmax 1
RoPE
Full
YaRN
353.45
SSMax
RoPE
Full
YaRN ∗
353.45
Entmax ( α=1.5 )
RoPE
Full
YaRN ∗
353.45
Sigmoid
RoPE
Full
YaRN ∗ + method-specific
353.45
ReLU2
RoPE
Full
YaRN ∗ + method-specific
353.55
Appendix
Table 7: Configurations in the controlled attention-mechanism comparison. All listed variants use the same query, key, and value dimensions. The Context column describes formal support, so mechanisms with distance-dependent attenuation but no hard support boundary are classified as Full. Parameter counts are reported in millions and rounded to two decimal places.
Variant
Variable total mass
Pairwise independent
Bounded
Exact rejection
Empty selection
No explicit length term
No inference scaling
Softmax
–
–
✓
–
–
✓
–
Softmax 1
✓
–
✓
–
–
✓
–
SSMax
–
–
✓
–
–
–
–
Entmax ( α=1.5 )
–
–
✓
✓
–
✓
–
Sigmoid
✓
✓
✓
–
–
–
–
ReLU2
✓
✓
–
✓
✓
–
–
Appendix
Table 8: Properties relevant to absolute relevance across context lengths for the mechanisms evaluated in the controlled comparison. “Bounded” indicates that each relevance value has a finite upper bound. Structural zeros introduced by causal or fixed-window masks are excluded from “Exact rejection.”
Figure 19: Long-context perplexity for 286M Multiscreen with and without TanhNorm under the default inference rule and with as-learned w . For the TanhNorm baseline, curves show the mean across three independently trained runs, and shaded regions indicate one standard deviation; the variant without TanhNorm uses a single run. The gray dashed line marks the pretraining context length ( 212 ).
Figure 20: ABCDigits retrieval accuracy for the same 286M Multiscreen variants and inference-window treatments as in fig. 19 . Results with TanhNorm are averaged over three independently trained runs, whereas results without TanhNorm use a single run. The default rule applies the post-training conversion to w=∞ , whereas the as-learned condition leaves every learned window unchanged. The gray dashed lines mark the pretraining context length ( 212 ).
Figure 21: Long-context perplexity for 286M Multiscreen using MiPE, NoPE, or RoPE. The MiPE curve shows the mean across three independently trained runs, and its shaded region, largely obscured by the mean curve, indicates one standard deviation; the NoPE and RoPE curves each use a single run. MiPE and NoPE use no inference-time positional scaling; the colored RoPE curves use the indicated fixed YaRN scaling factors. The gray dashed line marks the pretraining context length ( 212 ).
Figure 22: ABCDigits retrieval accuracy for the same 286M Multiscreen positional-encoding variants. The first panel applies context-matched YaRN scaling to RoPE, whereas the second uses unscaled RoPE; MiPE and NoPE require no inference-time positional scaling. The MiPE results are averaged over three independently trained runs, whereas the NoPE and RoPE results each use a single run. The gray dashed lines mark the pretraining context length ( 212 ).
Figure 23: Validation loss over 216 training steps ( 236 tokens) for standard Multiscreen and the closely capacity-matched variant without the output gate. The Multiscreen curve averages three independently trained runs, and its shaded region, largely obscured by the mean curve, indicates one standard deviation; the no-gate curve uses a single run.
Figure 24: Long-context perplexity for standard Multiscreen and the closely capacity-matched variant without the output gate. The standard Multiscreen curve shows the mean across three independently trained runs, and its shaded region, largely obscured by the mean curve, indicates one standard deviation; the no-gate curve uses a single run. The gray dashed line marks the pretraining context length ( 212 ).
Figure 25: ABCDigits retrieval accuracy for standard Multiscreen and the closely capacity-matched variant without the output gate. Results for standard Multiscreen are averaged over three independently trained runs, whereas the no-gate results use a single run. The gray dashed lines mark the pretraining context length ( 212 ).
Long contexts have become standard in pretrained LLMs, yet they remain expensive to run: prefill compute grows quadratically with sequence length, and every decode step re-reads a key-value cache that grows linearly with it. Sparse attention cuts these costs by attending only to a relevant subset of past tokens, but selecting that subset is itself expensive. We present SpotAttention, a lightweight selector that attaches to a frozen pretrained transformer and learns by KL distillation to estimate its attention distribution. The selector picks the top-K keys each query attends to, and because its estimate is a calibrated distribution, a dual top-p rule reads the per-query, per-layer budget directly from it. Across Qwen3 (dense, 4B-32B) and Qwen3.5 (hybrid linear/full attention, 4B-9B), SpotAttention matches dense accuracy at contexts up to 128K tokens, eight times the training length. Decode at L=128K runs 3.9x faster than FlashAttention and 1.8x faster than Twilight, the strongest training-free baseline. Quantizing the selector's K-cache to INT4 or FP4 microscale shrinks it 3.5x at no accuracy cost.
The quadratic computational cost of traditional attention mechanisms poses a major bottleneck to the scalability and practical deployment of large language models (LLMs), particularly in long-context scenarios. To improve efficiency, existing approaches often enforce rigid structural constraints such as local attention windows. However, these strategies typically lead to substantial performance degradation on tasks requiring precise long-range recall. In this work, we propose MATCH, a scalable and efficient framework that augments sparsified attention mechanisms with dynamically integrated in-context information through an efficient retrieval system. Empirical results show that MATCH significantly improves the performance of sparse-attention models on both synthetic and real-world natural-language tasks. These findings highlight the versatility of MATCH as a general approach for enhancing in-context retrieval capabilities while maintaining the efficiency benefits of sparse attention architectures.
Linrui Ma, Chun Hei Lo, Xinyu Wang +12
Huawei Canada · McGill University · Université de Montréal +1
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.
Xingyue Huang, Xueying Ding, Mingxuan Ju +3
University of Oxford · Carnegie Mellon University · Snap Inc.