On Vanishing Variance in Transformer Length Generalization

Abstract

It is a widely known issue that Transformers, when trained on shorter sequences, fail to generalize robustly to longer ones at test time. This raises the question of whether Transformer models are real reasoning engines, despite their impressive abilities in mathematical problem solving and code synthesis. In this paper, we offer a vanishing variance perspective on this issue. To the best of our knowledge, we are the first to demonstrate that even for today’s frontier models, a longer sequence length results in a decrease in variance in the output of the multi-head attention modules. On the $\operatorname{argmax}$ retrieval and dictionary lookup tasks, our experiments show that applying layer normalization after the attention outputs leads to significantly better length generalization. Our analyses attribute this improvement to a reduction—though not a complete elimination—of the distribution shift caused by vanishing variance. Project page: ruiningli.com/vanishing-variance.

Figure 1: Standard deviation of a fixed component in attention outputs from the first layer of Llama-3.2-1B (log-log scale) over multiple input sequences of fixed length $N$. Even in the latest LLMs, increasing sequence length $N$ reduces the variance of attended outputs, significantly degrading accuracy on long sequences.

1 Background: Vanishing Variance

It is no exaggeration to say that Transformers  (Vaswani et al., 2017) is the most important architecture in modern deep learning. It is widely adopted in almost every domain, ranging from natural language (Vaswani et al., 2017; Devlin et al., 2019; Brown et al., 2020) and vision (Dosovitskiy et al., 2021; Peebles & Xie, 2023) to audio (Radford et al., 2023) and protein design (Jumper et al., 2021). Despite its successes, recent studies (Press et al., 2022; Zhou et al., 2023; 2024; Kazemnejad et al., 2024; Veličković et al., 2024) in large language models (LLMs) have shown that transformer-based models often struggle with length generalization, an ability that requires the model to generalize to longer sequences than seen during training. Several prior works have proposed to either refine position encodings (Ruoss et al., 2023; Zhou et al., 2024; Kazemnejad et al., 2024) or adapt the softmax function (Press et al., 2022; Veličković et al., 2024) to improve length generalization. However, these methods are ad-hoc and lack interpretability, making it more of an art than a science to understand when and why they work.

In this paper, we study the distribution shift that occurs in the intermediate outputs when an attention module trained on shorter sequences is subsequently exposed to longer ones in a zero-shot manner. We hope that our findings will encourage future research on network architectures that are provably robust (e.g., invariant) to varying sequence lengths.

Background and notations.

At the core of Transformers is the attention mechanism  (Vaswani et al., 2017). The attention first projects the input sequence $\mathbf{X}=\left[\mathbf{x}_{1}\|\mathbf{x}_{2}\|\dots\|\mathbf{x}_{N}\right]^{\top}\in\mathbb{R}^{N\times D}$, where $N$ is the sequence length and each item $\mathbf{x}_{n}\in\mathbb{R}^{D}$ has $D$ features, into keys $\mathbf{K}=\mathbf{X}\mathbf{W}_{K}\in\mathbb{R}^{N\times D}$ and values $\mathbf{V}=\mathbf{X}\mathbf{W}_{V}\in\mathbb{R}^{N\times D}$ using learnable weight matrices $\mathbf{W}_{K},\mathbf{W}_{V}\in\mathbb{R}^{D\times D}$. Similarly, the query sequence $\mathbf{Y}=\left[\mathbf{y}_{1}\|\mathbf{y}_{2}\|\dots\|\mathbf{y}_{M}\right]^{\top}\in\mathbb{R}^{M\times D}$ is projected into queries $\mathbf{Q}=\mathbf{Y}\mathbf{W}_{Q}\in\mathbb{R}^{M\times D}$ using $\mathbf{W}_{Q}\in\mathbb{R}^{D\times D}$. The attention then computes $\mathbf{O}=\operatorname{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^{\top}}{\sqrt{D}}\right)\mathbf{V}\in\mathbb{R}^{M\times D}$ and projects it using another weight matrix $\mathbf{W}_{O}\in\mathbb{R}^{D\times D}$ to yield the final result $\operatorname{Attn}(\mathbf{X},\mathbf{Y})=\mathbf{O}\mathbf{W}_{O}$. In this paper, we use the term “attention weights” to refer to the softmax score, i.e., $\operatorname{softmax}(\frac{\mathbf{Q}\mathbf{K}^{\top}}{\sqrt{D}})$, “attention outputs” the intermediate $\mathbf{O}$, and $\operatorname{Attn}(\mathbf{X},\mathbf{Y})$ the final result.

Our main observation is the vanishing variance problem: as the sequence length $N$ increases, the variance of attention outputs (computed over multiple input sequences of length $N$) decreases. We formalize this as Proposition 1.

Proposition 1 (The vanishing variance problem).

