Length Generalization Needs Proper Regularization

Pavlo Vasylenko; Matthias Lindemann; André F. T. Martins; Marcos Treviso

cs.LGOct 3, 2026

Length Generalization Needs Proper Regularization

Authors: Pavlo Vasylenko, Matthias Lindemann, André F. T. Martins, Marcos Treviso

Organizations: Instituto Superior Técnico, Universidade de Lisboa · Instituto de Telecomunicações · ELLIS Unit Lisbon · Gandara AI

Abstract

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 $\times$ 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.

\bm{r}

— residual stream,

\bm{u}

— output of a sublayer (e.g.

\mathrm{Attn}

\mathrm{FFN}

\mathrm{Mamba2}

), Linear — linear projection.

Figure 3: Comparison of the relative error $|\mathrm{Var}[y^{\mathrm{Dropout}}]-\mathrm{Var}[y^{\mathrm{test}}]|/\mathrm{Var}[y^{\mathrm{test}}]$ 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 6

	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\cdot 10^{-4}$	SWUA	$2.5\%$	$20\%$	$6\cdot 10^{-5}$
Mamba2-85M, State=128	0.1	12	768	-	24	1024	The Pile	25B	512K	$1\cdot 10^{-3}$	Cosine	$10\%$	-	$1\cdot 10^{-4}$
Mamba2-45M, State=128	0.1	12	768	-	24	50-100	Copy	0.3B	19K	$1\cdot 10^{-3}$	Cosine	$5\%$	-	0
CPT
SmolLM3, SW	0.1-0.3	36	2048	11008	16	4096	DCLM-Edu	1B	1M	$2\cdot 10^{-5}$	SWU	$10\%$	-	$2\cdot 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;

V_{1}\rightarrow V_{2}

means that the left bound of the range has value

V_{1}

and the right one

V_{2}

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

Explore similar work

Relative Positions Generalize, Absolute Positions Memorize: An Implicit-Bias Account of Length Generalization in Attention
Jul 21, 2026Subham Singh, Ashutosh Mishra, Subha RautSoftmax Attention Rotary Positional Embeddings
Department of Mathematics and Statistics, Mississippi State University · Independent Researcher
Randomized YaRN Improves Length Generalization for Long-Context Reasoning
Jun 22, 2026Manas Mehta, Fangcong Yin, Greg DurrettLong-Context Language Modeling LLM Fine-Tuning
New York University
Distance generalization in transformers: why bother with positional encoding?
Sep 10, 2026Daniel Henrik Nevermann, Claudius GrosOOD Generalization Positional Encoding
Institute for Theoretical Physics, Goethe University Frankfurt, Germany

Jul 21, 2026J/K move · Enter open · S save

Relative Positions Generalize, Absolute Positions Memorize: An Implicit-Bias Account of Length Generalization in Attention

Subham Singh, Ashutosh Mishra, Subha Raut

Department of Mathematics and Statistics, Mississippi State University · Independent Researcher

Softmax Attention Rotary Positional Embeddings Length Generalization Implicit Bias

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Relative Positions Generalize, Absolute Positions Memorize: An Implicit-Bias Account of Length Generalization in Attention

Randomized YaRN Improves Length Generalization for Long-Context Reasoning

Distance generalization in transformers: why bother with positional encoding?