cs.LGApr 1, 2026
SaveScreening Is Enough
Organizations: Center for Emergent Matter Science (CEMS), RIKEN · Graduate School of Science, The University of Tokyo
Abstract
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 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 residual layers with parallel gated screening tiles per layer; is the row-wise unit-normalized embedding matrix shared by the input and language-modeling head, and and 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 . (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 , becoming inactive when .
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 ( ).
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 ( ).
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 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 , scaling factors and above exceed the plotted perplexity range. Dashed lines mark the pretraining context length ( ).
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 , which reuses the three runs from the scaling experiments and shows their mean and standard deviation. (b) Training outcomes at learning rate after tokens ( 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 tokens. (c) Model-side TTFT across model sizes at context length . (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 |
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 | ||
| Number of heads | ||
| Embedding dimension | ||
| Query/key dimension | 16 | |
| Value dimension | 64 | |
| MiPE threshold | 256 |
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 |
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 ( ) | 6 | 6 | 24 | 24 |
| Number of heads ( ) | 4 | 8 | 16 | 16 |
| Embedding dimension ( ) | 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 |
Table 4: Architecture configurations and parameter counts for the Transformer baseline models used in the scaling experiments.
| Parameter | Shape | Initialization |
|---|---|---|
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 -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 and gradient clipping at ; 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 |
Table 6: Training outcomes at learning rate under controlled modifications to Transformer and Multiscreen. “Stable” indicates completion of the -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 . PLNorm denotes pre-logit normalization, and each curve shows one -step training run.
Figure 14: Global gradient norms for five selected configurations that remain stable at learning rate , together with the standard 45M Transformer at learning rate as a low-learning-rate reference. PLNorm denotes pre-logit normalization, and each panel shows one -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 , 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 , the acceptance width , and the fraction of nonzero relevance values 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 and vertical position indicating the learned acceptance width . The horizontal axis is shown on a logarithmic scale. Marginal histograms show the corresponding distributions of and . The dotted vertical line marks the MiPE threshold ( ), at or above which MiPE becomes inactive, and the dashed vertical line marks the pretraining context length ( ), beyond which the default inference rule sets . 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 ( ). 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 accuracy across the grid.
| Variant | Pos. enc. | Context | Scaling | # Params (M) |
|---|---|---|---|---|
| Softmax | RoPE | Full | YaRN | |
| Softmax 1 | RoPE | Full | YaRN | |
| SSMax | RoPE | Full | YaRN ∗ | |
| Entmax ( ) | RoPE | Full | YaRN ∗ | |
| Sigmoid | RoPE | Full | YaRN ∗ + method-specific | |
| RoPE | Full | YaRN ∗ + method-specific |
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 ( ) | – | – | – | – | |||
| Sigmoid | – | – | – | – | |||
| – | – | – |
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 . 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 ( ).
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 , whereas the as-learned condition leaves every learned window unchanged. The gray dashed lines mark the pretraining context length ( ).
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 ( ).
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 ( ).
Figure 23: Validation loss over training steps ( 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 ( ).
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 ( ).