Consider a trained attention module with weights $\mathbf{W}_{Q},\mathbf{W}_{K},\mathbf{W}_{V},\mathbf{W}_{O}$. Let $\mathbf{X}=\left[\mathbf{x}_{1}\|\mathbf{x}_{2}\|\dots\|\mathbf{x}_{N}\right]^{\top}$ denote an input sequence of length $N$. If (1) $\mathbf{x}_{1},\mathbf{x}_{2},\dots,\mathbf{x}_{N}\overset{\text{i.i.d}}{\sim}\mathcal{X}$, a distribution over a finite vocabulary, and (2) $\mathbb{E}_{\mathbf{x}\sim\mathcal{X}}[\mathbf{W}_{V}\mathbf{x}]=\mathbf{0}$, then for a fixed query $\mathbf{y}$ and a fixed feature $d$,

$$\lim_{N\to\infty}\operatorname{Var}_{(\mathbf{x}_{1},\mathbf{x}_{2},\ldots,\mathbf{x}_{N})\sim\mathcal{X}^{N}}\left(\left[\operatorname{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^{\top}}{\sqrt{D}}\right)\mathbf{V}\right]_{d}\right)=0,$$

where $\mathbf{x}_{n},\mathbf{y}\in\mathbb{R}^{D}$ and $\mathbf{Q}\in\mathbb{R}^{1\times D},\mathbf{K}\in\mathbb{R}^{N\times D},\mathbf{V}\in\mathbb{R}^{N\times D}$ are intermediate results in $\operatorname{Attn}(\mathbf{X},[\mathbf{y}])$.

Informally, for a fixed component $d$ in the attention outputs, its variance over input sequences of length $N$, where each sequence consists of $N$ independently and identically distributed (i.i.d.) tokens, vanishes as $N\to\infty$.

Proof.

Please refer to Appendix A. ∎

Note that assumptions of Proposition 1 are violated in practice. In particular, the independence assumption does not hold in LLMs because of (1) the introduction of positional encoding, and more significantly (2) the nature of language, where preceding words provide important context for those that follow. In addition, $\mathbb{E}[\mathbf{W}_{V}\mathbf{x}_{i}]=\mathbf{0}$ is not strictly enforced. Nevertheless, we find that even for today’s frontier LLMs, the decay in attention output variance, as established in Proposition 1, remains pronounced. In Fig. 1, we plot the standard deviation $\sigma$ of a fixed component of the attention outputs from the first layer of Llama-3.2-1B (AI@Meta, 2024) as a function of input sequence length $N$. $\sigma$ is computed over $100$ length-$N$ sequences sampled randomly with $3$ strategies: ( Random Tokens w/o P.E.) We sample single tokens i.i.d uniformly at random from the tokenizer’s vocabulary, and remove the positional encoding for inference; ( Random Tokens w/ P.E.) We still sample single tokens i.i.d uniformly at random, but keep the positional encoding at inference time; ( Sentences w/ P.E.) We sample consecutive sentences from a long paragraph, and truncate the token sequences to length $N$—such sequences lie within the LLM’s training distribution. As can be seen in the log-log plot, for Random Tokens w/o P.E., where the independence assumption does hold, $\sigma$ scales with input sequence length $N$ roughly as $\sigma\propto N^{-0.5}$. For Random Tokens w/ P.E. and Sentences w/ P.E., where such assumption is no longer valid, the downward trend is still obvious.

2 Layer Normalization for Length Generalization

As variance vanishes with longer sequence lengths in attention outputs, we are intrigued to investigate the causes of performance degradation observed in LLMs. To this end, we perform a toy study on the statistical behavior of attention output values.

For simplicity, we consider a one-layer Transformer with single-head attention, omitting residual connections and normalization, following Veličković et al. (2024). We adopt this architecture as our Baseline. The model receives a single query token and an input sequence of varying length to perform simple algorithmic tasks. To eliminate confounds from positional encodings, we focus on order-invariant tasks, where the output depends only on the multiset (not the order) of input tokens, including $\operatorname{argmax}$ retrieval and dictionary lookup. Our goal is to evaluate models trained on shorter sequences using longer (i.e., out-of-distribution in length) sequences to study length generalization. More details of the model architecture and synthetic data generation are provided in Appendix B.

In Fig. 2, we visualize the distribution of $5$ individual components in attention outputs $\mathbf{O}$ across multiple input sequences of lengths $2^{4}$, $2^{12}$ and $2^{14}$, obtained with a model checkpoint trained on sequences of up to length $2^{4}$. As can be seen in the top row, testing on out-of-distribution sequence lengths leads to vanishing variance, causing a distribution shift where each individual component of $\mathbf{O}$ becomes more concentrated around its mean.

Figure 2: Distribution of 5 individual features in attention outputs $\mathbf{O}$ across batches. Each color represents a different feature. As input sequence length $N$ increases from $2^{4}$ to $2^{14}$, feature variance decreases, and values concentrate around their mean. Layer normalization (bottom) scales and shifts features to maintain relatively constant global variance, likely explaining its superior length generalization compared to the Baseline (top).

