Length generalization is the ability of sequential models to perform well on context lengths unseen during training. In this work, we show that the challenge of achieving length generalization is related not only to architectural choices such as positional encoding and the attention mechanism but also to the training procedure itself. We study how regularization affects length generalization and find that weight decay can hinder extrapolation. In contrast, dropout improves extrapolation when its placement within the architecture is reconsidered. In that regard, we show that the standard placement of dropout before layer normalization introduces a systematic distributional mismatch, and that applying dropout just before the linear projection resolves this issue. For example, a modified SmolLM3 with sliding window attention, continually pre-trained with dropout, can extrapolate perfectly to 64× on Needle-in-a-Haystack and far beyond the pre-training context size on RULER and HELMET. Mamba2 also benefits from dropout, suggesting an architecture-agnostic nature of the problem. We further propose Variance-Preserving Affine Dropout (VPAD), a new dropout strategy that substantially reduces the resulting pre-activation variance mismatch, leading to further extrapolation improvements in transformers.
Figures & tables
Figure 1: Transformer performance across training steps. Left : per-position CE loss (512-window moving average); the dotted line denotes the pre-training context length (2K). Right : RULER extrapolation, RBA averaged over 4–8K. Training with WD leads to overfitting to the pre-training context length. In contrast, dropout (VPAD) substantially mitigates the issue.
Figure 2: Dropout placements. Left : Pre-Norm dropout (Eq. 6 ). Right : Pre-Linear dropout. r — residual stream, u — output of a sublayer (e.g. Attn , FFN , Mamba2 ), Linear — linear projection.
Figure 3: Comparison of the relative error ∣Var[yDropout]−Var[ytest]∣/Var[ytest] of the pre-activation variances across layers of SmolLM3 fine-tuned with Left : Pre-Linear Dropout; Right : VPAD. VPAD substantially reduces the error. Moreover, the error decreases with depth for VPAD but grows for PL-D (first/last 18 layers, averaged over matrices and layers: 0.09/0.06 vs. 0.29/0.37).
Figure 4: Results for transformers trained from scratch, evaluated on Left: ArXiv papers; Middle: NIAH-1; Right: RULER. The dotted lines represent the pre-trained context length.
NIAH-1
RULER
Model
4K
8K
16K
32-128K
256K
4K
8K
16K
32K
Orig.
100.0
0.0
0.0
— 0.0 —
0.0
83.5
0.1
0.0
0.0
SW
100.0
0.0
0.0
— 0.0 —
0.0
83.4
0.5
0.0
0.0
CPT
SW, WD
100.0
0.0
0.0
— 0.0 —
0.0
83.5
1.1
0.0
0.0
SW, no-WD
100.0
0.0
0.0
— 0.0 —
0.0
83.5
1.1
0.0
0.0
Table 1: SmolLM3 performance on NIAH-1 and on RULER. Grey — the ID length. Bold — best.
Table 6Table 7
NIAH-1
RULER
Model
2K
4K
8K
16K
32K
2K
4K
8K
16K
32K
Orig.
100.0
100.0
100.0
0.0
0.0
38.3
14.3
8.5
0.4
0.0
CPT
WD
100.0
100.0
100.0
0.0
0.0
42.0
16.2
9.6
0.7
0.0
no-WD
100.0
100.0
100.0
0.0
0.0
42.2
16.0
9.6
0.0
0.0
PN-D, p=0.2
100.0
100.0
100.0
11.8
0.0
50.7
22.5
11.5
2.6
0.8
Table 5: Mamba2-1.3B performance on NIAH-1 and on RULER. Bold — best.
Appendix figures & tables10 assets
Supplementary material from the paper’s appendix.
Appendix
Model
p
Ls
Dim
FFN-D
H
Len
Data
TT
BT
LR
LRS
Warm-up
Ann-R
min-LR
Trained from scratch
Transformers, NAPE
0.1
24
768
3072
12
2048
DCLM-Edu
40B
512K
6⋅10−4
SWUA
2.5%
20%
6⋅10−5
Mamba2-85M, State=128
0.1
12
768
-
24
1024
The Pile
25B
512K
1⋅10−3
Cosine
10%
-
1⋅10−4
Mamba2-45M, State=128
0.1
12
768
-
24
50-100
Copy
0.3B
19K
1⋅10−3
Cosine
5%
-
0
CPT
SmolLM3, SW
0.1-0.3
36
2048
11008
16
4096
DCLM-Edu
1B
1M
2⋅10−5
SWU
10%
-
2⋅10−5
Appendix
Table 6: Hyperparameters of the trained models. Ls — number of layers, Dim — model’s dimension, FFN-D — hidden dimension of FFN sublayers, H — number of heads, Len — context length, TT — training tokens, BT — tokens per batch, LR — learning rate, LRS — LR scheduler, Ann-R — annealing ratio, min-LR — annealing min LR. LR schedulers: Stable with Warm-Up and Annealing (SWUA), Stable with Warm-Up (SWU).
Model
p
ARC-E
ARC-C
HS
PIQA
WG
OPQA
BoolQ
LMBD
Avg.
Original
-
77.8
53.2
75.3
79.0
69.6
44.2
80.4
70.8
68.8
CPT
WD
-
77.7
53.2
75.5
79.2
68.9
43.4
79.9
70.5
68.5
no-WD
-
77.6
53.3
75.4
79.1
69.6
43.6
80.0
70.6
68.7
PN-D
0.1
78.3
53.2
74.7
79.2
69.4
45.6
78.3
69.7
68.5
PL-D
0.1
79.0
54.3
74.9
79.4
69.1
44.6
81.0
70.4
69.1
Appendix
Table 7: SmolLM3-3B results on short benchmarks. HS - HellaSwag, WG - Winogrande, LMBD - Lambada, OPQA - OpenBookQA
NIAH-1
RULER
Model
4K
8K
16K
32-128K
256K
4K
8K
16K
32K
Orig.
100.0
0.0
0.0
— 0.0 —
0.0
83.5
0.1
0.0
0.0
SW
100.0
0.0
0.0
— 0.0 —
0.0
83.4
0.5
0.0
0.0
CPT
SW, WD
100.0
0.0
0.0
— 0.0 —
0.0
83.5
1.1
0.0
0.0
SW, no-WD
100.0
0.0
0.0
— 0.0 —
0.0
83.5
1.1
0.0
0.0
Appendix
Table 8: SmolLM3 performance on NIAH-1 and on RULER for all dropout probabilities. Notation as in Table 1 . “— V —” means that all entries are equal to V in the range; V1→V2 means that the left bound of the range has value V1 and the right one V2 . Grey columns — ID length.
Figure 6: Performance of models trained from scratch before and after applying checkpoint averaging. Left: RULER: in-distribution length (2K); Middle: RULER: out-of-distribution—average of performance at 4K and 8K; Right: Cross-Entropy loss on ArXiv papers from The Common Pile.
Model
4K
8K
16K
Orig.
100.0
0.0
0.0
Checkpoint Avg.
SW, K=10
100.0
41.2
0.8
SW, K=20
97.2
66.4
11.4
SW, K=30
80.4
68.2
2.2
CPT: Attention Dropout
Appendix
Table 9: Performance on NIAH-1 of SmolLM3 fine-tuned with attention dropout
Model
4K
8K
16K
Orig.
83.5
0.1
0.0
Checkpoint Avg.
SW, K=10
79.4
11.4
0.7
SW, K=20
61.6
13.8
1.6
SW, K=30
43.2
10.7
1.7
CPT: Attention Dropout
Appendix
Table 10: Performance on RULER of SmolLM3 fine-tuned with attention dropout
Model
2K
4K
8K
16K
32K
Softmax WD
40.39
3.09
0.06
0.04
0.04
Softmax VPAD
44.24
34.76
12.30
4.54
1.93
ASEntmax WD
43.96
20.94
5.44
2.97
2.44
ASEntmax PN-D
42.76
26.85
7.72
2.36
1.59
ASEntmax PL-D
47.30
36.80
20.11
11.39
6.73
ASEntmax VPAD
48.09
37.32
21.64
13.51
6.98
Appendix
Table 11: Performance of ASEntmax models on RULER
Model
4K
8K
16K
32K
64K
128K
256K
SmolLM3
100.0
0.0
0.0
0.0
0.0
0.0
0.0
SmolLM3-SW
100.0
0.0
0.0
0.0
0.0
0.0
0.0
Checkpoint Averaging
SmolLM3-SW, K=10
100.0
41.2
0.8
0.0
0.0
0.0
0.0
SmolLM3-SW, K=20
97.20
66.4
11.4
0.0
0.0
0.0
0.0
SmolLM3-SW, K=30
80.40
68.2
2.2
0.0
0.0
0.0
0.0
Appendix
Table 12: Full SmolLM3 results on NIAH-1.
Table 17
NIAH-1
RULER
Model
2K
4K
8K
16K
32K
2K
4K
8K
16K
32K
Orig.
100.0
100.0
100.0
0.0
0.0
38.3
14.3
8.5
0.4
0.0
CPT
WD
100.0
100.0
100.0
0.0
0.0
42.0
16.2
9.6
0.7
0.0
no-WD
100.0
100.0
100.0
0.0
0.0
42.2
16.0
9.6
0.0
0.0
PN-D p=0.2
100.0
100.0
100.0
11.8
0.0
50.7
22.5
11.5
2.6
0.8
Appendix
Table 15: Mamba2-1.3B performance on NIAH-1 and on RULER.
Transformers with relative positional encodings often extrapolate to sequences longer than those seen during training, whereas transformers with learned absolute encodings typically do not. This is a robust empirical regularity, and the explanations offered for it so far are chiefly about expressivity, that is, about whether a length-generalizing solution exists. We give an optimization explanation. On a minimal fixed-offset retrieval task that isolates positional selection, the gap is governed by the implicit bias of the trained attention head: among the many solutions that fit short sequences, which one gradient descent actually selects. We prove that rotary encodings make the attention logit a function of relative offset alone, an exact equivariance, so whatever selection rule is learned at training lengths is reproduced verbatim at every longer length. Learned absolute encodings instead leave out-of-range positions unconstrained, and the trained head pins to a fixed absolute position inside the training range. We characterize the learned rotary rule as a low-rank ``carrier'' kernel aligned with the target offset, and we derive the resulting graceful accuracy decay as an attention-dilution law; both predictions are confirmed across seeds and offsets. A linear-attention control shows the mechanism is specific to softmax: without normalization, training selects a min-norm interpolant that does not extrapolate. The phenomenon, the equivariance, and the carrier all transfer to a multi-layer, multi-head transformer trained on a full-sequence length-generalization task. The account connects the implicit bias of attention, implicit bias for extrapolation in recurrent models, and the learning side of the RASP-L conjecture.
Subham Singh, Ashutosh Mishra, Subha Raut
Department of Mathematics and Statistics, Mississippi State University · Independent Researcher
Large language models (LLMs) are typically pretrained on short sequences and then extended to work on longer sequences with additional training. However, such LLMs still struggle to further generalize to very long sequences. We propose Randomized YaRN, a training method that improves length generalization by combining YaRN-based positional extrapolation with randomized positional encoding and a length curriculum. During training on short context data, tokens are assigned YaRN positional encodings sampled from a larger position range, exposing the model to out-of-distribution positional representations even on short-context inputs. We evaluate Randomized YaRN on three challenging long-context reasoning benchmarks, BABILong, Multi-Round Coreference Resolution (MRCR), and LongBench v2. When training on data with short context, Randomized YaRN consistently improves reasoning performance on context lengths from 16K to 128K and outperforms standard fine-tuning, with the largest gains appearing at far out-of-distribution lengths. Our results suggest that progressively exposing models to OOD positional distributions provides an effective recipe for generalizable long-context reasoning.
Out-of-distribution length generalization, namely to extrapolate a task from short to longer context, has been studied intensively for transformers. Here we focus on distance generalization, which probes performance when inter-token distances are changed between training and inference, while keeping a fixed context length. We construct two synthetic delay copy tasks, both involving finite distances between source and recall, where tokens are copied either fully or selectively, and test models on delays unseen during training. We address three questions: (A) Do positional encoding schemes such as RoPE and ALiBi improve distance resolution relative to no positional encoding (NoPE)? (B) How does data diversity, the number of inter-token distances seen in training, affect performance? (C) When is distance transfer learning positive or negative? We present a thorough investigation, finding that it is paramount to improve our understanding of the underlying mechanisms.
Daniel Henrik Nevermann, Claudius Gros
Institute for Theoretical Physics, Goethe University Frankfurt, Germany