While this distribution shift of individual features is expected according to Proposition 1, we are more interested in the distribution shift of the entire feature vector in $\mathbb{R}^{D}$, as the whole vector is subsequently input to an MLP to predict the final result. As input sequence length $N$ increases, each feature is less likely to have extreme values (as its distribution is more centered). Consequently, the global feature variance, defined as $\sigma_{\text{global}}^{2}=\frac{1}{D}\Sigma_{d=1}^{D}(\mathbf{o}_{d}-\mu_{\text{global}})^{2}$ where $\mu_{\text{global}}=\frac{1}{D}\Sigma_{d=1}^{D}\mathbf{o}_{d}$ is the global mean, also decreases. We illustrate this observation in Fig. 3 (right), where the global variance decays as $N$ increases. In Fig. 3 (left), we show that in addition to the global variance, the global mean $\mu_{\text{global}}$ also exhibits drift. Such a distribution shift in attention outputs (and thus MLP inputs) hinders generalization, since the MLP is only trained on features with larger global variance and a different global mean (Zhou et al., 2022).

Figure 3: Layer normalization helps mitigate distribution shift in attention outputs. (Left) shows the drift in global mean as input sequence length deviates from the training distribution. The mean is normalized by the training global variance to eliminate scale differences. (Right) shows the decay in global variance. All results are averaged across $32$k random input sequences of the fixed length.

To mitigate this distribution shift, we explore applying layer normalization (Ba et al., 2016) immediately after the attention outputs, i.e., $\operatorname{LayerNorm}(\mathbf{O})_{b,t,d}=\gamma_{d}\cdot\frac{\mathbf{o}_{b,t,d}-\mathbf{\mu}_{b,t}}{\mathbf{\sigma}_{b,t}+\epsilon}+\beta_{d}$, where $\mathbf{O}$ is the batched attention outputs, $\mu_{b,t}=\frac{1}{D}\Sigma_{d=1}^{D}\mathbf{O}_{b,t,d}$, $\sigma_{b,t}=\sqrt{\frac{1}{D}\Sigma_{d=1}^{D}(\mathbf{O}_{b,t,d}-\mu_{b,t})^{2}}$, and $\gamma,\beta\in\mathbb{R}^{D}$ are learnable scale and shift parameters. While variance decay in individual features is inevitable (bottom row of Fig. 2), standardization and learnable scale and shift parameters help stabilize the feature distribution. This adjustment preserves the global mean and variance more effectively as sequence length increases $1000\times$ (Fig. 3). This enhances length generalization, as discussed next.

3 Experiments

We consider two tasks: $\operatorname{argmax}$ retrieval and dictionary lookup. The former has been considered by Veličković et al. (2024). The latter closely resembles the core function of the attention mechanism (i.e., to retrieve the most relevant information based on the similarity between queries and keys). As detailed in Appendix B, the order of input tokens does not affect the target output in either task. By deliberately selecting such tasks, we isolate and examine the length generalization capabilities of the attention mechanism itself, independent of any effects introduced by positional encodings (Zhou et al., 2024). We generate synthetic data (of input sequence length up to $16$) to train the models, and evaluate them on sequences of length up to $2^{14}$.

3.1 Results and Analysis

The results presented in Table 1 and Table 2 indicate that applying layer normalization to attention outputs leads to consistently better accuracy on out-of-distribution sequence lengths, with statistical significance confirmed by a paired $t$-test over $100$ training runs from different random seeds.

Test-time adaptation and fine-tuning are common techniques for improving length generalization in transformers (Anil et al., 2022; Veličković et al., 2024). To show that the benefits of layer normalization are orthogonal to these techniques, we implement the adaptive temperature method from Veličković et al. (2024) in both architectures, with and without layer normalization. Combined with test-time adaptation, layer normalization still yields a significant improvement. In Fig. 4, we demonstrate that layer normalization also mitigates dispersion (Veličković et al., 2024).

Figure 4: Heatmap of the largest $16$ attention weights, computed over $32$ examples. Layer normalization mitigates dispersion, which is inevitable as sequence length increases (Veličković et al., 2024).

Table 1: Results (%) on the $\operatorname{argmax}$ retrieval task. Results are averaged over $100$ runs with different random seeds. $p$-values are computed using a paired $t$-test. Entries highlighted in green indicate those with in-distribution sequence lengths.

Model	$2^{4}$	$2^{5}$	$2^{6}$	$2^{7}$	$2^{8}$	$2^{9}$	$2^{10}$	$2^{11}$	$2^{12}$	$2^{13}$	$2^{14}$
w.o. test-time adaptation
Baseline	$\mathbf{99.6}$	$99.2$	${98.4}$	$96.8$	$93.7$	$88.0$	$78.0$	$62.5$	$44.2$	$29.7$	$20.8$
Baseline (+ LN)	$\mathbf{99.6}$	$\mathbf{99.3}$	$\mathbf{98.6}$	$\mathbf{97.4}$	$\mathbf{94.8}$	$\mathbf{89.8}$	$\mathbf{81.0}$	$\mathbf{66.9}$	$\mathbf{49.2}$	$\mathbf{33.0}$	$\mathbf{22.6}$
$p$-value	$8/10^{1}$	$4/10^{2}$	$4/10^{2}$	$2/10^{2}$	$4/10^{5}$	$2/10^{4}$	$1/10^{4}$	$2/10^{5}$	$2/10^{6}$	$4/10^{5}$	$3/10^{3}$
w. test-time adaptation
Adaptive $\theta$	$99.6$	$99.2$	${98.5}$	$96.9$	$94.1$	$89.1$	$81.2$	$69.1$	$54.2$	$39.0$	$27.1$
Adaptive $\theta$ (+ LN)	$\mathbf{99.7}$	$\mathbf{99.4}$	$\mathbf{98.7}$	$\mathbf{97.5}$	$\mathbf{95.1}$	$\mathbf{91.0}$	$\mathbf{84.0}$	$\mathbf{73.6}$	$\mathbf{58.9}$	$\mathbf{43.1}$	$\mathbf{30.4}$
$p$-value	$7/10^{1}$	$5/10^{4}$	$1/10^{2}$	$5/10^{5}$	$4/10^{4}$	$5/10^{5}$	$1/10^{4}$	$2/10^{5}$	$8/10^{5}$	$4/10^{4}$	$7/10^{4}$

Table 2: Results (%) on the dictionary lookup task. Results are averaged over $100$ runs with different random seeds. $p$-values are computed using a paired $t$-test. Entries highlighted in green indicate those with in-distribution sequence lengths.

Model	$2^{4}$	$2^{5}$	$2^{6}$	$2^{7}$	$2^{8}$	$2^{9}$	$2^{10}$	$2^{11}$	$2^{12}$	$2^{13}$	$2^{14}$
w.o. test-time adaptation
Baseline	$99.3$	$98.6$	${97.3}$	$94.7$	$89.5$	$80.4$	$67.6$	$52.9$	$38.7$	$26.5$	$17.8$
Baseline (+ LN)	$\mathbf{99.4}$	$\mathbf{98.8}$	$\mathbf{97.6}$	$\mathbf{95.3}$	$\mathbf{90.7}$	$\mathbf{82.9}$	$\mathbf{71.7}$	$\mathbf{57.7}$	$\mathbf{44.1}$	$\mathbf{32.3}$	$\mathbf{22.4}$
$p$-value	$6/10^{2}$	$1/10^{2}$	$5/10^{2}$	$2/10^{3}$	$2/10^{4}$	$3/10^{8}$	$1/10^{11}$	$2/10^{12}$	$7/10^{15}$	$3/10^{21}$	$2/10^{19}$
w. test-time adaptation
Adaptive $\theta$	$99.3$	$98.6$	${97.2}$	$94.5$	$89.3$	$80.4$	$67.8$	$52.6$	$38.6$	$27.3$	$20.8$
Adaptive $\theta$ (+ LN)	$\mathbf{99.4}$	$\mathbf{98.8}$	$\mathbf{97.6}$	$\mathbf{95.4}$	$\mathbf{90.6}$	$\mathbf{82.9}$	$\mathbf{71.7}$	$\mathbf{57.8}$	$\mathbf{44.5}$	$\mathbf{33.4}$	$\mathbf{27.7}$
$p$-value	$6/10^{1}$	$5/10^{2}$	$3/10^{2}$	$1/10^{4}$	$3/10^{4}$	$1/10^{5}$	$8/10^{9}$	$6/10^{12}$	$2/10^{16}$	$9/10^{20}$	$1/10^{21}$

Does layer normalization alleviate distribution shift?

Layer normalization does alleviate—but not eliminate—distribution shift. With layer normalization, the global mean and global variance remain more stable on out-of-distribution sequence lengths (Fig. 3). However, the variance of fixed components in attention outputs still decays, regardless of layer normalization (Fig. 2).

3.2 Ablations

In addition to layer normalization, we explore an alternative normalization strategy in which we standardize (i.e., std. in Table 3) the attention outputs across the $D$ features without the learnable scale and shift parameters present in LN, i.e., $\operatorname{Standardize}(\mathbf{O})_{b,t,d}=\frac{\mathbf{O}_{b,t,d}-\mathbf{\mu}_{b,t}}{\mathbf{\sigma}_{b,t}+\epsilon},$ where $\mu_{b,t}$ and $\sigma_{b,t}$ are computed in the same manner as in layer normalization.

As shown in Table 3, where the relative accuracy gain over Baseline on the $\arg\max$ retrieval task is reported, standardization improves length generalization, even though it strictly constrains model capacity. This underscores the importance (and potential benefits) of addressing the observed distribution shift. LN outperforms standardization, as confirmed by the paired $t$-test. Similar ablation results on the dictionary lookup task can be found in Section B.3.

Table 3: Ablations on different normalization strategies on the $\operatorname{argmax}$ retrieval task. Relative results (%) compared to the Baseline ($\triangle$) are reported.

Model	$2^{4}$	$2^{5}$	$2^{6}$	$2^{7}$	$2^{8}$	$2^{9}$	$2^{10}$	$2^{11}$	$2^{12}$	$2^{13}$	$2^{14}$
$\triangle$ (+ std.)	$-0.05$	$+0.00$	$+0.09$	$+0.26$	$+0.60$	$+0.84$	$+1.62$	$+2.26$	$+3.05$	$+1.80$	$+0.70$
$\triangle$ (+ LN)	$\mathbf{+0.01}$	$\mathbf{+0.11}$	$\mathbf{+0.21}$	$\mathbf{+0.57}$	$\mathbf{+1.15}$	$\mathbf{+1.81}$	$\mathbf{+2.98}$	$\mathbf{+4.32}$	$\mathbf{+4.99}$	$\mathbf{+3.30}$	$\mathbf{+1.76}$
$p$-value	$2/10^{3}$	$7/10^{4}$	$1/10^{2}$	$5/10^{4}$	$3/10^{4}$	$3/10^{4}$	$7/10^{4}$	$2/10^{4}$	$3/10^{4}$	$1/10^{3}$	$5/10^{3}$

4 Related Work

Positional encoding for length generalization.

Many works have attributed the inability of Transformers to extrapolate to longer sequences to positional encoding. Several alternatives to the sinusoidal positional encoding originally introduced by Vaswani et al. (2017) have been proposed to enhance the performance of Transformer-based models in natural language processing (NLP) tasks, including relative positional encoding (Shaw et al., 2018; Dai et al., 2019), rotary positional encoding (Su et al., 2024), no positional encoding (Haviv et al., 2022) and randomized positional encoding (Ruoss et al., 2023). Authors have examined the impact of different variants of positional encoding on length generalization (Chi et al., 2022; Ruoss et al., 2023; Kazemnejad et al., 2024; Li et al., 2024; Peng et al., 2024; Zhou et al., 2024). Unlike prior work that explores positional encoding for length generalization, we focus on algorithmic tasks that are order-invariant. We present a vanishing variance perspective on length generalization which is orthogonal to the extrapolability of positional encoding.

Alternatives to softmax attention.

The $\operatorname{softmax}$ output has been utilized to interpret the inner workings of Transformers (Xu et al., 2015; Choi et al., 2016; Martins & Astudillo, 2016). More recently, Veličković et al. (2024) demonstrated that the attention weights output by $\operatorname{softmax}$ will disperse as sequence length increases, attributing this phenomenon to the Transformer’s limited capability in length generalization. In this paper, we show that this dispersion leads to the vanishing variance problem in the intermediate attention outputs. While many variants of softmax attention have been introduced (Correia et al., 2019; Press et al., 2022; Tan et al., 2024; Ye et al., 2024), they are motivated mostly by interpretability, rather than the distribution of attention outputs for length generalization. To the best of our knowledge, none of the existing works have fundamentally eliminated the vanishing variance problem we presented in this paper. We hope our study can motivate designs of network architectures that are provably invariant to sequence length variations.

5 Conclusion

In this paper, we have introduced the vanishing variance problem and provided both theoretical analysis and empirical evidence demonstrating its role in inducing distribution shift in attention outputs. This shift hinders the ability of Transformers to generalize effectively to out-of-distribution sequence lengths. We demonstrated that mitigating this distribution shift through techniques like layer normalization and standardization—despite potential trade-offs in model expressiveness—significantly improves length generalization in attention models.

Future work.

We conduct our experiments using a single-layer, single-head attention architecture for simplicity, while real-world models typically use multi-layer, multi-head attention. Our conclusions may not fully generalize to these more complex architectures. Future work may validate the normalization strategies on larger benchmarks like CLRS (Veličković et al., 2022) and real-world LLMs. Moreover, layer normalization only partially mitigates distribution shift presented in this paper, and is already widely adopted in Transformers (though not immediately after attention outputs). Future work may design architectures that are provably invariant to sequence length variations.

Appendix A Proof of Proposition 1

Proposition 1 (The vanishing variance problem).

Consider a trained attention module with weights $\mathbf{W}_{Q},\mathbf{W}_{K},\mathbf{W}_{V},\mathbf{W}_{O}$. Let $\mathbf{X}=\left[\mathbf{x}_{1}\|\mathbf{x}_{2}\|\dots\|\mathbf{x}_{N}\right]^{\top}$ denote an input sequence of length $N$. If (1) $\mathbf{x}_{1},\mathbf{x}_{2},\dots,\mathbf{x}_{N}\overset{\text{i.i.d}}{\sim}\mathcal{X}$, a distribution over a finite vocabulary, and (2) $\mathbb{E}_{\mathbf{x}\sim\mathcal{X}}[\mathbf{W}_{V}\mathbf{x}]=\mathbf{0}$, then for a fixed query $\mathbf{y}$ and a fixed feature $d$,

$$\lim_{N\to\infty}\operatorname{Var}_{(\mathbf{x}_{1},\mathbf{x}_{2},\ldots,\mathbf{x}_{N})\sim\mathcal{X}^{N}}\left(\left[\operatorname{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^{\top}}{\sqrt{D}}\right)\mathbf{V}\right]_{d}\right)=0,$$

where $\mathbf{x}_{n},\mathbf{y}\in\mathbb{R}^{D}$ and $\mathbf{Q}\in\mathbb{R}^{1\times D},\mathbf{K}\in\mathbb{R}^{N\times D},\mathbf{V}\in\mathbb{R}^{N\times D}$ are intermediate results in $\operatorname{Attn}(\mathbf{X},[\mathbf{y}])$.

Informally, for a fixed component $d$ in the attention outputs, its variance over input sequences of length $N$, where each sequence consists of $N$ independently and identically distributed (i.i.d.) tokens, vanishes as $N\to\infty$.

Proof.

Let $\mathbf{V}_{n,d}=[\mathbf{W}_{V}\mathbf{x}_{n}]_{d}$. Firstly, we argue that $\mathbf{V}_{k,d}$ and $\mathbf{V}_{l,d}$ are independent for $k\neq l$. Let $\pi_{d}:\mathbb{R}^{D}\to\mathbb{R}$ be the projection onto the $d$-th coordinate, namely $\pi_{d}(\mathbf{v})=\mathbf{v}_{d}$ for $\mathbf{v}\in\mathbb{R}^{D}$. Observe that $\mathbf{V}_{k,d}=\pi_{d}(\mathbf{W}_{V}\mathbf{x}_{k})=(\pi_{d}\circ\mathbf{W}_{V})(\mathbf{x}_{k})$ and $\mathbf{V}_{l,d}=\pi_{d}(\mathbf{W}_{V}\mathbf{x}_{l})=(\pi_{d}\circ\mathbf{W}_{V})(\mathbf{x}_{l})$. By assumption (1), $\mathbf{x}_{k}$ and $\mathbf{x}_{l}$ are independent for $k\neq l$. Since $\mathbf{V}_{k,d}$, $\mathbf{V}_{l,d}$ are measurable functions of independent random variables, they are independent for $k\neq l$. By assumption, $\mathbf{x}_{k}$ and $\mathbf{x}_{l}$ are identically distributed for every $k,l$, so $\mathbf{V}_{k,d}$ and $\mathbf{V}_{l,d}$ are identically distributed. Thus, $\operatorname{Var}[\mathbf{V}_{k,d}]$ depends only on $d$, not on $k$. Set $\sigma_{d}^{2}=\operatorname{Var}[\mathbf{V}_{k,d}]$. $\sigma_{d}^{2}$ is finite since the vocabulary is finite and thus compact and bounded.

Let $\mathbf{A}$ denote the attention weights $\operatorname{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^{\top}}{\sqrt{D}}\right)\in\mathbb{R}^{1\times N}$ as the query sequence $[\mathbf{y}]$ consists of only a single item. Let $\mathbf{A}_{n}$ denote the $n$-th element of $\mathbf{A}$. We have

$$\operatorname{Var}\left(\sum_{n=1}^{N}\mathbf{A}_{n}\mathbf{V}_{n,d}\right) =\mathbb{E}\left[\left(\sum_{n=1}^{N}\mathbf{A}_{n}\mathbf{V}_{n,d}\right)^{2}\right]-\left(\mathbb{E}\left[\sum_{n=1}^{N}\mathbf{A}_{n}\mathbf{V}_{n,d}\right]\right)^{2}$$

$$\leq\mathbb{E}\left[\left(\sum_{n=1}^{N}\mathbf{A}_{n}\mathbf{V}_{n,d}\right)^{2}\right]$$

$$=\mathbb{E}\left[\sum_{n=1}^{N}\mathbf{A}_{n}^{2}\mathbf{V}_{n,d}^{2}\right]+\mathbb{E}\left[\sum_{1\leq k,l\leq N,k\neq l}\mathbf{A}_{k}\mathbf{A}_{l}\mathbf{V}_{k,d}\mathbf{V}_{l,d}\right]$$

$$\leq\mathbb{E}\left[\sum_{n=1}^{N}(\operatorname{max}_{1\leq n\leq N}\mathbf{A}_{n})^{2}\mathbf{V}_{n,d}^{2}\right]+\mathbb{E}\left[\sum_{1\leq k,l\leq N,k\neq l,\mathbf{V}_{k,d}\mathbf{V}_{l,d}\geq 0}\mathbf{A}_{k}\mathbf{A}_{l}\mathbf{V}_{k,d}\mathbf{V}_{l,d}\right]$$

$$\leq(\operatorname{max}_{1\leq n\leq N}\mathbf{A}_{n})^{2}\sum_{n=1}^{N}\mathbb{E}\left[\mathbf{V}_{n,d}^{2}\right]+\mathbb{E}\left[\sum_{1\leq k,l\leq N,k\neq l,\mathbf{V}_{k,d}\mathbf{V}_{l,d}\geq 0}\mathbf{V}_{k,d}\mathbf{V}_{l,d}\right]$$

$$\leq N(\operatorname{max}_{1\leq n\leq N}\mathbf{A}_{n})^{2}\sigma_{d}^{2}+\sum_{1\leq k,l\leq N,k\neq l,\mathbf{V}_{k,d}\mathbf{V}_{l,d}\geq 0}\mathbb{E}\left[\mathbf{V}_{k,d}\mathbf{V}_{l,d}\right]$$

$$\leq N(\operatorname{max}_{1\leq n\leq N}\mathbf{A}_{n})^{2}\sigma_{d}^{2}+\sum_{1\leq k,l\leq N,k\neq l,\mathbf{V}_{k,d},\mathbf{V}_{l,d}\geq 0}\mathbb{E}\left[\mathbf{V}_{k,d}\right]\mathbb{E}\left[\mathbf{V}_{l,d}\right]$$

$$\leq N(\operatorname{max}_{1\leq n\leq N}\mathbf{A}_{n})^{2}\sigma_{d}^{2}$$

In the derivation above, we used the fact that $\mathbf{A}_{n}\in[0,1]$ and that for $k\neq l$, $\mathbf{V}_{k,d}$ and $\mathbf{V}_{l,d}$ are independent. We also used the assumption that for every $1\leq k\leq N$, $\mathbb{E}\left[\mathbf{V}_{k,d}\right]=0$. Since the tokens come from a finite dictionary, and since $\mathbf{x}\to\mathbf{W}_{Q}\mathbf{x}$ and $\mathbf{x}\to\mathbf{W}_{K}\mathbf{x}$ are continuous functions on compact domain (dictionary is finite), the logits $\langle\mathbf{W}_{Q}\mathbf{x}_{i},\mathbf{W}_{K}\mathbf{x}_{j}\rangle$ are bounded, because they are continuous image of a compact set and every compact set on the real line is closed and bounded. By Lemma 2.1 of Veličković et al. (2024), there exist a constant $C>0$ and $N_{0}\in\mathbb{N}$, such that for every $N\geq N_{0}$, $(\operatorname{max}_{1\leq n\leq N}\mathbf{A}_{n})^{2}<\frac{C}{N^{2}}$. Then for every $N\geq N_{0}$,

$$\operatorname{Var}\left(\sum_{n=1}^{N}\mathbf{A}_{n}\mathbf{V}_{n,d}\right) \leq N\sigma^{2}\frac{C}{N^{2}}=\sigma^{2}\frac{C}{N}.$$

Let $\epsilon>0$. There exists $N_{1}\in\mathbb{N}$ such that for every $N\geq N_{1}$, $\sigma^{2}\frac{C}{N}<\epsilon$. Then for every $N\geq\operatorname{max}(N_{0},N_{1})$,

$$\operatorname{Var}\left(\sum_{n=1}^{N}\mathbf{A}_{n}\mathbf{V}_{n,d}\right) <\epsilon.$$

∎

Table 4: Results (%) on the $\operatorname{argmax}$ retrieval task. Results are averaged over $100$ runs with different random seeds. $p$-values are computed using a paired $t$-test. Entries highlighted in green indicate those with in-distribution sequence lengths.

Model	$2^{4}$	$2^{5}$	$2^{6}$	$2^{7}$	$2^{8}$	$2^{9}$	$2^{10}$	$2^{11}$	$2^{12}$	$2^{13}$	$2^{14}$
Baseline	$99.8$	$99.8$	${99.6}$	$99.3$	$98.5$	$97.1$	$94.4$	$89.1$	$79.9$	$65.8$	$47.8$
Baseline (+ LN)	$\mathbf{99.9}$	$\mathbf{99.9}$	$\mathbf{99.8}$	$\mathbf{99.5}$	$\mathbf{99.1}$	$\mathbf{98.1}$	$\mathbf{96.2}$	$\mathbf{92.8}$	$\mathbf{86.3}$	$\mathbf{75.1}$	$\mathbf{58.7}$
$p$-value	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}1/10^{8}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}6/10^{6}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}1/10^{4}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}2/10^{4}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}9/10^{8}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}5/10^{8}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}3/10^{8}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}5/10^{10}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}1/10^{11}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}3/10^{12}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}1/10^{13}}$

Table 5: Results (%) on the dictionary lookup task. Results are averaged over $100$ runs with different random seeds. $p$-values are computed using a paired $t$-test. Entries highlighted in green indicate those with in-distribution sequence lengths.

Model	$2^{4}$	$2^{5}$	$2^{6}$	$2^{7}$	$2^{8}$	$2^{9}$	$2^{10}$	$2^{11}$	$2^{12}$	$2^{13}$	$2^{14}$
Baseline	$\mathbf{99.9}$	$\mathbf{99.9}$	${99.8}$	$99.7$	$99.5$	$99.1$	$98.3$	$96.5$	$93.5$	$87.7$	$77.8$
Baseline (+ LN)	$\mathbf{99.9}$	$\mathbf{99.9}$	$\mathbf{99.9}$	$\mathbf{99.8}$	$\mathbf{99.6}$	$\mathbf{99.3}$	$\mathbf{98.7}$	$\mathbf{97.5}$	$\mathbf{95.0}$	$\mathbf{90.4}$	$\mathbf{82.6}$
$p$-value	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}2/10^{1}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}1/10^{2}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}5/10^{2}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}2/10^{3}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}3/10^{3}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}6/10^{5}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}4/10^{7}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}8/10^{9}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}2/10^{10}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}2/10^{10}}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}2/10^{13}}$

Table 6: Ablations on different normalization strategies on the dictionary lookup task. Relative results (%) compared to the Baseline ($\triangle$) are reported.

Model	$2^{4}$	$2^{5}$	$2^{6}$	$2^{7}$	$2^{8}$	$2^{9}$	$2^{10}$	$2^{11}$	$2^{12}$	$2^{13}$	$2^{14}$
$\triangle$ (+ std.)	$\mathbf{+0.09}$	$+0.14$	$+0.23$	$+0.54$	$+1.08$	$+2.27$	$+3.49$	$+4.78$	$+5.14$	$+5.65$	$\mathbf{+4.57}$
$\triangle$ (+ LN)	$\mathbf{+0.09}$	$\mathbf{+0.20}$	$\mathbf{+0.30}$	$\mathbf{+0.64}$	$\mathbf{+1.22}$	$\mathbf{+2.55}$	$\mathbf{+4.06}$	$\mathbf{+4.86}$	$\mathbf{+5.38}$	$\mathbf{+5.72}$	$+4.51$
$p$-value	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}1.0}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}0.3}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}0.6}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}0.6}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}0.6}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}0.5}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}0.3}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}0.9}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}0.6}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}0.9}$	${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}0.9}$

Appendix B More Experimental Details

B.1 Implementation Details

$\operatorname{argmax}$ retrieval.

We follow Veličković et al. (2024) and train the same neural network architecture in PyTorch for $100,000$ gradient steps with the same hyper-parameter setup. We also follow Veličković et al. (2024) to generate data of varying number of items to train and test the model.

$\operatorname{dictionary}$ lookup.

The network architecture is the same as the $\arg\max$ retrieval task. We generate data for training and evaluation in the following way: for each item of the length-$N$ sequence, we sample a value class $c_{V}\sim\mathcal{U}\{1,\dots,C_{V}\}$ i.i.d at random; each item also has a key class $1\leq c_{K}\leq C_{K}$. The key classes of all $N$ items in the sequence are sampled without replacement. In our experiments, $C_{K}=16384$ and $C_{V}=64$.

The features of each item $\mathbf{x}_{i}$ is defined as $\operatorname{Emb}_{K}(c_{K})\parallel\operatorname{Emb}_{V}(c_{V})$, i.e., the concatenation of the embeddings of the key class and the value class. The embedding vectors for each (key and value) class are optimized jointly with the attention network.

The query sequence in our case is guaranteed to be of length $1$. We sample a key class present in the input sequence and use its embedding vector as the query.

For this task, we found that the optimization usually converges within $10,000$ gradient steps. We train the attention network, together with the embedding vectors, in PyTorch for $10,000$ steps with the same hyper-parameter setup as the $\arg\max$ task.

B.2 Results When Training on More Diverse Sequence Lengths

To validate the utility of normalization when the length gap between the training sequences and the test sequences is smaller, we follow the same experimental setup as in Section 3, but sample sequences of up to $256$ items during training. We found it beneficial to gradually increase the length of the sequences sampled throughout training, as is commonly done during pre-training of frontier LLMs (AI@Meta, 2024). The results are reported in Table 4 and Table 5. With layer normalization, the accuracies on out-of-distribution sequence lengths are significantly higher than without on both tasks, demonstrating the importance of normalization for length generalization over various training settings.

B.3 Ablations on the Dictionary Lookup Task

Ablation results on the dictionary lookup task are shown in Table 6, which are consistent with the results on the $\arg\max$ retrieval task presented in Section 3.2. However, on this task, the performance of standardization and layer normalization is more similar, as indicated by the larger $p$-values, suggesting weaker statistical evidence for a significant difference.