Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMs

## Abstract

Abstract Practitioners have consistently observed three puzzling phenomena in transformer-based large language models (LLMs): attention sinks , value-state drains , and residual-state peaks , collectively referred to as extreme-token phenomena . These phenomena are characterized by certain so-called “sink tokens” receiving disproportionately high attention weights, exhibiting significantly smaller value states, and having much larger residual-state norms than those of other tokens. These extreme tokens give rise to various challenges in LLM inference, quantization, and interpretability. We elucidate the mechanisms behind extreme-token phenomena. First, we show that these phenomena arise in very simple architectures—transformers with one to three layers—trained on a toy model, the Bigram-Backcopy (BB) task. In this setting, we identify an active-dormant mechanism , where attention heads become sinks for specific input domains while remaining non-sinks for others. Our theoretical analysis of the training dynamics reveals that these phenomena are driven by a mutual reinforcement mechanism . Building on these insights, we propose strategies to mitigate extreme-token phenomena during pretraining, including replacing softmax with ReLU and Adam with SGD. Next, we extend our analysis to pretrained LLMs, including Llama and OLMo, showing that many attention heads exhibit a similar active-dormant mechanism as in the BB task, and that the mutual reinforcement mechanism also governs the emergence of extreme-token phenomena during LLM pretraining. Our results reveal that many of the static and dynamic properties of extreme-token phenomena predicted by the BB task align with observations in pretrained LLMs. † † Code for our experiments is available at https://github.com/GuoTianYu2000/Active-Dormant-Attention .

## 1 Introduction

Recent analyses of transformer-based open-source large language models (LLMs), such as GPT-2 (Radford et al., ), Llama-2 (Touvron et al., ), Llama-3 (Dubey et al., ), Mixtral (Jiang et al., ), and Pythia (Biderman et al., ), have revealed three intriguing phenomena:

- -
Attention sinks (Xiao et al., ): In many attention heads, the initial token consistently attracts a large portion of the attention weights. Other special tokens, such as the delimiter token, can also draw significant attention weights. These tokens are collectively referred to as *sink tokens*.
- -
Value-state drains (Guo et al., ): For the attention heads that exhibit attention sinks, the value states of sink tokens are consistently much smaller than those of other tokens.
- -
Residual-state peaks (Sun et al., ): The residual states of sink tokens, excluding those from the first and last layers, exhibit significantly larger norms compared to other tokens.

These phenomena often appear together and consistently occur in various pretrained LLMs, which we collectively refer to as the *extreme-token phenomena*. Figure 1 illustrates these phenomena in Llama-3.1-8B-Base, using a fixed prompt sentence: “$\langle\texttt{s}\rangle$Summer is warm$\langle$$\mathtt{period}$$\rangle$ Winter is cold$\langle$$\mathtt{period}$$\rangle$”. Here, the first token, $\langle\texttt{s}\rangle$ (the Beginning-of-Sequence token), serves as the sink token. As shown in the figure, the sink token receives disproportionately high attention weights, exhibits significantly smaller value states, and has much larger residual state norms compared to other tokens. It is important to note that the first token does not have to be $\langle\texttt{s}\rangle$ to act as a sink token; other tokens appearing first in the sequence can also serve this role. Additionally, in models such as Llama-2, a delimiter token can also function as the sink token.

Figure: (a) Attention weights at L24
Refer to caption: x1.png

The extreme-token phenomena have posed several challenges for pretrained transformers in downstream tasks. For instance, sink tokens require special treatment during long-context inference (Xiao et al., ; Han et al., ; Yu et al., ; Chen et al., ) and model quantization (Dettmers et al., ; Liu et al., ; Son et al., ) to maintain high levels of performance. Additionally, attention sinks have reduced the interpretability of attention maps in vision transformers (Darcet et al., ). To address these issues, Sun et al. () and Darcet et al. () propose adding a “special token” to transformers to serve as the sink token, preventing other tokens from becoming sinks. However, even this special token still exhibits extreme-token phenomena. Despite these efforts, no prior work has satisfiably explained the mechanisms behind the extreme-token phenomena. Xiao et al. () proposes a hypothesis for why they occur, suggesting that models tend to dump unnecessary attention values to specific tokens.

This work aims to demystify the extreme-token phenomena in LLMs. We demonstrate that these phenomena arise from an active-dormant mechanism in attention heads (cf. Claim 1), coupled with a mutual-reinforcement mechanism during pretraining (cf. Claim 2). We support these statements through studies on simplified transformer architectures and tasks, a dynamical theory for these models, and experiments on pretrained LLMs. The structure of the paper and our key contributions are outlined as follows:

- 1.
In Section 2, we train one- to three-layer transformers on a simple task called the Bigram-Backcopy (BB) task, which also displays extreme-token phenomena similar to those observed in LLMs. We show that attention sinks and value-state drains are a consequence of the active-dormant mechanism (cf. Claim 1). Both theoretically and empirically, we demonstrate that mutual reinforcement mechanism (cf. Claim 2) dynamically drives these phenomena: attention sinks and value-state drains reinforce one another, leading to a stable phase where all query tokens generate near identical attention logits for the keys of extreme tokens. Additionally, empirical results reveal that residual-state peaks arise from the interaction between this mutual reinforcement mechanism and the Adam optimization algorithm.
- 2.
In Section 3, we demonstrate the active-dormant mechanism in pre-trained LLMs by showing that many attention heads transition between active and dormant phases based on the input domain. Specifically, we identify an interpretable active-dormant head (Layer 16, Head 25 in Llama 2-7B-Base (Touvron et al., )) that activates on GitHub data but remains dormant on Wikipedia data. Moreover, in examining the dynamics of OLMo-7B-0424 (Groeneveld et al., ), we observe the same mutual reinforcement mechanism and stable phase, consistent with those found in the BB task. This demonstrates that the simple BB model captures both the static and dynamic properties of extreme-token phenomena in LLMs and accurately predicts their behavior.
- 3.
Importantly, the quantitative properties of extreme-token dynamics show strong consistency among the theoretical and empirical results of the Bigram-Backcopy task and the empirical performance of OLMo. In particular, we consistently observe the sink-logits concentration phenomenon, where the logits corresponding to the key of the extreme token and the queries of all non-extreme tokens ($\text{logit}_{\cdot,\texttt{$\langle\texttt{s}\rangle$}}$) are nearly identical—an observation not previously documented in the literature. We summarize the aligned results between the theoretical and empirical findings of the Bigram-Backcopy task and the empirical performance of LLMs in Table 1.
- 4.
We propose architectural and optimization modifications to mitigate the extreme-token phenomena. Specifically, we demonstrate that replacing SoftMax with ReLU activations in attention heads eliminates extreme-token phenomena in the BB task, while switching from Adam to SGD removes the residual-state peak phenomenon. We discuss the possibility that similar modifications could mitigate extreme-token phenomena in LLMs.

**Table 1: Consistency of the quantitative properties across the theoretical and empirical results of the Bigram-Backcopy task and empirical results of LLMs. A ✓ denotes a consistent result, while a $\star$ denotes an inconclusive result. The $\text{logit}_{\cdot,\texttt{$\langle\texttt{s}\rangle$}}$ denotes logits corresponding to the key of the extreme token and queries of all non-extreme tokens, i.e., the ${\tt Qry}_{\cdot}^{\top}{\tt Key}_{\texttt{$\langle\texttt{s}\rangle$}}$. The $\Delta\text{logit}_{\cdot,\texttt{$\langle\texttt{s}\rangle$}}=\text{logit}_{ \cdot,\texttt{$\langle\texttt{s}\rangle$}}-\text{Mean}[\text{logit}_{\cdot, \text{others}}]$ is a progress measure for attention sinks. The $\|{\tt Val}_{\texttt{$\langle\texttt{s}\rangle$}}\|$ denotes the value state norm of the extreme token, and $\|{\tt Res}_{\texttt{$\langle\texttt{s}\rangle$}}\|$ denotes the residual state norm of the extreme token. See Section 1.2 for the definitions of these notations.**
|  | BB-task Theory | BB-task Experiments | LLM Experiments |
| --- | --- | --- | --- |
| $\Delta\text{logit}_{\cdot,\texttt{$\langle\texttt{s}\rangle$}}$ $\log$-growth | ✓ | ✓ | $\star$ |
| $\|{\tt Val}_{\texttt{$\langle\texttt{s}\rangle$}}\|$ monotonic decrease | ✓ | ✓ | ✓ |
| $\|{\tt Res}_{\texttt{$\langle\texttt{s}\rangle$}}\|$ linear growth | $\star$ | ✓ | ✓ |
| $\text{logit}_{\cdot,\texttt{$\langle\texttt{s}\rangle$}}$ concentration | ✓ | ✓ | ✓ |

### 1.1 Related work

Several studies independently identified the “attention sink” phenomenon in language models and vision transformers, where attention weights were found to be concentrated on a few tokens (Xiao et al., ; Darcet et al., ; Han et al., ; Zhai et al., ; Elhage et al., ; Dettmers et al., ). Recent research has provided more detailed characterizations of this attention pattern and the attention sink phenomenon (Fu, ; Sun et al., ). Sun et al. () attributed the attention sink to the massive activation of the hidden representations of the corresponding tokens. Both Sun et al. () and Zhai et al. () discussed methods for mitigating the attention sink by modifying the model and training recipes. Additionally, recent studies have leveraged the attention sink phenomenon to develop improved quantization and more efficient inference algorithms (Liu et al., ; Chen et al., ; Yu et al., ; Son et al., ; Lin et al., ; Bondarenko et al., ; Hu et al., ). A concurrent work by Gu et al. () studied how optimization, data distribution, loss function, and model architecture in LM pre-training influence the emergence of attention sink, showing that replacing the softmax function with sigmoid can prevent attention sink emergence in models up to 1B parameters.

The dynamics of transformers are studied under various simplifications, including linear attention structures (Zhang et al., ; Ahn et al., ), reparametrizations (Tian et al., ), NTK (Deora et al., ), often in the setting of in-context linear regression (Ahn et al., ; Wu et al., ; Zhang et al., ) and structured sequences (Bietti et al., ; Nichani et al., ; Tian et al., ). Notably, Zhang et al. (); Huang et al. (); Kim et al. () demonstrate that a one-layer attention head trained via gradient descent converges to a model that effectively performs in-context regression. Bietti et al. () shows the fast learning of bigram memorization and the slow development of in-context abilities. Tian et al. () shows the scan and snap dynamics in reparametrized one-layer transformers. Reddy () simplifies the structure of the induction head, showing the connection between the sharp transitions of in-context learning dynamics and the nested nonlinearities of multi-layer operations.

Mechanistic interpretability is a growing field focused on understanding the internal mechanisms of language models in solving specific tasks (Elhage et al., ; Geva et al., ; Meng et al., ; Nanda et al., ; Olsson et al., ; Bietti et al., ; Wang et al., ; Feng and Steinhardt, ; Todd et al., ). This includes mechanisms like the induction head and function vector for in-context learning (Elhage et al., ; Olsson et al., ; Todd et al., ; Bietti et al., ), the binding ID mechanism for binding tasks (Feng and Steinhardt, ), association-storage mechanisms for factual identification tasks (Meng et al., ), and a complete circuit for indirect object identification tasks (Wang et al., ). The task addressed in this paper is closely related to Bietti et al. (), who explored synthetic tasks where tokens are generated from either global or context-specific bigram distributions. Several other studies have also employed synthetic tasks to explore neural network mechanisms (Charton, ; Liu et al., ; Nanda et al., ; Allen-Zhu and Li, ; Zhu and Li, ; Guo et al., ; Zhang et al., ; Lin et al., ).

A line of work focuses on quantizing neural networks using low-bit fixed-point representations (Jacob et al., ; Zafrir et al., ; Lin et al., ; Nagel et al., ; Gholami et al., ), such as INT8 (Lin et al., ; Dettmers et al., ) or INT4 (; Wu et al., ; Dettmers and Zettlemoyer, ) to save memory usage and computational cost. In LLMs, the extreme-token phenomena lead to substantial performance degradation after quantization (Bondarenko et al., ) and have become a key focus of recent research (Fan et al., ; Yao et al., ; Lin et al., ; Hu et al., ). Dettmers et al. () and Lin et al. () propose mixed-precision approaches, using FP16 for outlier values and INT8 for others, enabling large model quantization without performance loss. Xiao et al. () rescales the weights and activations to reduce magnitudes of outliers, and Bondarenko et al. () proposes modified attention structures to remove outliers, making language models easier to quantize.

We note that Gurnee et al. () proposed Attention Deactivation Neurons, Bondarenko et al. () proposed the “no-op” hypothesis, and Xiao et al. () proposed the “dump unnecessary attention” conjecture as mechanisms of attention sinks. In contrast, we explain the extreme-token phenomena through the active-dormant and mutual reinforcement mechanisms, offering the proof of their emergence within training dynamics in a toy model and providing empirical evidence of these mechanisms in LLMs.

### 1.2 Preliminaries and notations

While different LLMs may use slightly varying transformer architectures, most use the structure proposed by Vaswani (), with the key modification being the shift from post-norm to pre-norm. We represent the tokenized input sequence of length $n$, with positional embeddings included, as ${\mathbf{H}}=[{\mathbf{h}}_{1},\dots,{\mathbf{h}}_{n}]\in\mathbb{R}^{d\times n}$, where ${\mathbf{h}}_{i}$ denotes the $i$th input token, and $d$ is the embedding dimension. We denote the layer-normalization operation as ${\rm LN}$, the column-wise SoftMax operation as ${\sf SoftMax}$, the causal-mask as ${\sf mask}$, and the pointwise ReLU function as $\mathrm{ReLU}$.

The transformer architecture applies causal-attention and MLP layers iteratively to the input sequence ${\mathbf{H}}$. A causal-attention layer with $M$ heads is represented as ${\rm Attn}(\cdot)$, parameterized by ${\{({\mathbf{Q}}_{m},{\mathbf{K}}_{m},{\mathbf{V}}_{m},{\mathbf{O}}_{m})\}}_{m}$:

$$ $\textstyle{\rm Attn}({\mathbf{H}})\mathrel{\mathop{:}}=\sum_{m=0}^{M-1}{\tt attn }_{m}({\mathbf{H}})\in\mathbb{R}^{d\times n},$ (1) $$

where each attention head ${\tt attn}_{m}(\cdot)$ is given by

$$ $\textstyle{\tt attn}_{m}({\mathbf{H}})\mathrel{\mathop{:}}={\mathbf{O}}_{m}{ \mathbf{V}}_{m}{\rm LN}({\mathbf{H}}){\sf SoftMax}{\left({\sf mask}{\left({\rm LN }({\mathbf{H}})^{\top}{\mathbf{K}}_{m}^{\top}{\mathbf{Q}}_{m}{\rm LN}({\mathbf {H}})\right)}\right)}.$ (2) $$

We denote the attention map as ${\sf Map}={\sf SoftMax}{\left({\sf mask}{\left({\rm LN}({\mathbf{H}})^{\top}{
\mathbf{K}}_{m}^{\top}{\mathbf{Q}}_{m}{\rm LN}({\mathbf{H}})\right)}\right)}$, and typically plot its transpose, ${\sf Map}^{\top}$, in figures.

An MLP layer, denoted ${\tt mlp}(\cdot)$, has parameters $({\mathbf{W}}_{1},{\mathbf{W}}_{2})$:

$$ $\textstyle{\tt mlp}({\mathbf{H}})\mathrel{\mathop{:}}={\mathbf{W}}_{2}\mathrm{ ReLU}({\mathbf{W}}_{1}{\rm LN}({\mathbf{H}}))\in\mathbb{R}^{d\times n}.$ (3) $$

An $L$-layer transformer consists of a composition of $L$ self-attention and MLP layers with residual connection structure. Given an input ${\mathbf{H}}^{(0)}\in\mathbb{R}^{d\times n}$, the output of the $L$-layer transformer, ${\mathbf{H}}^{(L)}$, is computed as follows:

$$ $\textstyle{\mathbf{H}}^{(\ell+1)}={\mathbf{H}}^{(\ell+1/2)}+{\tt mlp}^{(\ell)} {\left({\mathbf{H}}^{(\ell+1/2)}\right)},~{}~{}~{}{\mathbf{H}}^{(\ell+1/2)}={ \mathbf{H}}^{(\ell)}+{\rm Attn}^{(\ell)}{\left({\mathbf{H}}^{(\ell)}\right)},~ {}~{}~{}\ell\in{\left\{0,\dots,L-1\right\}}.$ (4) $$

For consistency between the code and the text, we adopt zero-indexing throughout this paper, meaning that attention head and layer indices begin at 0 0 instead of $1$.

For the output ${\mathbf{H}}^{(\ell+1)}$ of layer $\ell$, we define the residual state ${\tt Res}_{v}$ of a token $v\in\{0,1,\ldots,n-1\}$ as the $v$th column of ${\mathbf{H}}^{(\ell+1)}$. For a specific layer $\ell$ with input ${\mathbf{H}}^{(\ell)}\in\mathbb{R}^{d\times n}$, and for a specific attention head $m$ with query, key, and value matrices $({\mathbf{Q}},{\mathbf{K}},{\mathbf{V}},{\mathbf{O}})$, we define the query, key, and value states $({\tt Qry}_{v},{\tt Key}_{v},{\tt Val}_{v})$ of a token $v\in[n]$ as the $v$th columns of ${\mathbf{Q}}{\mathbf{H}}^{(\ell)}$, ${\mathbf{K}}{\mathbf{H}}^{(\ell)}$, and ${\mathbf{O}}{\mathbf{V}}{\mathbf{H}}^{(\ell)}$, respectively(^1^11We define the value state as ${\mathbf{O}}{\mathbf{V}}{\mathbf{H}}$ rather than ${\mathbf{V}}{\mathbf{H}}$ since what is added to the residual state is essentially ${\mathbf{O}}{\mathbf{V}}{\mathbf{H}}$.). The attention logit $\text{logit}_{v^{\prime},v}$ is defined as the $(v^{\prime},v)$th element of $({\mathbf{H}}^{(\ell)})^{\top}{\mathbf{Q}}^{\top}{\mathbf{K}}{\mathbf{H}}^{(
\ell)}$. For notation simplicity, we omit the dependence on $\ell$ and $m$ in $({\tt Qry}_{v},{\tt Key}_{v},{\tt Val}_{v},\text{logit}_{v^{\prime},v})$, as these will be clear from context. Additionally, for a fixed token $v$, we use the shorthand $\text{logit}_{\cdot,v}$ for the set $\{\text{logit}_{v^{\prime},v}\mid v^{\prime}\in\mathcal{V}\}$.

We use $\langle\texttt{s}\rangle$ to refer to the "Beginning-of-Sequence" token. Since the $\langle\texttt{s}\rangle$ token consistently behaves as an extreme token in LLMs, we often refer to $\langle\texttt{s}\rangle$ and extreme tokens interchangeably. We also abuse the notation by writing $({\tt Qry}_{\texttt{$\langle\texttt{s}\rangle$}},{\tt Key}_{\texttt{$\langle
\texttt{s}\rangle$}},{\tt Val}_{\texttt{$\langle\texttt{s}\rangle$}})$ to represent the query, key, and value states of the $\langle\texttt{s}\rangle$ token.

## 2 Extreme-token Phenomena in the Bigram-Backcopy Task

In this section, we analyze simple transformers trained on the Bigram-Backcopy (BB) task, a simple model that exhibits extreme-token phenomena. We demonstrate the active-dormant mechanism (cf. Claim 1) and mutual reinforcement mechanism (cf. Claim 2) within the BB task and provide predictions for the behavior of sink tokens, which will be validated through LLM experiments in the following section.

The Bigram-Backcopy task is a data-generation model that consists of two sub-tasks: Bigram-transition and Backcopy. In this model, each sequence begins with the $\langle\texttt{s}\rangle$ token, followed by tokens sampled according to a pre-determined bigram transition probability ${\sf P}$ (in other words, a Markov chain). When specific trigger tokens are encountered, instead of sampling according to the transition ${\sf P}$, the preceding token is copied to the next position. An illustration of the Bigram-Backcopy task is provided in Figure 2. Following Bietti et al. (), we select the transition ${\sf P}$ and the vocabulary $\mathcal{V}$ with $|\mathcal{V}|=V=64$ based on the estimated character-level bigram distribution from the tiny Shakespeare dataset. In all experiments, the set of trigger tokens, $\mathcal{T}$, is fixed and consists of the $\left|\mathcal{T}\right|=3$ most frequent tokens from the unigram distribution. Consequently, the non-trigger token set, $\mathcal{V}\setminus\mathcal{T}$, comprises $61$ tokens.

### 2.1 One-layer transformer exhibits attention sinks and value-state drains

On the Bigram-Backcopy task, we pre-train a standard one-layer transformer with a single SoftMax attn head and one mlp layer. Unless otherwise specified, the model is trained using Adam for $10,000$ steps, achieving near-optimal prediction accuracy. Detailed training procedures are provided in Appendix C.1. Figure 2 shows that the trained transformer exhibits the attention sink phenomenon, where the $\langle\texttt{s}\rangle$ token captures a significant proportion of the attention weights. More importantly, the attention weights display interpretable patterns: all non-trigger tokens exhibit attention sinks, while the attention for trigger tokens is concentrated on their preceding positions. Additionally, Figure 2 reveals a value-state drain phenomenon similar to that observed in LLMs, suggesting that, for non-trigger tokens, the attn head contributes minimal value to the residual stream. We provide additional attention patterns on different input sequences in Section C.2.

Figure: (d) The Bigram-Backcopy task
Refer to caption: x4.png

#### The active-dormant mechanism of the attention head.

Inspired by the interpretable attention weight patterns observed, we propose the active-dormant mechanism. For any given token, an attention head is considered active if it makes a significant contribution to the residual state, and dormant if its contribution is minimal. As illustrated in Figure 2, when trained on the BB task, the attention head is active for trigger tokens and dormant for non-trigger tokens.

Figure 4 demonstrates that the mlp layer is responsible for the Bigram task whereas the attn head takes care of the Backcopy task. When the mlp layer is zeroed out, the backcopy loss remains significantly better than a random guess, but the bigram loss degrades to near-random levels. Conversely, when the attn layer is zeroed out, the backcopy loss becomes worse than a random guess, while the bigram loss remains unaffected. This indicates that on trigger tokens, the attn head is active and handles the backcopy task, whereas on non-trigger tokens, the attn head is dormant, allowing the mlp layer to handle the Bigram task. We summarize the active-dormant mechanism of the attn head in Claim 1.

Figure: Figure 3: Active-dormant mechanism
Refer to caption: x7.png

Figure: (a) Excess risk after interventions
Refer to caption: x8.png

###### Claim 1 (Active-dormant mechanism) .

#### The growth of attention logits on the ⟨ s ⟩ delimited-⟨⟩ s \langle\texttt{s}\rangle ⟨ s ⟩ token and the decrease in its value state norms.

Figure 4 illustrates the training dynamics of excess risks, attention weights, attention logits (for each token $v_{n}$ at position $n$ in the prompt, we compute $\Delta\text{logit}_{\cdot,\texttt{$\langle\texttt{s}\rangle$}}\equiv\mathtt{
mean}_{n}[\langle{\tt Qry}_{v_{n}},{\tt Key}_{\texttt{$\langle\texttt{s}
\rangle$}}\rangle-\mathtt{mean}_{i}(\langle{\tt Qry}_{v_{n}},{\tt Key}_{v_{i}}
)\rangle]$, which serves as a progress measure for attention sinks), and value state norms for the $\langle\texttt{s}\rangle$ token. All values are rescaled to the 0 0 to $1$ range to highlight trends rather than absolute values. Both the Bigram and Backcopy excess risks decrease to nearly zero within the first 1000 steps, with the Bigram excess risk approaching zero faster than the Backcopy risk. As the Backcopy risk decreases, the attention weights on the $\langle\texttt{s}\rangle$ token begin to increase, suggesting a connection between the formation of attention sinks and the backcopy function in the attention heads. After the first $1000$ steps, although both Bigram and Backcopy excess risks have nearly reached zero, the attention logits and weights on the $\langle\texttt{s}\rangle$ token continue to increase, while the value state norm of the $\langle\texttt{s}\rangle$ token continues to decrease. While this is an intriguing phenomenon, our next goal is to understand why the attention logits and value state norms continue to evolve toward extreme values.

### 2.2 Analysis of a minimally-sufficient transformer architecture

Figure: Figure 5: Simplified transformer architecture. The output logits are computed by summing the contributions from both the mlp layer and the attn head. The predicted probabilities are obtained by applying the SoftMax function to these output logits. The mlp layer is assumed to provide the Markov transition probabilities for non-trigger tokens, while the attn head is parameterized by attention logits and value states, as described in Eq. (6), (7), and (8). Additionally, the trainable variables, denoted by $(\bm{\alpha},\bm{\beta})\in\mathbb{R}^{V}\times\mathbb{R}^{V}$, represent the attention logits and value states of the $\langle\texttt{s}\rangle$ token.
Refer to caption: x10.png

In this section, we analyze the training dynamics of transformers on the BB task, focusing on a simplified architecture that retains the attention sinks and value-state-drains phenomena. We analyze the regime when the Bigram transition probability is fully learned, and the Backcopy task is partially learned (i.e., after step $200$ in Figure 4), and we focus on the dynamics of the attention logits and value states. Readers who are more interested in the results than the theoretical analysis can skip the detailed analysis and proceed directly to the statement of the mutual reinforcement mechanism in Claim 2.

Let $\mathcal{V}$ (of size $V$) denote the set of all tokens excluding the $\langle\texttt{s}\rangle$ token, and let $\mathcal{T}$ represent the set of all trigger tokens. For any $v\in\mathcal{V}$, we define $p_{vk}={\sf P}(k|v)$ as the next-token Markov transition probability, and $\bm{p}_{v}=(p_{v1},\ldots,p_{vV})^{\top}\in\Delta(\mathcal{V})$ as the transition vector in the simplex. The embedding map is denoted by ${\tt ebd}:[n]\times\mathcal{V}\to\mathbb{R}^{D}$, where for a token $v\in\mathcal{V}$ at position $i\in[n]$, the embedded vector is ${\tt ebd}_{i}(v)$. The $\langle\texttt{s}\rangle$ token always appears at position 0 0, and we denote its embedding vector by ${\tt ebd}(\texttt{$\langle\texttt{s}\rangle$})$. For simplicity, we abuse the notation and use the sequence itself, $[\texttt{$\langle\texttt{s}\rangle$},v_{1},\ldots,v_{n}]$ where $\{v_{k}\}_{k\in[n]}\subseteq\mathcal{V}$, to represent the embedding of the sequence.

Given an input sequence ${\mathbf{H}}=[\texttt{$\langle\texttt{s}\rangle$},v_{1:n}]\in\mathbb{R}^{D
\times(n+1)}$ with $\langle\texttt{s}\rangle$ as the zeroth token, we define the predicted probability of the next token as ${\sf SoftMax}({\rm TF}({\mathbf{H}})_{n})$, where ${\rm TF}({\mathbf{H}})_{n}\in\mathbb{R}^{D}$ is the last column of ${\rm TF}({\mathbf{H}})\in\mathbb{R}^{D\times(n+1)}$, defined as

$$ ${\rm TF}(\cdot)={\tt attn}(\cdot)+{\tt mlp}(\cdot),~{}~{}{\tt attn}({\mathbf{H }})={\mathbf{V}}{\mathbf{H}}{\sf SoftMax}({\sf mask}({\mathbf{H}}^{\top}{ \mathbf{K}}^{\top}{\mathbf{Q}}{\mathbf{H}})),~{}~{}{\tt mlp}({\mathbf{H}})={ \mathbf{W}}_{2}\mathrm{ReLU}({\mathbf{W}}_{1}{\mathbf{H}}).$ (5) $$

The simplified transformer architecture ${\rm TF}$ is a parallel summation of the ${\tt attn}$ head and the ${\tt mlp}$ layer, with no layer normalization. This parallel summation is a reasonable simplification, as sequential ${\tt attn}$ and ${\tt mlp}$ layers can effectively simulate parallel ${\tt attn}$ and ${\tt mlp}$ operations. Notice that we have redefined the notations of ${\tt attn}$ and ${\tt mlp}$ in this section, which are simplified versions of Eq. (2) and (3).

#### Simplification and reparameterization of the model.

To simplify the analysis of the training dynamics, we further reduce the model by restricting the $({\mathbf{K}},{\mathbf{Q}},{\mathbf{V}},{\mathbf{W}}_{1},{\mathbf{W}}_{2})$ matrices to follow the patterns observed in the later training stages (i.e., after step 200 of the training in Figure 4).

- •
Restricted Attention Pattern. Based on the intuition from Figure 2, we know that eventually only a few attention logits are non-trivial. Thus, we assume that the model has learned the attention pattern by this stage (which is reasonable given that the Backcopy risk is already small after step 200 in Figure 4). We parameterize the attention logits on the $\langle\texttt{s}\rangle$ key-token as $(\alpha_{\texttt{$\langle\texttt{s}\rangle$}};\alpha_{v_{1}};\ldots;\alpha_{v_
{n}})$, restrict the attention logits for any trigger query-token to $(0,\ldots,\lambda,0)$ (where the second-to-last coordinate is $\lambda$), and set all other logits to zero. Specifically, we restrict:
$\displaystyle~{}{\tt ebd}(\texttt{$\langle\texttt{s}\rangle$})^{\top}{\mathbf{
K}}^{\top}{\mathbf{Q}}\cdot{\tt ebd}_{i}(v)=\alpha_{v}\cdot 1\{v\not\in
\mathcal{T}\}~{}~{}~{}\text{for }v\in\mathcal{V},i\in[n],$
(6)
$\displaystyle~{}{\tt ebd}_{i}(\overline{v})^{\top}{\mathbf{K}}^{\top}{\mathbf{
Q}}\cdot{\tt ebd}_{j}(v)=\lambda\cdot 1\{v\in\mathcal{T},i=j-1\}~{}~{}~{}\text
{for }v,\overline{v}\in\mathcal{V},i,j\in[n].$
Notice that this naturally implies $\alpha_{v}=0$ for $v\in\mathcal{T}$.
- •
Restricted Value Pattern. At later stages of the training dynamics, we observe that the value states for each token are nearly a scaled version of the one-hot encoding vector. We assume this observed pattern and parameterize the value state of $v$ by $\xi_{v}\bm{e}_{v}\in\mathbb{R}^{V}$. For the $\langle\texttt{s}\rangle$ token, we parameterize its value state by $\bm{\beta}\in\mathbb{R}^{V}$. Specifically, we restrict
$\displaystyle~{}{\mathbf{V}}\cdot{\tt ebd}(\texttt{$\langle\texttt{s}\rangle$}
)=\bm{\beta}\in\mathbb{R}^{V},$
(7)
$\displaystyle~{}{\mathbf{V}}\cdot{\tt ebd}_{i}(v)=\xi_{v}\bm{e}_{v}\in\mathbb{
R}^{V},~{}~{}~{}\text{with $\xi_{v}=0$ for $v\in\mathcal{T}$, and $\xi_{v}\geq
0
$ for $v\in\mathcal{V}\setminus\mathcal{T}$.}$
- •
MLP Layer Perfectly Predicts the Transition Probability. Notice that the mlp layer handles the Bigram task. By step 200 in Figure 4, the Bigram risk has nearly vanished. Therefore, we assume that the ${\tt mlp}$ layer outputs the Markov transition probabilities $\bm{p}_{v}$ for non-trigger tokens $v$, and zero for trigger tokens. Specifically, we restrict:
${\tt mlp}({\tt ebd}_{i}(v))=\log\bm{p}_{v}\cdot 1\{v\not\in\mathcal{T}\}~{}~{}
~{}\text{for }v\in\mathcal{V}.$
(8)

These reparameterizations are illustrated in Figure 5. Theorem 1 establishes the existence of a transformer architecture that satisfies the restrictions and reparameterizations outlined above. Furthermore, this restricted transformer can generate the ground-truth transitions of the BB model when certain parameters diverge.

###### Theorem 1 (Existence of reparameterization that solves the BB task; informal) .

For any parameters $(\bm{\alpha}\in\mathbb{R}^{V},\bm{\beta}\in\mathbb{R}^{V},\bm{\xi}\in\mathbb{R
}^{V},\lambda\in\mathbb{R})$, there exists a one-layer transformer as described in (5) with weight matrices $({\mathbf{Q}},{\mathbf{K}},{\mathbf{V}},{\mathbf{W}}_{1},{\mathbf{W}}_{2})$ such that Eq. (6), (7), and (8) hold. Furthermore, there exists a sequence of parameters where $\min_{v\in\mathcal{V}}\alpha_{v}\to\infty$, $\min_{v\in\mathcal{V}}\xi_{v}\to\infty$, $\lambda\to\infty$, and $\bm{\beta}=0$, such that this transformer generates the ground-truth transitions of the BB model in the limit.

The formal statement and proof of Theorem 1 are provided in Appendix A.1.

###### Theorem 1 (Existence of reparameterization that solves the BB task; informal) .

#### Dynamic analyses of the reparameterized model.

To analyze the later stage training dynamics, we adopt the reparameterization given in Eq. (6), (7), and (8) as our assumption. We further define $M_{k}=\sum_{i=1}^{n}1\{v_{i}=k\}$, $\bm{M}=(M_{1},\ldots,M_{V})$, and $M=\sum_{k\in\mathcal{V}}M_{k}=n$. Substituting these into Eq. (5), for a non-trigger token $v\in\mathcal{V}\setminus\mathcal{T}$, the output of the attention layer with input sequence ${\mathbf{H}}=[\texttt{$\langle\texttt{s}\rangle$},v_{1:n-1},v]$ is given by

$$ ${\rm TF}({\mathbf{H}})_{n}=\log\bm{p}_{v}+\frac{e^{\alpha_{v}}}{e^{\alpha_{v}} +M}\bm{\beta}+\sum_{k=1}^{V}\frac{M_{k}\xi_{k}}{e^{\alpha_{v}}+M}\cdot\bm{e}_{ k}.$ (9) $$

Therefore, for the non-trigger token $v$, the cross-entropy loss between the true Markov transition $\bm{p}_{v}$ and the predicted transition ${\sf SoftMax}({\rm TF}({\mathbf{H}})_{n})$ is given by

$$ $\mathsf{loss}_{v}(\alpha_{v},\bm{\beta})=\sum_{k=1}^{V}p_{vk}\Big{\{}\log\Big{ [}\sum_{i=1}^{V}p_{vi}\exp\Big{(}\frac{e^{\alpha_{v}}\beta_{i}+M_{i}\xi_{i}}{e ^{\alpha_{v}}+M}\Big{)}\Big{]}-\frac{e^{\alpha_{v}}\beta_{k}+M_{k}\xi_{k}}{e^{ \alpha_{v}}+M}-\log p_{vk}\Big{\}}.$ (10) $$

For simplicity, we neglect the loss on trigger tokens and assume that $(\{M_{i}\}_{i\in[V]},M)$ remain fixed across different positions in the input sequences.(^2^22We note that Reddy () makes a similar simplification in analyzing induction heads.) We then consider the total loss as the average of the losses on each non-trigger token, weighted by its proportion in the stable distribution $\{\pi_{v}\}_{v\in\mathcal{V}}$, given by

$$ $\textstyle\mathsf{loss}(\bm{\alpha},\bm{\beta})=\sum_{v\in\mathcal{V}\setminus \mathcal{T}}\pi_{v}\cdot\mathsf{loss}_{v}(\alpha_{v},\bm{\beta}).$ (11) $$

We assume that $\bm{\xi}$ and $\lambda$ are fixed, and that $\bm{\alpha}$ (the attention logits of the $\langle\texttt{s}\rangle$ token) and $\bm{\beta}$ (the value state norms of the $\langle\texttt{s}\rangle$ token) are trainable variables, as we are interested in the dynamics of the attention logits and value state norm for the $\langle\texttt{s}\rangle$ token. The following theorem illustrates the logarithmic growth of the attention logits $\bm{\alpha}$, the shrinkage of value states $\bm{\beta}$, and the stable phase of these two variables.

###### Theorem 2 .

Consider the gradient flow of the loss function $\mathsf{loss}(\bm{\alpha},\bm{\beta})$. Assume $\xi_{v}\geq 0$ for any $v$ and $\pi_{v}>0$ for any $v\in\mathcal{V}$, and $\{M_{i}\cdot\xi_{i}\}_{i\in\mathcal{V}}$ are not all equal.

- (a)
(Attention logits grow logarithmically, reinforced by small value states) Fix $\bm{\beta}=\beta\cdot\bm{1}$ for a constant $\beta$, and consider the gradient flow over $\bm{\alpha}$. With any initial value $\bm{\alpha}(0)$, there exists $\bm{r}(t)$ with norm uniformly bounded in time, such that
$\textstyle\bm{\alpha}(t)=\frac{1}{2}\log t\cdot\bm{1}+\bm{r}(t).$
(12)
- (b)
(Value state shrinks to a small constant vector, reinforced by large attention logits) Fix $\bm{\alpha}=\alpha\cdot\bm{1}$ for a constant $\alpha$, define $\overline{\beta}(0)=V^{-1}[\sum_{v}\beta_{v}(0)]$ and $\overline{B}=V^{-1}[\sum_{v}M_{v}\xi_{v}]$. Consider the gradient flow over $\bm{\beta}$. As $t\to\infty$, we have
$\bm{\beta}(t)\to\bm{\beta}^{\star}=[\overline{\beta}(0)+e^{-\alpha}\overline{B
}]\cdot\bm{1}-e^{-\alpha}\cdot\bm{M}\circ\bm{\xi}.$
(13)
- (c)
(Stable phase: Sink-logits concentration) Consider the gradient flow over the variables $(\bm{\alpha},\bm{\beta})$. Any vector of the following form
$\bm{\alpha}=\alpha\cdot\bm{1},\quad\bm{\beta}=c\cdot\bm{1}-e^{-\alpha}\cdot\bm
{M}\circ\bm{\xi},~{}~{}~{}\alpha,c\in\mathbb{R}$
(14)
is a stationary point. These are all global minimizers of $\mathsf{loss}(\bm{\alpha},\bm{\beta})$.

The proof of Theorem 2 is provided in Appendix A.2, A.3, and A.4. We offer two key remarks: (1) As $\alpha_{v}\to\infty$, a Taylor expansion of the gradient $\partial\mathsf{loss}/\partial\alpha_{v}$ suggests that $\mathrm{d}\alpha_{v}/\mathrm{d}t\propto\exp(-2\alpha_{v})$, which leads to the logarithmic growth of $\alpha_{v}$. Similar logarithmic growth has been reported in the literature under different setups (Tian et al., ; Zhu et al., );
(2) The stable phase described in Theorem 2(c) seems to imply that the system can remain stable without attention sinks, as it does not require $\alpha$ to be large. However, in practice, models trained on the BB task tend to converge to a stable phase where $\alpha$ is relatively large.

###### Theorem 2 .

#### The formation of attention sinks and value-state drains.

Below, we explain how Theorem 2 reveals the mutual reinforcement mechanism behind the formation of attention sinks and value-state drains.

- (a)
When the value states of the $\langle\texttt{s}\rangle$ token are small and constant, $\bm{\beta}=\beta\cdot\bm{1}$, Theorem 2(a) shows that the attention logits on the $\langle\texttt{s}\rangle$ token $\bm{\alpha}(t)\approx\alpha(t)\bm{1}$ for $\alpha(t)=(1/2)\log t$, grow logarithmically. This demonstrates that the presence of a small constant value state ($\bm{\beta}=\beta\cdot\bm{1}$) reinforces the formation of attention sinks ($\bm{\alpha}(t)\approx\alpha(t)\cdot\bm{1}$ for $\alpha(t)$ increases logarithmically).
- (b)
When the attention logits of the $\langle\texttt{s}\rangle$ token are large and constant, $\bm{\alpha}=\alpha\cdot\bm{1}$ for $\alpha\to\infty$, Theorem 2(b) shows that the value states of the $\langle\texttt{s}\rangle$ token $\bm{\beta}(t)\to\overline{\beta}(0)\cdot\bm{1}$. Starting with a random Gaussian initialization for $\bm{\beta}(0)$, we have $\|\bm{\beta}(t)\|_{2}\approx\|\overline{\beta}(0)\cdot\bm{1}\|_{2}\approx\|\bm
{\beta}(0)\|_{2}/\sqrt{V}$, where $V$ is the vocabulary size, typically large. This indicates that attention sinks ($\bm{\alpha}=\alpha\cdot\bm{1}$ for large $\alpha$) reinforces the formation of value-state drains ($\bm{\beta}(t)\to\beta\cdot\bm{1}$ for small $\beta$).
- (c)
In the later stages of the dynamics, both the attention logits and value states of the $\langle\texttt{s}\rangle$ token stabilize, as described in 2(c). The attention logits remain constant at $\bm{\alpha}=\alpha\cdot\bm{1}$ with large $\alpha$, while the value states become small, $\bm{\beta}=[\overline{\beta}(0)+e^{-\alpha}\overline{B}]\cdot\bm{1}-e^{-\alpha
}\cdot\bm{M}\circ\bm{\xi}$.

Based on these theoretical insights, we summarize the dynamical mechanism underlying attention sinks and value-state drains: For any attention head given a specific prompt, if the model can accurately predict the next token without using the attention head, but adding any value state from previous tokens—except for certain special tokens—worsens the prediction, the attention head will become dormant, forming an attention sink at those special tokens. This phenomenon is induced by the mutual reinforcement mechanism, as described below:

Figure: Figure 6: Mutual reinforcement mechanism
Refer to caption: x11.png

###### Claim 2 (Mutual reinforcement mechanism) .

#### Experimental verification of the quantitative prediction.

Revisiting Figure 4, which illustrates the dynamics of a single-layer transformer model trained with Adam on the BB task, we observe that $\Delta\text{logit}_{\cdot,\texttt{$\langle\texttt{s}\rangle$}}$ exhibits growth rates consistent with Theorem 2. In this context, $\Delta\text{logit}_{\cdot,\texttt{$\langle\texttt{s}\rangle$}}$ corresponds to $\alpha$, as all other attention logits are assumed to be zero under the assumptions of Theorem 2. When plotted on a logarithmic scale, the $\Delta\text{logit}_{\cdot,\texttt{$\langle\texttt{s}\rangle$}}$ curve grows approximately linearly between 1,000 and 10,000 steps, then accelerates before stabilizing around 100,000 steps. Meanwhile, the norm of the value state $\|{\tt Val}_{\texttt{$\langle\texttt{s}\rangle$}}\|_{2}$ decreases monotonically. The simultaneous increase in attention weights and decrease in value-state norms demonstrate the mutual reinforcement mechanism during the training process.

To further validate that Theorem 2 accurately captures the dynamics of the original model, we constructed a simplified model based on Eq. (6), (7), and (8), and trained the parameters $(\bm{\alpha}\in\mathbb{R}^{V},\bm{\beta}\in\mathbb{R}^{V},\bm{\xi}\in\mathbb{R
}^{V},\lambda\in\mathbb{R})$ using Adam. The resulting training curves closely resemble those of the one-layer transformer, also displaying the mutual reinforcement mechanism. A detailed description of the experiment can be found in Appendix C.3.

#### Generality of the theoretical prediction.

Although Theorem 2 focuses on a specific BB task with a simplified architecture and loss function, the underlying principles are broadly applicable to more general settings. In particular, we expect that the formation of extreme tokens in LLMs follows a similar mutual reinforcement mechanism. Indeed, Theorem 2 is essentially based on the following two key assumptions: (1) even with a specific attention head ${\tt attn}$ zeroed out, the LLM can still accurately predict the next token, implying that the attention head is better off dormant; and (2) for the attention head ${\tt attn}$, value states of previous tokens—except for certain special tokens—remain relevant for specific tasks and therefore do not vanish. Under these assumptions, we anticipate the formation of attention sinks and value-state drains for the attention head ${\tt attn}$ and such special tokens. In Section 3, we explore how these phenomena are formed during the training dynamics of LLMs, finding that the empirical results align with the theory.

Figure: (a) ReLU attention
Refer to caption: x12.png

#### Replacing SoftMax by ReLU attention removes attention sinks and value-state drains.

As a consequence of our theory, we predict that training using ReLU attention in place of SoftMax attention will prevent the mutual reinforcement mechanism. Without SoftMax, the training dynamics no longer push the attention weights toward the $\langle\texttt{s}\rangle$ token, which remains zero throughout training. In the absence of attention sinks, the dynamics no longer push down the value state norm, and the mutual reinforcement mechanism breaks. Figure 7 presents the training dynamics on the BB task using ReLU instead of SoftMax attention, showing that both the Bigram and Backcopy risk converge to the Bayes risk after 200 training steps, but the attention logits of $\langle\texttt{s}\rangle$ do not increase, and the value state does not shrink, confirming our prediction.

### 2.3 The emergence of residual-state peaks

In this section, we experimentally investigate the residual-state peaks phenomenon. We observe that no residual-state peaks occur in the single-layer transformer trained on the BB task. To explore this further, we train slightly deeper transformers on the BB task and track the residual state norm after layer 0 0. We observe that two-layer models do not exhibit residual-state peaks, while models with three or more layers do.
Additional experimental results are provided in Appendix B.1 and B.2.

#### Massive residual state at layer 0 output induces attention sinks and value-state drains in the middle layer.

To investigate the relationship between massive residual states and attention sinks, we train on the BB task using the “attn+mlp+attn+mlp+mlp” model, which is the minimal structure that shows the massive residual states phenomena. We perform intervention by analyzing how the model’s behavior changes after zeroing out layer 0 (the first “attn+mlp” block). Before and after zeroing, we compute the difference in $\|{\tt Res}_{\texttt{$\langle\texttt{s}\rangle$}}\|$ and $\text{Mean}_{v}[\|{\tt Res}_{v}\|]$ at the layer 0 output, and compute $\text{logit}_{\cdot,\texttt{$\langle\texttt{s}\rangle$}}$ and $\|{\tt Val}_{\texttt{$\langle\texttt{s}\rangle$}}\|$ in the middle layer. After zeroing out, the residual state norm becomes non-massive, and attention logits and the value state norm return to a normal level. This confirms that the residual-state peak contributes to the attention sink and value-state-drain phenomena in the middle layer of pre-trained transformers.

#### Linear growth of residual-state norm with Adam training.

Figure 7 shows the residual-state norms of the $\langle\texttt{s}\rangle$ token at the layer 0 output of three-layer transformers during pre-training on the BB task. The results indicate that training the transformer with Adam leads to a linear increase in residual state norms.

#### Switching from Adam to SGD and switching from SoftMax to ReLU attention eliminates the residual-state peaks.

Figure 7 also illustrates the dynamics of residual-state norms in other training setups. When switching the training algorithm from Adam to SGD, attention sinks remain, but residual-state peaks disappear. Similarly, switching to ReLU attention, which lacks the mutual reinforcement mechanism, also eliminates residual-state peaks. These findings highlight the dependence of residual-state peaks on SoftMax attention and the Adam optimization algorithm. We propose a potential explanation of this phenomenon in Appendix B.3.

## 3 Extreme-token Phenomena in pretrained LLMs

In this section, we investigate extreme-token phenomena in open-source pretrained LLMs. In Section 3.1, we analyze the static behavior of these phenomena in Llama 2-7B-Base (Touvron et al., ), confirming the existence of the active-dormant mechanism in LLMs. Notably, we identify a specific head that is active on GitHub samples but dormant on Wikipedia samples. In Section 3.2, we examine the dynamic behavior of extreme-token phenomena during the pretraining of OLMo-7B (Groeneveld et al., ). We show that the attention logits, value states norm, and residual states norm of the sink token(s) in OLMo reflect behavior similar to that of the simpler BB model. Specifically, the simultaneous formation of attention sinks and value-state drains gives evidence for the mutual reinforcement mechanism.

### 3.1 Active-dormant mechanism in LLMs

Figure: (a) Attention weights for GitHub/Wikipedia data
Refer to caption: extracted/5985254/Figures/L16_H25/attn_github_head25.png

Our study of the BB model leads to the following prediction with respect to the extreme-token phenomena, which we hypothesize also applies to LLMs:

Attention heads are controlled by an active-dormant mechanism (cf. Claim 1). The presence of attention sinks and value-state drains indicates that an attention head is in a dormant phase.

This hypothesis suggests that in LLMs, whether an attention head becomes a sink depends on the context. Specifically, the attention head may become entirely irrelevant for selecting the next tokens in certain contexts or tasks, but not in others. When this irrelevance occurs, the attention head transitions into an attention sink. This hypothesis was confirmed in small transformers and the BB task, as demonstrated in Section 2.

Accordingly, we aim to identify instances of attention heads in pretrained LLMs that exhibit this active-dormant behavior, i.e., heads that are dormant in some domains but active in others. In Figure 8, we display a particular attention head—Layer 16 Head 25 (L16H25) of Llama 2-7B-Base (Touvron et al., )—which demonstrates a clear active-dormant distinction across two distinct contexts (e.g., tokens from the GitHub subset versus the Wikipedia subset of RedPajama (Computer, )). While many attention heads show similar context-dependent behavior (see Appendix D), we focus on this one because the conditions for its activation are straightforward and interpretable, whereas other heads may have more nuanced criteria.

Figure 8(a) shows the attention maps of L16H25 on samples from both the GitHub and Wikipedia subsets of RedPajama. It demonstrates that L16H26 is dormant (i.e., an attention sink) on samples from Wikipedia, which resemble prose, and active (i.e., not an attention sink) on samples from GitHub, which resemble code. Additionally, Figure 8(b) compares the loss difference when L16H25 is zeroed out for prompts from both domains. The results show that zeroing out this head significantly decreases model performance on GitHub sequences, while having minimal impact on Wikipedia sequences. This observation also confirms the head behaves as dormant in some contexts and active in others—in some contexts, removing this head has no effect on model performance, while in others, its removal causes significant performance drops.

### 3.2 Extreme-token phenomena along training dynamics of LLMs

Figure: (a) Attention sink dynamics
Refer to caption: x16.png

Figure: (a) Logit dynamics
Refer to caption: x19.png

Our study of the BB model leads to the following prediction about the dynamical behavior of the extreme-token phenomena, which we hypothesize also applies to LLMs:

Attention heads undergo an attention-increasing and value-state-shrinking phase driven by the mutual reinforcement mechanism (cf. Claim 2). This is followed by a stable phase, where all non-trigger tokens have large, nearly identical attention logits on the extreme token. Simultaneously, the residual state norms of the extreme tokens increase linearly during pretraining.

We confirm these predictions below. To observe the training dynamics of a large-scale LLM, we use the setup of OLMo-7B-0424 (Groeneveld et al., ) (henceforth just referred to as OLMo), which provides open-sourced weights at various stages of their training.(^3^33We did not analyze Llama for dynamics, as they do not provide open-source intermediate checkpoints along pretraining.) For our analysis, we inspect OLMo at multiple checkpoints: every 500 steps for the first 10,000 steps, then at 25,000 steps, 50,000 steps, and every 50,000 steps up to 449,000 steps (approximately the end of their training).(^4^44For the single 150,000-step checkpoint, we observed that its statistics were outliers, which we hypothesize is due to a system failure. We address this by using the average of nearby checkpoints to represent its statistics.) The input we use for this analysis is again “Summer is warm$\langle$$\mathtt{period}$$\rangle$ Winter is cold$\langle$$\mathtt{period}$$\rangle$”(^5^55Note that OLMo does not have a $\langle\texttt{s}\rangle$ token, but attention sinks still form in the majority of heads. In particular, the first token always behaves as an attention sink. We discuss this further in Section G.2.) In this prompt, the “$\mathrm{Delim}$” token, namely “$\langle$$\mathtt{period}$$\rangle$”, also becomes a sink token along with $\langle\texttt{s}\rangle$. We believe this occurs because the period is not semantically meaningful and is not useful for predicting future tokens (cf.  Section G.2)

Figure 9 illustrates the dynamics of attention weights, value state norms, and the residual state norms for attention heads in Layer 24 of OLMo. The figure shows that the average attention on extreme tokens ($\langle\texttt{s}\rangle$ and $\mathrm{Delim}$) increases rapidly at the beginning of training before stablizing, while the value state norms of these extreme tokens decrease during training steps 10k-100k. The synchronized evolution of attention weights and value state norms aligns with the prediction of the mutual reinforcement mechanism. Additionally, the residual states of $\langle\texttt{s}\rangle$ increase linearly, while those of other tokens converge to a small number. Figure 10 provides a more detailed examination of the attention logits in Layer 24 of OLMo. Figure 10(a) presents the dynamics of the difference in attention logits, showing that $\Delta\text{logit}_{\cdot,\texttt{$\langle\texttt{s}\rangle$}}$ increase during training steps 10k-100k, matching the decreasing phase of the value states.
Figure 10(b) also demonstrates the sink-logits concentration phenomenon. Specifically, it shows that the sink logits will eventually converge to a stable phase, in which logits corresponding to the key of the sink token and queries of all non-sink tokens are nearly identical. These findings coincide with the dynamical behavior predicted by the BB model, as outlined in Theorem 2(c) and corroborated by the experimental results in Figure 4.

## 4 Conclusions

In this work, we investigated the extreme-token phenomena, specifically attention sinks, value-state drains, and residual-state peaks. We analyzed simple transformers trained on the Bigram-Backcopy (BB) task, both theoretically and empirically, demonstrating that these models exhibit the same extreme-token phenomena observed in large language models (LLMs). Building on the insights from the BB task, we made several detailed predictions about the behavior of extreme-token phenomena in LLMs. In particular, we identified the active-dormant mechanism governing attention heads in both the BB model and LLMs, with attention sinks and value-state drains serving as indicators of dormant phase, and a mutual reinforcement mechanism that induces these phenomena during pretraining. Using insights from these mechanisms, we applied simple modifications to the model architecture and optimization procedure, effectively mitigating the extreme-token phenomena in the BB model. Overall, our work uncovers the underlying mechanisms of extreme-token phenomena and suggests potential pathways to mitigate these issues during LLM pretraining.

We believe the most compelling direction for future work is to explore whether eliminating the extreme-token phenomena is essential or beneficial for building powerful transformer-based LLMs. While it is possible to mitigate these phenomena through simple modifications to the architecture or training algorithms, it remains unclear whether their elimination significantly improves downstream tasks such as inference and quantization. Given the resource-intensive nature of pretraining large-scale LLMs, we anticipate that pretraining a model at the scale of GPT-2 could both provide valuable insight into this issue and help point the way to architectures that can reduce the pretraining burden.

## Acknowledgements

TG thanks Yaodong Yu, Licong Lin, and Ruiqi Zhang for insightful discussions.
YB thanks Caiming Xiong and Huan Wang for the many insightful discussions in the early stages of this work.
This project is supported by NSF DMS-2210827, CCF-2315725, CAREER DMS-2339904, ONR N00014-24-S-B001, a UC Berkeley College of Engineering fellowship, an Amazon Research Award, a Google Research Scholar Award, an Okawa Foundation Research Grant, and the European Union (ERC-2022-SYG-OCEAN-101071601).

## Appendix A Proofs of Theorem 1 and 2

We introduce new notations that are frequently used in the proofs. Recall that in Eq. (11), we used $\{\pi_{v}\}_{v\in\mathcal{V}}$ to denote the stable distribution across all tokens. We further define the stable distribution excluding trigger tokens as follows:

$$ $\tilde{\bm{\pi}}\in\mathbb{R}^{V},~{}~{}~{}\tilde{\pi}_{i}=\pi_{i}\bm{1}\{i\in \mathcal{V}\setminus\mathcal{T}\}.$ (15) $$

Section 2.2 defines the bigram transition probability in the Bigram-Backcopy task as $p_{vk}=\sf{P}(k\mid v)$.
We further define the bigram transition probability matrix as

$$ $\mathbf{P}=\left(\begin{matrix}p_{11}&\ldots&p_{1V}\\ \vdots&\ddots&\vdots\\ p_{V1}&\ldots&p_{VV}\\ \end{matrix}\right)=\left(\begin{matrix}\bm{p}_{1}^{\top}\\ \vdots\\ \bm{p}_{V}^{\top}\\ \end{matrix}\right).$ (16) $$

Given a token $v$, define the predicted probability at token $v$ as the logit output passed through the softmax activation. Let ${\mathbf{H}}=[\texttt{$\langle\texttt{s}\rangle$};v_{1:n-1};v]$. Using the form of ${\rm TF}({\mathbf{H}})_{n}$ defined in Eq. (9), we denote

$$ $\bm{l}_{v}={\sf SoftMax}({\rm TF}({\mathbf{H}})_{n})=(l_{v1},\ldots,l_{vV}),~{ }~{}~{}\text{with}~{}~{}~{}l_{vi}=\frac{p_{vi}\exp\Big{[}\frac{M_{i}\xi_{i}+e^ {\alpha}\beta_{i}}{e^{\alpha}+M}\Big{]}}{\sum_{k=1}^{V}p_{vk}\exp\Big{[}\frac{ M_{k}\xi_{k}+e^{\alpha}\beta_{k}}{e^{\alpha}+M}\Big{]}}.$ (17) $$

Similar to Eq. (16), we define the full output probability matrix as

$$ $\mathbf{L}=\left(\begin{matrix}l_{11}&\ldots&l_{1V}\\ \vdots&\ddots&\vdots\\ l_{V1}&\ldots&l_{VV}\\ \end{matrix}\right)=\left(\begin{matrix}\bm{l}_{1}^{\top}\\ \vdots\\ \bm{l}_{V}^{\top}\\ \end{matrix}\right).$ (18) $$

Using the notation $\bm{l}_{v}$ and $\tilde{\pi}_{v}$, we can rewrite the loss functions defined in Eq. (10) and Eq. (11) as follows:

$$ $\mathsf{loss}_{v}(\alpha_{v},\bm{\beta})=-\sum_{k=1}^{V}p_{vk}\log l_{vk},~{}~ {}~{}~{}\mathsf{loss}_{v}(\bm{\alpha},\bm{\beta})=\sum_{v=1}^{V}\tilde{\pi}_{v }\mathsf{loss}_{v}(\alpha_{v},\bm{\beta}).$ (19) $$

We always have that $\sum_{k}p_{vk}=1$ and $\sum_{k}l_{vk}=1$. The total variation norm and KL-divergence are then defined as:

$$ $\|\bm{p}_{v}-\bm{l}_{v}\|_{\text{TV}}=\sum_{k}|p_{vk}-l_{vk}|,~{}~{}~{}~{} \text{KL}(\bm{p}_{v}~{}||~{}\bm{l}_{v})=-\sum_{k}p_{vk}\log(l_{vk}/p_{vk}).$ (20) $$

Given any vector $\bm{u}=[u_{1};\ldots;u_{d}]$, define the corresponding diagonal matrix as

$$ $\operatorname{diag}(\bm{u})=\left(\begin{matrix}u_{1}&0&\ldots&0\\ \vdots&\ddots&&\vdots\\ \vdots&&\ddots&\vdots\\ 0&\ldots&0&u_{d}\\ \end{matrix}\right).$ $$

Given any $\bm{p}_{v}$ defined in Eq. (16), denote

$$ ${\mathbf{G}}_{v}^{\mathbf{P}}=\operatorname{diag}(\bm{p}_{v})-\bm{p}_{v}\bm{p} _{v}^{\top},\quad{\mathbf{G}}_{v}^{\mathbf{L}}=\operatorname{diag}(\bm{l}_{v}) -\bm{l}_{v}\bm{l}_{v}^{\top}.$ (21) $$

We now present technical lemmas concerning ${\mathbf{G}}_{v}^{\mathbf{P}}$ and ${\mathbf{G}}_{v}^{l}$.

###### Lemma A.0 .

The matrices ${\mathbf{G}}^{\mathbf{P}}_{v}\in\mathbb{R}^{V\times V}$ and ${\mathbf{G}}^{\mathbf{L}}_{v}\in\mathbb{R}^{V\times V}$ are positive semi-definite for any $v\in\mathcal{V}$.

###### Proof of Lemma A.0 .

Since $\sum_{k=1}^{V}p_{vk}=1$ and $\sum_{k=1}^{V}l_{vk}=1$ for any $v$, we have that

$$ $\displaystyle({\mathbf{G}}^{\mathbf{P}}_{v})_{ii}=p_{i}-p_{i}^{2}$ $\displaystyle=p_{i}(\sum_{k\neq i}p_{k})\geq\sum_{k\neq i}|({\mathbf{G}}^{ \mathbf{P}}_{v})_{ik}|,$ $\displaystyle({\mathbf{G}}^{\mathbf{L}}_{v})_{ii}=l_{i}-l_{i}^{2}$ $\displaystyle=l_{i}(\sum_{k\neq i}l_{k})\geq\sum_{k\neq i}|({\mathbf{G}}^{ \mathbf{L}}_{v})_{ik}|.$ $$

This shows that both ${\mathbf{G}}^{\mathbf{P}}_{v}$ and ${\mathbf{G}}^{\mathbf{L}}_{v}$ are diagonally dominant matrices. By Corollary 6.2.27 in Horn and Johnson (), they are positive semi-definite.
∎

###### Lemma A.0 .

Suppose that $\tilde{\pi}_{v}>0$ for any $v\in\mathcal{V}\setminus\mathcal{T}$. For any $\bm{\eta}\in\mathbb{R}^{V}$ with $\bm{\eta}\perp\bm{1}$, there exists $\omega>0$ such that

$$ $\bm{\eta}^{\top}\Big{[}\sum_{k=1}^{V}\tilde{\pi}_{k}{\mathbf{G}}_{k}^{\mathbf{ P}}\Big{]}\bm{\eta}\geq\omega\|\bm{\eta}\|_{2}^{2}.$ $$

###### Proof of Lemma A.0 .

Denote the null spaces of ${\mathbf{G}}_{v}^{\mathbf{P}}$ for $v\in\mathcal{V}$ as $\sf S_{v}$. We solve for each $\sf S_{v}$. Setting ${\mathbf{G}}_{v}^{\mathbf{P}}\bm{\eta}=0$ gives that

$$ $[p_{vj}-p_{vj}(\sum_{k}p_{vk})]\eta_{j}=0~{}\text{for any $j\in\mathcal{V}$.}$ $$

If $p_{vj}\neq 0$, we divide each side with $p_{vj}$ and get that $\eta_{j}=\sum_{k}p_{vk}\eta_{k}$. As a result, we get that

$$ $\text{\sf{S}}_{v}={\left\{\bm{\eta}\mid\eta_{j}\text{ is constant for }p_{vj} \neq 0\right\}}.$ $$

Since all ${\pi}_{k}>0$, for any $k\in\mathcal{V}\setminus\mathcal{T}$, there is $v\in\mathcal{V}\setminus\mathcal{T}$ such that $p_{vk}>0$, we get that

$$ $\cap_{v\in\mathcal{V}\setminus\mathcal{T}}\text{\sf{S}}_{v}={\left\{c\cdot\bm{ 1}\mid c\in\mathbb{R}\right\}}.$ $$

Since $\bm{\eta}\perp\bm{1}$, we get that $\bm{\eta}\perp\cap_{v\in\mathcal{V}\setminus\mathcal{T}}\sf{S}_{v}$. We denote the minimal non-zero eigenvalues of ${\mathbf{G}}_{v}^{\mathbf{L}}$ for $v\in\mathcal{V}\setminus\mathcal{T}$ as $\lambda$. We get that

$$ $\bm{\eta}^{\top}\Big{[}\sum_{k=1}^{V}\tilde{\pi}_{k}{\mathbf{G}}_{k}^{\mathbf{ P}}\Big{]}\bm{\eta}\geq\Big{[}\min_{v\in\mathcal{V}\setminus\mathcal{T}}\tilde {\pi}_{v}\Big{]}\lambda\|\bm{\eta}\|_{2}^{2}.$ $$

Setting $\omega=\lambda\cdot\min_{v\in\mathcal{V}\setminus\mathcal{T}}\tilde{\pi}_{v}>0$, this proves Lemma A.0.
∎

###### Lemma A.0 .

Given $\omega$ defined in Lemma A.0,
suppose that

$$ $\max_{v,k}\left|p_{vk}-l_{vk}\right|=\delta\leq\min{\left\{\omega/(6V),1\right \}}.$ (22) $$

For any $\bm{\eta}\in\mathbb{R}^{V}$ with $\bm{\eta}\perp\bm{1}$, we have that

$$ $\bm{\eta}^{\top}\Big{[}\sum_{k=1}^{V}\tilde{\pi}_{k}{\mathbf{G}}_{k}^{\mathbf{ L}}\Big{]}\bm{\eta}\geq\frac{\omega}{2}\|\bm{\eta}\|_{2}^{2}.$ $$

###### Proof of Lemma A.0 .

Denote $\delta=\max_{v,k}\left|p_{vk}-l_{vk}\right|$. Suppose that $\delta\leq 1$.
For any $k\in\mathcal{V}\setminus\mathcal{T}$, we can verify that

$$ $\Big{|}({\mathbf{G}}_{k}^{\mathbf{P}})_{ij}-({\mathbf{G}}_{k}^{\mathbf{L}})_{ ij}\Big{|}\leq 3\delta,$ $$

for any $i,j\in[V]$.
We denote

$$ ${\mathbf{E}}=\sum_{k=1}^{V}\tilde{\pi}_{k}{\mathbf{G}}_{k}^{\mathbf{P}}-\sum_{ k=1}^{V}\tilde{\pi}_{k}{\mathbf{G}}_{k}^{\mathbf{L}}.$ $$

Therefore, $|{\mathbf{E}}_{ij}|\leq 3\delta$ for any $i,j\in[V]$.
This means that

$$ $\bm{\eta}^{\top}{\mathbf{E}}\bm{\eta}\leq\|{\mathbf{E}}\|_{2}\|\bm{\eta}\|_{2} ^{2}\leq\|{\mathbf{E}}\|_{F}\|\bm{\eta}\|_{2}^{2}\leq V\cdot 3\delta\cdot\|\bm {\eta}\|_{2}^{2}.$ $$

As a result, when $\delta\leq\min{\left\{\omega/(6V),1\right\}}$, we get that

$$ $\bm{\eta}^{\top}\Big{[}\sum_{k=1}^{V}\tilde{\pi}_{k}{\mathbf{G}}_{k}^{\mathbf{ L}}\Big{]}\bm{\eta}\geq\omega\|\bm{\eta}\|_{2}^{2}-\bm{\eta}^{\top}{\mathbf{E} }\bm{\eta}\geq\frac{\omega}{2}\|\bm{\eta}\|_{2}^{2}.$ $$

This proves Lemma A.0.
∎

###### Lemma A.0 .

###### Proof of Lemma A.0 .

###### Lemma A.0 .

###### Proof of Lemma A.0 .

###### Lemma A.0 .

###### Proof of Lemma A.0 .

### A.1 Proof of Theorem 1

We denote the hidden dimension as $d$ and the sequence length as $N$. Recall that the token $v$ at position $i$ is encoded as ${\tt ebd}_{i}(v)$. We begin with the assumption regarding the transformer’s embedding dimension:

###### Assumption .

We have ${\left\{{\tt ebd}_{0}(\texttt{$\langle\texttt{s}\rangle$})\right\}}\cup{\left
\{{\tt ebd}_{i}(v)\right\}}_{i\in\{0\}\cup[N-1],v\in\mathcal{V}}\subseteq
\mathbb{R}^{d}$, where the embedding dimension $d\geq VN+1$.

Assumption Assumption requires a large embedding dimension $d\geq VN+1$. This assumption is used to ensure that there are enough orthonormal bases in the embedding space. Given the fact that there are $\text{O}(\exp(d))$ approximately linearly independent vectors for large $d$ (Vershynin, ), it is possible to relax the assumption to be $d\gg\log(VN)$. However, since Assumption Assumption pertains only to the construction of $\lambda$ for trigger tokens and is unrelated to Theorem 2, we adopt it to simplify the proof of Theorem 1.

###### Theorem A.0 (Formal statement of Theorem 1 ) .

Let Assumption Assumption hold. For any parameters $(\bm{\alpha}\in\mathbb{R}^{V},\bm{\beta}\in\mathbb{R}^{V},\bm{\xi}\in\mathbb{R
}^{V},\lambda\in\mathbb{R})$, there exists a one-layer transformer (5) with weight matrices $({\mathbf{Q}},{\mathbf{K}},{\mathbf{V}},{\mathbf{W}}_{1},{\mathbf{W}}_{2})$ such that Eq. (6), (7), and (8) hold. Consider the Bigram-Backcopy task, where given an input ${\mathbf{H}}=[\texttt{$\langle\texttt{s}\rangle$};v_{1:n-1},v]$, the ground-truth transition gives ${\sf P}(v^{\prime}\mid{\mathbf{H}})=p_{vv^{\prime}}$ for $v\in\mathcal{V}\setminus\mathcal{T}$, and ${\sf P}(v^{\prime}\mid{\mathbf{H}})=1\{v^{\prime}=v_{n-1}\}$ for $v\in\mathcal{T}$.
There exists a sequence $\min_{v\in\mathcal{V}}\alpha_{v}\to\infty$, $\min_{v\in\mathcal{V}}\xi_{v}\to\infty$, $\lambda\to\infty$, and $\bm{\beta}=\bm{0}$ such that this transformer generates the ground-truth transition in the limit, i.e.,

$$ ${\sf SoftMax}({\rm TF}({\mathbf{H}})_{n})\to{\sf P}(\,\cdot\,|{\mathbf{H}}).$ (23) $$

###### Proof of Theorem A.0 .

Step 1. Construction for the attention head. We let $\{{\tt ebd}_{0}(\texttt{$\langle\texttt{s}\rangle$})\}\cup\{{\tt ebd}_{i}(v)\}
_{i\in\{0\}\cup[N-1],v\in\mathcal{V}}\cup{\left\{\bm{e}_{v}\right\}}_{v\in
\mathcal{V}}$ to be a set of orthonormal basis in $\mathbb{R}^{d}$, and denote ${\left\{\bm{\eta}_{i}\right\}}_{i\in\{0\}\cup[N-1]}\subseteq\mathbb{R}^{d}$ by a set of orthonormal basis in $\mathbb{R}^{d}$ (the existence is guaranteed by Assumption Assumption). Therefore, for any parameters $(\bm{\alpha}\in\mathbb{R}^{V},\bm{\beta}\in\mathbb{R}^{V},\bm{\xi}\in\mathbb{R
}^{V},\lambda\in\mathbb{R})$, there exists a query matrix ${\mathbf{Q}}\in\mathbb{R}^{d\times N}$ such that

$$ $\displaystyle~{}{\mathbf{Q}}\cdot{\tt ebd}_{i}(v)=\lambda\bm{\eta}_{i-1}~{}~{} ~{}\text{for }i>1,~{}~{}v\in\mathcal{T},$ (24) $\displaystyle~{}{\mathbf{Q}}\cdot{\tt ebd}_{i}(v)=\alpha_{v}\bm{\eta}_{0}~{}~{ }~{}\text{for }i>0,~{}~{}v\in\mathcal{V}\setminus\mathcal{T}.$ $$

Meanwhile, there is a key matrix ${\mathbf{K}}\in\mathbb{R}^{d\times N}$ such that

$$ $\displaystyle~{}{\mathbf{K}}\cdot{\tt ebd}_{i}(v)=\bm{\eta}_{i}~{}~{}~{}\text{ for }i>0,~{}~{}v\in\mathcal{V},$ (25) $\displaystyle~{}{\mathbf{K}}\cdot{\tt ebd}_{0}(\texttt{$\langle\texttt{s} \rangle$})=\bm{\eta}_{0}.$ $$

Denote ${\left\{\bm{e}_{v}\right\}}_{v\in\mathcal{V}}$ as an orthonormal basis in $\mathbb{R}^{V}$.
There is a matrix ${\mathbf{V}}\in\mathbb{R}^{d\times V}$ such that

$$ $\displaystyle~{}{\mathbf{V}}\cdot{\tt ebd}_{i}(v)=\xi_{v}\bm{e}_{v}\in\mathbb{ R}^{V},~{}~{}~{}\text{with $\xi_{v}=0$ for $v\in\mathcal{T}$, and $\xi_{v}\geq 0 $ for $v\in\mathcal{V}\setminus\mathcal{T}$.}$ (26) $\displaystyle~{}{\mathbf{V}}\cdot{\tt ebd}_{0}(\texttt{$\langle\texttt{s} \rangle$})=\bm{\beta}\in\mathbb{R}^{V}.$ $$

This construction matches Eq. (6) and (7).

As a result, for $v_{n}\in\mathcal{V}\setminus\mathcal{T}$, by Eq. (5), denoting ${\mathbf{H}}=[\texttt{$\langle\texttt{s}\rangle$};v_{1:n-1};v_{n}]$ and ${\tt attn}({\mathbf{H}})_{n}$ to be the last column of ${\tt attn}({\mathbf{H}})$, we have

$$ $\displaystyle{\tt attn}({\mathbf{H}})_{n}=$ $\displaystyle~{}\sum_{i=0}^{n}\frac{\exp[{\tt ebd}_{n}(v_{n})^{\top}{\mathbf{Q }}^{\top}{\mathbf{K}}\cdot{\tt ebd}_{i}(v_{i})]{\mathbf{V}}\cdot{\tt ebd}_{i}( v_{i})}{\sum_{j=0}^{n}\exp[{\tt ebd}_{n}(v_{n})^{\top}{\mathbf{Q}}^{\top}{ \mathbf{K}}\cdot{\tt ebd}_{j}(v_{j})]}$ $\displaystyle=$ $\displaystyle~{}\frac{\exp[\alpha_{v_{n}}\bm{\eta}_{0}^{\top}\bm{\eta}_{0}] \cdot{\bm{\beta}}+\sum_{i=1}^{n}{\left\{\exp[\alpha_{v_{n}}\bm{\eta}_{0}^{\top }\bm{\eta}_{i}]\xi_{v_{i}}\cdot\bm{e}_{v_{i}}\right\}}}{\exp[\alpha_{v_{n}}\bm {\eta}_{0}^{\top}\bm{\eta}_{0}]+\sum_{j=1}^{n}\exp[\alpha_{v_{n}}\bm{\eta}_{0} ^{\top}\bm{\eta}_{j}]}$ $\displaystyle=$ $\displaystyle~{}\frac{e^{\alpha_{v_{n}}}}{e^{\alpha_{v_{n}}}+n}\cdot{\bm{\beta }}+\sum_{i=1}^{n}\frac{1}{e^{\alpha_{v_{n}}}+n}\cdot\xi_{v_{i}}\bm{e}_{v_{i}}.$ $$

For $v_{n}\in\mathcal{T}$, we have

$$ $\displaystyle{\tt attn}({\mathbf{H}})_{n}=$ $\displaystyle~{}\sum_{i=0}^{n}\frac{\exp[{\tt ebd}_{n}(v_{n})^{\top}{\mathbf{Q }}^{\top}{\mathbf{K}}\cdot{\tt ebd}_{i}(v_{i})]{\mathbf{V}}\cdot{\tt ebd}_{i}( v_{i})}{\sum_{j=0}^{n}\exp[{\tt ebd}_{n}(v_{n})^{\top}{\mathbf{Q}}^{\top}{ \mathbf{K}}\cdot{\tt ebd}_{j}(v_{j})]}$ $\displaystyle=$ $\displaystyle~{}\frac{\exp[\lambda\bm{\eta}_{n-1}^{\top}\bm{\eta}_{0}]\cdot{ \bm{\beta}}+\sum_{i=1}^{n}{\left\{\exp[\lambda\bm{\eta}_{n-1}^{\top}\bm{\eta}_ {i}]\xi_{v_{i}}\cdot\bm{e}_{v_{i}}\right\}}}{\exp[\lambda\bm{\eta}_{n-1}^{\top }\bm{\eta}_{0}]+\sum_{j=1}^{n}\exp[\lambda\bm{\eta}_{n-1}^{\top}\bm{\eta}_{j}]}$ $\displaystyle=$ $\displaystyle~{}\frac{1}{e^{\lambda}+n}\cdot{\bm{\beta}}+\sum_{i\neq n-1}\frac {1}{e^{\lambda}+n}\cdot\xi_{v_{i}}\bm{e}_{v_{i}}+\frac{e^{\lambda}}{e^{\lambda }+n}\cdot\xi_{v_{n-1}}\bm{e}_{v_{n-1}}.$ $$

Step 2. Construction for the MLP layer. Further, define the weights for the mlp layer such that

$$ $\displaystyle~{}{\mathbf{W}}_{1}\cdot{\tt ebd}_{i}(v)=\bm{e}_{v}\in\mathbb{R}^ {V},~{}~{}~{}{\mathbf{W}}_{2}\bm{e}_{v}=\log\bm{p}_{v}\cdot 1\{v\not\in \mathcal{T}\}\in\mathbb{R}^{V}~{}~{}~{}\text{for }i\in[N],~{}~{}v\in\mathcal{V},$ (27) $$

where ${\left\{\bm{e}_{v}\right\}}$ is the eorthonormal basis in $\mathbb{R}^{V}$ and $\bm{p}_{v}\in\mathbb{R}^{V}$ is defined in Eq. (16). As a result, ${\tt mlp}({\mathbf{H}})_{n}={\mathbf{W}}_{2}\text{ReLU}({\mathbf{W}}_{1}{\tt
ebd
}_{n}(v))={\mathbf{W}}_{2}\bm{e}_{v}=\log\bm{p}_{v}\cdot\bm{1}{\left\{v\notin
\mathcal{T}\right\}}$.
This matches the Eq. (8).

Step 3. The output of the transformer. By Eq. (5) again, on non-trigger token $v\in\mathcal{V}\setminus\mathcal{T}$, the transformer output gives that

$$ $\displaystyle{\rm TF}({\mathbf{H}})_{n}$ $\displaystyle={\tt mlp}({\tt ebd}_{n}(v))+{\tt attn}({\mathbf{H}})_{n}$ $\displaystyle=\log\bm{p}_{v}+\frac{e^{\alpha_{v_{n}}}}{e^{\alpha_{v_{n}}}+n} \cdot{\bm{\beta}}+\sum_{i=0}^{n}\frac{1}{e^{\alpha_{v_{n}}}+n}\cdot\xi_{v_{i}} \bm{e}_{v_{i}}.$ $$

On trigger token $v\in\mathcal{T}$, the transformer output gives that

$$ $\displaystyle{\rm TF}({\mathbf{H}})_{n}$ $\displaystyle={\tt mlp}({\tt ebd}_{n}(v))+{\tt attn}({\mathbf{H}})_{n}$ $\displaystyle=\frac{1}{e^{\lambda}+n}\cdot{\bm{\beta}}+\sum_{i\neq n-1}\frac{1 }{e^{\lambda}+n}\cdot\xi_{v_{i}}\bm{e}_{v_{i}}+\frac{e^{\lambda}}{e^{\lambda}+ n}\cdot\xi_{v_{n-1}}\bm{e}_{v_{n-1}}.$ $$

There exists a sequence $\min_{v\in\mathcal{V}}\alpha_{v}\to\infty$, $\min_{v\in\mathcal{V}}\xi_{v}\to\infty$, $\lambda\to\infty$, and $\bm{\beta}=0$, we get that

$$ ${\sf SoftMax}[{\rm TF}({\mathbf{H}})_{n}]\to\bm{p}_{v_{n}}~{}~{}~{}\text{for } n>0,~{}~{}v_{n}\in\mathcal{V}\setminus\mathcal{T},$ $$

$$ ${\sf SoftMax}[{\rm TF}({\mathbf{H}})_{n}]\to(1\{v=v_{n-1}\})_{v\in\mathcal{V}} ~{}~{}~{}\text{for }n>0,~{}~{}v_{n}\in\mathcal{T}.$ $$

This proves Eq. (23), indicating that the transformer output matches the ground truth transition. This finishes the proof of Theorem A.0.
∎

###### Assumption .

###### Theorem A.0 (Formal statement of Theorem 1 ) .

###### Proof of Theorem A.0 .

### A.2 Proof of Theorem 2 (c): Stable phase

We first state Lemma A.0 and Proposition A.0 that are used to prove Theorem 2(c). Lemma A.0 computes the gradients of $l_{ik}$ as defined in Eq. (17).

###### Lemma A.0 .

Given $l_{ik}$ defined in Eq. (17), for any $i$, $k$, $v$, and any value of $\alpha_{v}$ and $\beta_{v}$, we have that

$$ $\displaystyle\frac{\partial l_{ik}}{\partial\alpha_{v}}=$ $\displaystyle~{}\frac{\bm{1}\{i=v\}l_{ik}e^{\alpha_{i}}}{(e^{\alpha_{i}}+M)^{2 }}\Big{[}M\beta_{k}-M_{k}\xi_{k}-\sum_{j=1}^{V}l_{ij}(M\beta_{j}-M_{j}\xi_{j}) \Big{]},$ $\displaystyle\frac{\partial l_{ik}}{\partial\beta_{v}}=$ $\displaystyle~{}\frac{e^{\alpha_{i}}}{e^{\alpha_{i}}+M}[l_{ik}\bm{1}\{k=v\}-l_ {ik}l_{iv}].$ $$

Furthermore, we have

$$ $\sum_{k=1}^{V}\frac{\partial l_{ik}}{\partial\alpha_{v}}=0~{}~{}~{}\text{for any }i,\ v,\ \bm{\alpha},\ \text{and }\bm{\beta},\quad\sum_{v=1}^{V}\frac{ \partial l_{ik}}{\partial\beta_{v}}=0~{}~{}~{}\text{for any }i,\ k,\ \bm{ \alpha},\ \text{and }\bm{\beta}.$ $$

###### Proof of Lemma A.0 .

We repeatedly use the following two facts:

$$ $\displaystyle\frac{\partial\Big{\{}\exp\Big{[}\frac{M_{k}\xi_{k}+e^{\alpha_{i} }\beta_{k}}{e^{\alpha_{i}}+M}\Big{]}\Big{\}}}{\partial\alpha_{v}}=$ $\displaystyle~{}\frac{\bm{1}\{i=v\}e^{\alpha_{i}}(M\beta_{k}-M_{k}\xi_{k})}{(e ^{\alpha_{i}}+M)^{2}}\exp\Big{[}\frac{M_{k}\xi_{k}+e^{\alpha_{i}}\beta_{k}}{e^ {\alpha_{i}}+M}\Big{]},$ $\displaystyle\frac{\partial\Big{\{}\exp\Big{[}\frac{M_{k}\xi_{k}+e^{\alpha_{i} }\beta_{k}}{e^{\alpha_{i}}+M}\Big{]}\Big{\}}}{\partial\beta_{v}}=$ $\displaystyle~{}\frac{\bm{1}\{k=v\}e^{\alpha_{i}}}{e^{\alpha_{i}}+M}\exp\Big{[ }\frac{M_{k}\xi_{k}+e^{\alpha_{i}}\beta_{k}}{e^{\alpha_{i}}+M}\Big{]}.$ $$

When $i\neq v$, $l_{ik}$ has zero gradients with respect to $\alpha_{v}$. When $i=v$, we have that

$$ $\displaystyle\frac{\partial l_{vk}}{\partial\alpha_{v}}=$ $\displaystyle~{}l_{vk}e^{\alpha_{v}}\Big{[}\frac{M\beta_{k}-M_{k}\xi_{k}}{(e^{ \alpha_{v}}+M)^{2}}\Big{]}-\frac{l_{vk}\sum_{i=1}^{V}p_{vi}e^{\alpha_{v}}\Big{ [}\frac{M\beta_{i}-M_{i}\xi_{i}}{(e^{\alpha_{v}}+M)^{2}}\Big{]}\exp\Big{[} \frac{M_{i}\xi_{i}+e^{\alpha_{v}}\beta_{i}}{e^{\alpha_{v}}+M}\Big{]}}{\sum_{i= 1}^{V}p_{vi}\exp\Big{[}\frac{M_{i}\xi_{i}+e^{\alpha_{v}}\beta_{i}}{e^{\alpha_{ v}}+M}\Big{]}}$ $\displaystyle=$ $\displaystyle~{}\frac{e^{\alpha_{v}}}{(e^{\alpha_{v}}+M)^{2}}\Big{\{}l_{vk}[M \beta_{k}-M_{k}\xi_{k}]-l_{vk}\sum_{j=1}^{V}l_{vj}(M\beta_{j}-M_{j}\xi_{j}) \Big{\}},$ $$

and

$$ $\displaystyle\frac{\partial l_{ik}}{\partial\beta_{v}}=$ $\displaystyle~{}\Big{[}\frac{e^{\alpha_{i}}}{e^{\alpha_{i}}+M}\Big{]}l_{ik}\bm {1}\{k=v\}-\frac{\Big{[}\frac{e^{\alpha_{i}}}{e^{\alpha_{i}}+M}\Big{]}p_{iv} \exp\Big{[}\frac{M_{v}\xi_{v}+e^{\alpha_{i}}\beta_{v}}{e^{\alpha_{i}}+M}\Big{] }p_{ik}\exp\Big{[}\frac{M_{k}\xi_{k}+e^{\alpha_{i}}\beta_{k}}{e^{\alpha_{i}}+M }\Big{]}}{\Big{(}\sum_{j=1}^{V}p_{ij}\exp\Big{[}\frac{M_{j}\xi_{j}+e^{\alpha_{ i}}\beta_{j}}{e^{\alpha_{i}}+M}\Big{]}\Big{)}^{2}}$ $\displaystyle=$ $\displaystyle~{}\Big{[}\frac{e^{\alpha_{i}}}{e^{\alpha_{i}}+M}\Big{]}[l_{ik} \bm{1}\{k=v\}-l_{ik}l_{iv}].$ $$

We can verify that

$$ $\displaystyle\sum_{k=1}^{V}\frac{\partial l_{ik}}{\partial\alpha_{v}}=$ $\displaystyle~{}\frac{e^{\alpha_{v}}}{(e^{\alpha_{v}}+M)^{2}}\sum_{k=1}^{V} \Big{\{}l_{vk}[M\beta_{k}-M_{k}\xi_{k}]-l_{vk}\sum_{j=1}^{V}l_{vj}(M\alpha_{j} -M_{j}\xi_{j})\Big{\}}$ $\displaystyle=$ $\displaystyle~{}\frac{e^{\alpha_{v}}}{(e^{\alpha_{v}}+M)^{2}}\Big{\{}\sum_{k=1 }^{V}l_{vk}[M\beta_{k}-M_{k}\xi_{k}]-\sum_{j=1}^{V}l_{vj}(M\alpha_{j}-M_{j}\xi _{j})\Big{\}}$ $\displaystyle=$ $\displaystyle~{}0,$ $$

and

$$ $\displaystyle\sum_{v=1}^{V}\frac{\partial l_{ik}}{\partial\beta_{v}}=$ $\displaystyle~{}\Big{[}\frac{e^{\alpha_{i}}}{e^{\alpha_{i}}+M}\Big{]}\sum_{v=1 }^{V}[l_{ik}\bm{1}\{k=v\}-l_{ik}l_{iv}]$ $\displaystyle=$ $\displaystyle~{}\Big{[}\frac{e^{\alpha_{i}}}{e^{\alpha_{i}}+M}\Big{]}[l_{ik}-l _{ik}]$ $\displaystyle=$ $\displaystyle~{}0.$ $$

This finishes the proof of Lemma A.0.
∎

Proposition A.0 computes the gradient of $\mathsf{loss}$ with respect to $\bm{\alpha}$ and $\bm{\beta}$, giving the ODE of the gradient flow.

###### Proposition A.0 .

Consider the gradient flow of optimizing $\mathsf{loss}(\bm{\alpha},\bm{\beta})$ given by

$$ $\dot{\bm{\alpha}}(t)=-\nabla_{\bm{\alpha}}\mathsf{loss}(\bm{\alpha}(t),\bm{ \beta}(t)),~{}~{}~{}\dot{\bm{\beta}}(t)=-\nabla_{\bm{\beta}}\mathsf{loss}(\bm{ \alpha}(t),\bm{\beta}(t)).$ (28) $$

Simplifying the dynamics using Lemma A.0 gives that

$$ $\displaystyle\dot{\alpha}_{v}(t)$ $\displaystyle=\frac{\tilde{\pi}_{v}e^{\alpha_{v}}}{(e^{\alpha_{v}}+M)^{2}}\sum _{i=1}^{V}({p}_{vi}-{l}_{vi})(M\beta_{i}-M_{i}\xi_{i}),$ $\displaystyle\dot{\beta}_{v}(t)$ $\displaystyle=\sum_{k=1}^{V}\Big{\{}\frac{\tilde{\pi}_{k}e^{\alpha_{k}}[p_{kv} -l_{kv}]}{e^{\alpha_{k}}+M}\Big{\}}.$ $$

###### Proof of Proposition A.0 .

Taking the derivative of $\mathsf{loss}(\bm{\alpha},\bm{\beta})$ gives that

$$ $\displaystyle\frac{\partial\mathsf{loss}(\bm{\alpha},\bm{\beta})}{\partial \alpha_{v}}=~{}$ $\displaystyle\tilde{\pi}_{v}\sum_{k=1}^{V}p_{vk}\cdot\frac{-1}{l_{vi}}\cdot \frac{\partial l_{vi}}{\partial\alpha_{v}}$ $\displaystyle=~{}$ $\displaystyle\frac{\tilde{\pi}_{v}e^{\alpha_{v}}}{(e^{\alpha_{v}}+M)^{2}}\Big{ \{}\sum_{i=1}^{V}l_{vi}[M\beta_{i}-M_{i}\xi_{i}]-\sum_{k=1}^{V}p_{vk}[M\beta_{ k}-M_{k}\xi_{k}]\Big{\}}$ $\displaystyle=~{}$ $\displaystyle\frac{\tilde{\pi}_{v}e^{\alpha_{v}}}{(e^{\alpha_{v}}+M)^{2}}\sum_ {k=1}^{V}\Big{\{}[l_{vk}-p_{vk}][M\beta_{k}-M_{k}\xi_{k}]\Big{\}}.$ $$

Similarly, we have that

$$ $\displaystyle\frac{\partial\mathsf{loss}(\bm{\alpha},\bm{\beta})}{\partial \beta_{v}}=~{}$ $\displaystyle\sum_{j=1}^{V}\tilde{\pi}_{j}\sum_{k=1}^{V}p_{jk}\Big{\{}\frac{e^ {\alpha_{j}}l_{jv}}{e^{\alpha_{j}}+M}-\frac{e^{\alpha_{j}}\bm{1}\{k=v\}}{e^{ \alpha_{j}}+M}\Big{\}}$ $\displaystyle=~{}$ $\displaystyle\sum_{j=1}^{V}\Big{\{}\frac{\tilde{\pi}_{j}e^{\alpha_{j}}[l_{jv}- p_{jv}]}{e^{\alpha_{j}}+M}\Big{\}}.$ $$

Plug them in Eq. (28) proves Proposition A.0.
∎

###### Theorem A.0 (Restatement the stable phase part in Theorem 2 (c)) .

Assume $\xi_{v}\geq 0$ for any $v$, $\pi_{v}>0$ for any $v\in\mathcal{V}$, and $\{M_{i}\cdot\xi_{i}\}_{i\in\mathcal{V}}$ are not all equal. Consider the gradient flow over the variables $(\bm{\alpha},\bm{\beta})$, i.e., $(\dot{\bm{\alpha}}(t),\dot{\bm{\beta}}(t))=-\nabla_{\bm{\alpha},\bm{\beta}}
\mathsf{loss}(\bm{\alpha}(t),\bm{\beta}(t))$. Any vector of the following form

$$ $\bm{\alpha}^{\star}=\alpha\cdot\bm{1},\quad\bm{\beta}^{\star}=c\cdot\bm{1}-e^{ -\alpha}\cdot\bm{M}\circ\bm{\xi},~{}~{}~{}\alpha,c\in\mathbb{R}$ (29) $$

is a stationary point. These are all global minimizers of $\mathsf{loss}(\bm{\alpha},\bm{\beta})$.

###### Proof of Theorem A.0 .

When $\bm{\alpha}=\bm{\alpha}^{\star}$ and $\bm{\beta}=\bm{\beta}^{\star}$, given $l_{vi}$ defined in Eq. (17) with any $v$ and $i$, we have that

$$ $\displaystyle l_{vi}=$ $\displaystyle~{}\frac{p_{vi}\exp\Big{[}\frac{M_{i}\xi_{i}+e^{\alpha}\beta_{i}} {e^{\alpha}+M}\Big{]}}{\sum_{k=1}^{V}p_{vk}\exp\Big{[}\frac{M_{k}\xi_{k}+e^{ \alpha}\beta_{k}}{e^{\alpha}+M}\Big{]}}$ $\displaystyle=$ $\displaystyle~{}\frac{p_{vi}\exp\Big{[}\frac{e^{\alpha}c}{e^{\alpha}+M}\Big{]} }{\sum_{k=1}^{V}p_{vk}\exp\Big{[}\frac{e^{\alpha}c}{e^{\alpha}+M}\Big{]}}$ $\displaystyle=$ $\displaystyle~{}p_{vi}.$ $$

Plug $l_{vi}$ into $\partial\mathsf{loss}(\bm{\alpha},\bm{\beta})/\partial\bm{\alpha}$ and $\partial\mathsf{loss}(\bm{\alpha},\bm{\beta})/\partial\bm{\beta}$, we have

$$ $\displaystyle\frac{\partial\mathsf{loss}(\bm{\alpha},\bm{\beta})}{\partial \alpha_{v}}\Big{|}_{\bm{\alpha}^{\star},\bm{\beta}^{\star}}$ $\displaystyle=\frac{\tilde{\pi}_{v}e^{\alpha_{v}}}{(e^{\alpha_{v}}+M)^{2}}\sum _{k=1}^{V}\Big{\{}(l_{vk}-p_{vk})[W\beta_{k}-M_{k}\xi_{k}]\Big{\}}=0,$ $\displaystyle\frac{\partial\mathsf{loss}(\bm{\alpha},\bm{\beta})}{\partial \beta_{v}}\Big{|}_{\bm{\alpha}^{\star},\bm{\beta}^{\star}}$ $\displaystyle=\sum_{k=1}^{V}\Big{\{}\frac{\tilde{\pi}_{k}e^{\alpha_{k}}[l_{kv} -p_{kv}]}{e^{\alpha_{k}}+M}\Big{\}}=0.$ $$

This shows that $\bm{\alpha}=\bm{\alpha}^{\star}$ and $\bm{\beta}=\bm{\beta}^{\star}$ are stationary points. We further compute the second-order derivative using Lemma A.0. To simplify the notation, we use $z_{k}=W\beta_{k}-M_{k}\xi_{k}$ and $\bm{z}=[z_{1},\ldots,z_{V}]$. We have that

$$ $\displaystyle\frac{\partial^{2}\mathsf{loss}(\bm{\alpha},\bm{\beta})}{\partial \alpha_{i}\partial\alpha_{v}}\Big{|}_{\bm{\alpha}^{\star},\bm{\beta}^{\star}}= ~{}$ $\displaystyle\bm{1}\{v=i\}\cdot\frac{\tilde{\pi}_{v}e^{\alpha}}{(e^{\alpha}+M) ^{2}}\sum_{k=1}^{V}\Big{\{}\frac{\partial l_{ik}}{\partial\alpha_{v}}z_{k}\Big {\}}$ $\displaystyle=~{}$ $\displaystyle\bm{1}\{v=i\}\cdot\frac{\tilde{\pi}_{v}e^{2\alpha}}{(e^{\alpha}+M )^{4}}\Big{\{}\sum_{k=1}^{V}l_{ik}z_{k}^{2}-\Big{[}\sum_{k=1}^{V}l_{ik}z_{k} \Big{]}^{2}\Big{\}}$ $\displaystyle=~{}$ $\displaystyle\bm{1}\{v=i\}\cdot\frac{\tilde{\pi}_{v}e^{2\alpha}}{(e^{\alpha}+M )^{4}}\Big{\{}\sum_{k=1}^{V}p_{ik}z_{k}^{2}-\Big{[}\sum_{k=1}^{V}p_{ik}z_{k} \Big{]}^{2}\Big{\}},$ $$

where in the last line, we plugged in $l_{vi}=p_{vi}$ for any $v$ and $i$. Similarly, we compute the second order derivatives with respect to $\alpha_{i}$ and $\beta_{v}$,

$$ $\displaystyle\frac{\partial^{2}\mathsf{loss}(\bm{\alpha},\bm{\beta})}{\partial \alpha_{i}\partial\beta_{v}}\Big{|}_{\bm{\alpha}^{\star},\bm{\beta}^{\star}}=~{}$ $\displaystyle\frac{\tilde{\pi}_{i}e^{\alpha}}{(e^{\alpha}+M)^{2}}\sum_{k=1}^{V }\Big{\{}\frac{\partial l_{ik}}{\partial\beta_{v}}z_{k}\Big{\}}$ $\displaystyle=~{}$ $\displaystyle\frac{\tilde{\pi}_{i}e^{2\alpha}}{(e^{\alpha}+M)^{3}}\Big{\{}p_{ iv}z_{k}-p_{iv}\sum_{k=1}^{V}p_{ik}z_{k}\Big{\}}.$ $$

With the same manner, we compute the second order derivatives with respect to $\beta_{i}$ and $\beta_{v}$,

$$ $\displaystyle\frac{\partial^{2}\mathsf{loss}(\bm{\alpha},\bm{\beta})}{\partial \beta_{i}\partial\beta_{v}}\Big{|}_{\bm{\alpha}^{\star},\bm{\beta}^{\star}}=~{}$ $\displaystyle\sum_{k=1}^{V}\Big{\{}\frac{\partial l_{ki}}{\partial\beta_{v}} \frac{\tilde{\pi}_{k}e^{\alpha}}{e^{\alpha}+M}\Big{\}}$ $\displaystyle=~{}$ $\displaystyle\frac{e^{2\alpha}}{(e^{\alpha}+M)^{2}}\sum_{k=1}^{V}\{\tilde{\pi} _{k}[\bm{1}\{v=i\}p_{kv}-p_{ki}p_{kv}]\}.$ $$

Combining the above computations gives that

$$ $\displaystyle\text{Hessian}(\mathsf{loss}(\bm{\alpha}^{\star},\bm{\beta}^{ \star}))=\left(\begin{matrix}\nabla^{2}_{\bm{\alpha}}\mathsf{loss}(\bm{\alpha} ,\bm{\beta})&\nabla_{\bm{\alpha}}\nabla_{\bm{\beta}}\mathsf{loss}(\bm{\alpha}, \bm{\beta})\\ \nabla_{\bm{\beta}}\nabla_{\bm{\alpha}}\mathsf{loss}(\bm{\alpha},\bm{\beta})& \nabla^{2}_{\bm{\alpha}}\mathsf{loss}(\bm{\alpha},\bm{\beta})\\ \end{matrix}\right),$ $$

with

$$ $\displaystyle\nabla^{2}_{\bm{\alpha}}\mathsf{loss}(\bm{\alpha},\bm{\beta})=$ $\displaystyle~{}\frac{e^{2\alpha}}{(e^{\alpha}+M)^{4}}\operatorname{diag}\Big{ \{}\tilde{\pi}\circ[{\mathbf{z}}^{\top}{\mathbf{G}}^{\mathbf{P}}_{1}{\mathbf{z }};\ldots;{\mathbf{G}}^{\mathbf{P}}_{V}{\mathbf{z}}]\Big{\}},$ $\displaystyle\nabla_{\bm{\alpha}}\nabla_{\bm{\beta}}\mathsf{loss}(\bm{\alpha}, \bm{\beta})=$ $\displaystyle~{}\frac{e^{2\alpha}}{(e^{\alpha}+M)^{3}}\operatorname{diag}\Big{ \{}\tilde{\pi}\Big{\}}[{\mathbf{z}}^{\top}{\mathbf{G}}^{\mathbf{P}}_{1};\ldots ;{\mathbf{z}}^{\top}{\mathbf{G}}^{\mathbf{P}}_{V}],$ $\displaystyle\nabla^{2}_{\bm{\beta}}\mathsf{loss}(\bm{\alpha},\bm{\beta})=$ $\displaystyle~{}\frac{e^{2\alpha}}{(e^{\alpha}+M)^{2}}\sum_{k=1}^{V}\tilde{\pi }_{k}{\mathbf{G}}^{\mathbf{P}}_{k},$ $$

where ${\mathbf{G}}^{\mathbf{P}}_{k}$ is defined in Eq. (21). Furthermore, there exists $\bm{U}$ such that $\bm{U}\text{Hessian}(\mathsf{loss}(\bm{\alpha}^{\star},\bm{\beta}^{\star}))\bm
{U}^{\top}=\text{Diag-Hessian}(\mathsf{loss}(\bm{\alpha}^{\star},\bm{\beta}^{
\star}))$, with

$$ $\text{Diag-Hessian}(\mathsf{loss}(\bm{\alpha}^{\star},\bm{\beta}^{\star}))= \left(\begin{matrix}\nabla^{2}_{\bm{\alpha}}\mathsf{loss}(\bm{\alpha},\bm{ \beta})&0\\ 0&\frac{e^{2\alpha}}{(e^{\alpha}+M)^{2}}{\mathbf{B}}\\ \end{matrix}\right),$ $$

where the ${\mathbf{B}}$ is given by

$$ $\displaystyle{\mathbf{B}}=\sum_{k=1}^{V}\tilde{\pi}_{k}\Big{(}{\mathbf{G}}^{ \mathbf{P}}_{k}-(\bm{z}^{\top}{\mathbf{G}}_{k}^{\mathbf{P}}\bm{z})^{-1}{ \mathbf{G}}^{\mathbf{P}}_{k}\bm{z}\bm{z}^{\top}{\mathbf{G}}^{\mathbf{P}}_{k} \Big{)}.$ $$

To prove that ${\mathbf{B}}$ is positive semi-definite, consider any vector $\bm{\eta}$ with $\|\bm{\eta}\|_{2}=1$:

$$ $\displaystyle\bm{\eta}^{\top}{\mathbf{B}}\bm{\eta}=$ $\displaystyle~{}\sum_{k=1}^{V}\tilde{\pi}_{k}\Big{(}\bm{\eta}^{\top}{\mathbf{G }}^{\mathbf{P}}_{k}\bm{\eta}-\frac{\bm{\eta}^{\top}{\mathbf{G}}^{\mathbf{P}}_{ k}\bm{z}\bm{z}^{\top}{\mathbf{G}}^{\mathbf{P}}_{k}\bm{\eta}}{\bm{z}^{\top}{ \mathbf{G}}_{k}^{\mathbf{P}}\bm{z}}\Big{)}.$ $$

Since ${\mathbf{G}}^{\mathbf{P}}_{k}$ is positive semi-definite, the Cauchy inequality gives that

$$ $\bm{z}^{\top}{\mathbf{G}}^{\mathbf{P}}_{k}\bm{\eta}\leq\sqrt{\bm{z}^{\top}{ \mathbf{G}}_{k}^{\mathbf{P}}\bm{z}\bm{\eta}^{\top}{\mathbf{G}}_{k}^{\mathbf{P} }\bm{\eta}}.$ $$

As a result, we have that

$$ $\displaystyle\bm{\eta}^{\top}{\mathbf{B}}\bm{\eta}\geq$ $\displaystyle~{}\sum_{k=1}^{V}\tilde{\pi}_{k}\Big{(}\bm{\eta}^{\top}{\mathbf{G }}^{\mathbf{P}}_{k}\bm{\eta}-\frac{\bm{z}^{\top}{\mathbf{G}}_{k}^{\mathbf{P}} \bm{z}\bm{\eta}^{\top}{\mathbf{G}}_{k}^{\mathbf{P}}\bm{\eta}}{\bm{z}^{\top}{ \mathbf{G}}_{k}^{\mathbf{P}}\bm{z}}\Big{)}=0.$ $$

This shows that ${\mathbf{B}}$ is positive semi-definite. Therefore, $\text{Hessian}(\mathsf{loss}(\bm{\alpha}^{\star},\bm{\beta}^{\star}))$ is positive semi-definte. This proves Theorem A.0.
∎

###### Lemma A.0 .

###### Proof of Lemma A.0 .

###### Proposition A.0 .

###### Proof of Proposition A.0 .

###### Theorem A.0 (Restatement the stable phase part in Theorem 2 (c)) .

###### Proof of Theorem A.0 .

### A.3 Proof of Theorem 2 (a): Attention sinks

###### Theorem A.0 (Restatement of the attention sink part in Theorem 2 (a)) .

Assume $\xi_{v}\geq 0$ for any $v$, $\pi_{v}>0$ for any $v\in\mathcal{V}$, and $\{M_{i}\cdot\xi_{i}\}_{i\in\mathcal{V}}$ are not all equal. Fix $\bm{\beta}=\beta\cdot\bm{1}$ for a constant $\beta$, and consider the gradient flow of the loss function $\mathsf{loss}(\bm{\alpha},\bm{\beta})$ over $\bm{\alpha}$, i.e., $\dot{\bm{\alpha}}(t)=-\nabla\mathsf{loss}(\bm{\alpha}(t),\bm{\beta})$. With any initial value $\bm{\alpha}(0)$, there exists $\bm{r}(t)$ with norm uniformly bounded in time, such that

$$ $\textstyle\bm{\alpha}(t)=\frac{1}{2}\log t\cdot\bm{1}+\bm{r}(t).$ (30) $$

###### Proof of Theorem A.0 .

We separately analyze each entry of $\bm{\alpha}$. Focusing on $\alpha_{v}$, to simplify the notation, we introduce a random variable $\varphi$ such that

$$ $\mathbb{P}(\varphi=M_{k}\xi_{k})=p_{vk}.$ $$

Denote

$$ $u=e^{\alpha_{v}}.$ $$

Therefore, using Lemma A.0, we get that

$$ $\displaystyle\frac{\mathrm{d}u}{\mathrm{d}t}=\frac{\tilde{\pi}_{v}e^{2\alpha_{ v}}}{(e^{\alpha_{v}}+M)^{2}}\sum_{i=1}^{V}({p}_{vi}-{l}_{vi})(M\beta_{i}-M_{i} \xi_{i}).$ $$

We take in $\bm{\beta}=c\cdot\bm{1}$ and expand the expression of $\mathrm{d}u/\mathrm{d}t$. This gives us that

$$ $\displaystyle\frac{\mathrm{d}u}{\mathrm{d}t}=$ $\displaystyle~{}\frac{\tilde{\pi}_{v}u^{2}}{(u+M)^{2}}\frac{\sum_{k=1}^{V}p_{ vk}e^{M_{k}\xi_{k}/(u+M)}M_{k}\xi_{k}-\sum_{k=1}^{V}p_{vk}e^{M_{k}\xi_{k}/(u+M )}\sum_{k=1}^{V}p_{vk}M_{k}\xi_{k}}{\sum_{k=1}^{V}p_{vk}e^{M_{k}\xi_{k}/(u+M)}}$ $\displaystyle=$ $\displaystyle~{}\frac{\tilde{\pi}_{v}u^{2}}{(u+M)^{2}}\frac{\mathrm{Cov}(e^{ \frac{\varphi}{u+M}},\varphi)}{\mathbb{E}e^{\frac{\varphi}{u+M}}}.$ $$

Since both $e^{x/(u+M)}$ and $x$ are monotonically increasing with respect to $x$, $\mathrm{d}u/\mathrm{d}t\geq 0$. Therefore, $u$ is monotonically increasing, and we have that

$$ $\frac{u(t)^{2}}{[u(t)+M]^{2}}\geq\frac{u(0)^{2}}{[u(0)+M]^{2}},\quad\mathbb{E} e^{\frac{\varphi}{u(t)+M}}\leq\mathbb{E}e^{\frac{\varphi}{u(0)+M}}.$ $$

Meanwhile, the first and second order Taylor expansions of $e^{\varphi/(u+M)}$ give that

$$ $e^{\frac{\varphi}{u+M}}=1+\frac{\theta_{1}(\varphi)\varphi}{u+M},\quad e^{ \frac{\varphi}{u+M}}=1+\frac{\varphi}{u+M}+\theta_{2}(\varphi)\Big{[}\frac{ \varphi}{u+M}\Big{]}^{2},$ $$

where both $\theta_{1}(\varphi)$ and $\theta_{2}(\varphi)\varphi^{2}$ are monotonically increasing functions of $\varphi$. We also have the bound that

$$ $\theta(\varphi)\leq\Big{[}\exp\Big{\{}\frac{\max_{k}M_{k}\xi_{k}}{u(0)+M}\Big{ \}}-1\Big{]}/\Big{[}\frac{\max_{k}M_{k}\xi_{k}}{u(0)+M}-1\Big{]}=C_{\theta}.$ $$

Therefore, we get two more inequalities:

$$ $\mathrm{Cov}(\theta_{1}(\varphi)\varphi,\varphi)\leq C_{\theta}\mathbb{E}( \varphi^{2}),\quad\mathrm{Cov}(\theta_{2}(\varphi)\varphi^{2},\varphi)\geq 0.$ $$

We bound $\mathrm{d}u/\mathrm{d}t$ and get that

$$ $\displaystyle\frac{\mathrm{d}u}{\mathrm{d}t}\leq$ $\displaystyle~{}\tilde{\pi}_{v}\mathrm{Cov}(e^{\frac{\varphi}{u+M}},\varphi)$ $\displaystyle=$ $\displaystyle~{}\tilde{\pi}_{v}\mathrm{Cov}(1+\frac{\theta_{1}(\varphi)\varphi }{u+M},\varphi)$ $\displaystyle\leq$ $\displaystyle~{}\frac{\tilde{\pi}_{v}C_{\theta}\mathbb{E}(\varphi^{2})}{u}.$ $$

By solving the ODE, we get that

$$ $\displaystyle u\leq\sqrt{2\tilde{\pi}_{v}C_{\theta}\mathbb{E}(\varphi^{2})t+C_ {1}}.$ $$

To give a lower bound, we have that

$$ $\displaystyle\frac{\mathrm{d}u}{\mathrm{d}t}\geq$ $\displaystyle~{}\frac{u(0)^{2}}{[u(0)+M]^{2}}\frac{\tilde{\pi}_{v}\mathrm{Cov} (e^{\frac{\varphi}{u+M}},\varphi)}{\mathbb{E}e^{\frac{\varphi}{u(0)+M}}}$ $\displaystyle=$ $\displaystyle~{}\frac{u(0)^{2}}{[u(0)+M]^{2}}\frac{\tilde{\pi}_{v}}{\mathbb{E} e^{\frac{\varphi}{u(0)+M}}}\mathrm{Cov}(1+\frac{\varphi}{u+M}+\theta_{2}( \varphi)\Big{[}\frac{\varphi}{u+M}\Big{]}^{2},\varphi)$ $\displaystyle\geq$ $\displaystyle~{}\frac{u(0)^{2}}{[u(0)+M]^{2}}\frac{\tilde{\pi}_{v}}{\mathbb{E} e^{\frac{\varphi}{u(0)+M}}}\frac{\operatorname{Var}(\varphi)}{u+M}$ $\displaystyle\geq$ $\displaystyle~{}\frac{u(0)^{2}}{[u(0)+M]^{2}}\frac{\tilde{\pi}_{v}}{\mathbb{E} e^{\frac{\varphi}{u(0)+M}}}\cdot\frac{u(0)}{u(0)+M}\cdot\frac{\operatorname{ Var}(\varphi)}{u}$ $\displaystyle=$ $\displaystyle~{}\frac{\tilde{C}}{u}.$ $$

Therefore, $u\geq\sqrt{\tilde{C}t+\tilde{C}_{2}}$. In conclusion, we have that

$$ $y_{v}=\log u=\frac{1}{2}\log t+r_{v},$ $$

with $r_{v}$ bounded. This proves Theorem A.0.
∎

###### Theorem A.0 (Restatement of the attention sink part in Theorem 2 (a)) .

###### Proof of Theorem A.0 .

### A.4 Proof of Theorem 2 (b): Value-state drains

###### Theorem A.0 (Restatement of Theorem 2 (b)) .

Assume $\xi_{v}\geq 0$ for any $v$, $\pi_{v}>0$ for any $v\in\mathcal{V}$, and $\{M_{i}\cdot\xi_{i}\}_{i\in\mathcal{V}}$ are not all equal. Fix $\bm{\alpha}=\alpha\cdot\bm{1}$ for a constant $\alpha$, define $\overline{\beta}(0)=V^{-1}[\sum_{v}\beta_{v}(0)]$ and $\overline{B}=V^{-1}[\sum_{v}M_{v}\xi_{v}]$. Consider the gradient flow of the loss function $\mathsf{loss}(\bm{\alpha},\bm{\beta})$ over $\bm{\beta}$ for fixed $\bm{\alpha}$, i.e., $\dot{\bm{\beta}}(t)=-\nabla_{\bm{\beta}}\mathsf{loss}(\bm{\alpha},\bm{\beta}(t))$. As $t\to\infty$, we have

$$ $\bm{\beta}(t)\to\bm{\beta}^{\star}=[\overline{\beta}(0)+e^{-\alpha}\overline{B }]\cdot\bm{1}-e^{-\alpha}\cdot\bm{M}\circ\bm{\xi}.$ (31) $$

###### Proof of Theorem A.0 .

We plug $\bm{\beta}^{\star}$ into the $\mathsf{loss}$ and get that $\mathsf{loss}(\bm{\alpha},\bm{\beta}^{\star})=\sum_{v=1}^{V}\tilde{\pi}_{v}
\sum_{k=1}^{V}p_{vk}\log p_{vk}$. Computing $\nabla^{2}_{\bm{\beta}}\mathsf{loss}(\bm{\alpha},\bm{\beta})$, we get that

$$ $\nabla^{2}_{\bm{\beta}}\mathsf{loss}(\bm{\alpha},\bm{\beta})=\sum_{k=1}^{V} \tilde{\pi}_{k}{\mathbf{G}}^{\mathbf{L}}_{k},$ $$

where ${\mathbf{G}}^{\mathbf{L}}_{k}$ is defined in Eq. (21).
Lemma A.0 indicates that it is positive semi-definite. Therefore, we have that

$$ $\mathsf{loss}(\bm{\alpha},\bm{\beta}(t))\to\mathsf{loss}(\bm{\alpha},\bm{\beta }^{\star})~{}~{}\text{as }t\to\infty.$ $$

We choose $\delta$ as defined in Eq.(22). When $t$ is sufficiently large, we have that

$$ $\mathsf{loss}(\bm{\alpha},\bm{\beta}(t))\leq\mathsf{loss}(\bm{\alpha},\bm{ \beta}^{\star})+\frac{1}{\min_{k\in\mathcal{V}\setminus\mathcal{T}}\tilde{\pi} _{k}}\cdot 2\delta^{2}.$ $$

The convexity further implies that for any $\tilde{\bm{\beta}}=\theta\bm{\beta}(t)+(1-\theta)\bm{\beta}^{\star}$ ($\theta\in(0,1)$), we have that

$$ $\mathsf{loss}(\bm{\alpha},\tilde{\bm{\beta}})\leq\mathsf{loss}(\bm{\alpha},\bm {\beta}^{\star})+\frac{1}{\min_{k\in\mathcal{V}\setminus\mathcal{T}}\tilde{\pi }_{k}}\cdot 2\delta^{2}.$ $$

Denote $\tilde{\bm{l}}_{v}={\bm{l}}_{v}(\bm{\alpha},\tilde{\bm{\beta}})$ as $\bm{l}$ evaluated on $(\bm{\alpha},\tilde{\bm{\beta}})$. Using the definition of the KL-divergence in Eq. (20), we have that

$$ $\sum_{v=1}^{V}\tilde{\pi}_{v}KL(\bm{p}_{v}~{}||~{}\tilde{\bm{l}}_{v})=\mathsf{ loss}(\bm{\alpha},\bm{\beta}(t))-\mathsf{loss}(\bm{\alpha},\bm{\beta}^{\star}) \leq\frac{1}{\min_{k\in\mathcal{V}\setminus\mathcal{T}}\tilde{\pi}_{k}}\cdot 2 \delta^{2}.$ $$

This further implies that $KL(\bm{p}_{v}~{}||~{}\tilde{\bm{l}}_{v})\leq 2\delta^{2}$ for any $v$. Using Pinsker’s inequality, we get that

$$ $\sum_{k=1}^{V}|p_{vk}-\tilde{l}_{vk}|=\|\bm{p}_{v}-\bm{l}_{v}\|_{\text{TV}} \leq\sqrt{KL(\bm{p}_{v}~{}||~{}\tilde{\bm{l}}_{v})/2}\leq\delta.$ $$

Therefore, $\max_{v,k}\left|p_{vk}-\tilde{l}_{vk}\right|\leq\delta$.
Lemma A.0 gives that $\sum_{v=1}^{V}\dot{\beta}_{v}(t)=0$. Therefore, $\sum_{v=1}^{V}\beta_{v}(t)/V=\overline{\beta}(0)$. The choice of $\bm{\beta}^{\star}$ guarantees that $\overline{\beta}^{\star}=\overline{\beta}(0)$. This shows that $\bm{\beta}(t)-\bm{\beta}^{\star}\perp\bm{1}$. Using Lemma A.0, there exists $\omega>0$ such that

$$ $(\bm{\beta}(t)-\bm{\beta}^{\star})^{\top}\nabla^{2}_{\bm{\beta}}\mathsf{loss}( \bm{\alpha},\bm{\beta})(\bm{\beta}(t)-\bm{\beta}^{\star})=(\bm{\beta}(t)-\bm{ \beta}^{\star})^{\top}\Big{[}\sum_{k=1}^{V}\tilde{\pi}_{k}{\mathbf{G}}^{ \mathbf{L}}_{k}\Big{]}(\bm{\beta}(t)-\bm{\beta}^{\star})\geq\frac{\omega}{2}\| \bm{\beta}(t)-\bm{\beta}^{\star}\|_{2}^{2}.$ $$

Using Taylor expansion, we have that

$$ $\displaystyle\mathsf{loss}(\bm{\alpha},\bm{\beta}^{\star})-\mathsf{loss}(\bm{ \alpha},\bm{\beta}(t))=$ $\displaystyle~{}-\nabla_{\beta}\mathsf{loss}(\bm{\alpha},\bm{\beta}(t))(\bm{ \beta}(t)-\bm{\beta}^{\star})+\frac{1}{2}(\bm{\beta}(t)-\bm{\beta}^{\star})^{ \top}\nabla^{2}_{\bm{\beta}}\mathsf{loss}(\bm{\alpha},\tilde{\bm{\beta}})(\bm{ \beta}(t)-\bm{\beta}^{\star})$ $\displaystyle\geq$ $\displaystyle~{}-\nabla_{\beta}\mathsf{loss}(\bm{\alpha},\bm{\beta}(t))(\bm{ \beta}(t)-\bm{\beta}^{\star})+\frac{\omega}{2}\|\bm{\beta}(t)-\bm{\beta}^{ \star}\|_{2}^{2}$ $\displaystyle\geq$ $\displaystyle~{}-\frac{1}{2\omega}\|\nabla_{\beta}\mathsf{loss}(\bm{\alpha}, \bm{\beta}(t))\|_{2}^{2}.$ $$

This shows that $\mathsf{loss}(\bm{\alpha},\bm{\beta}(t))$ satisfies the Polyak-Lojasiewicz (PL) condition (Karimi et al., ) when $t$ is sufficiently large.
This proves Theorem A.0.
∎

###### Theorem A.0 (Restatement of Theorem 2 (b)) .

###### Proof of Theorem A.0 .

## Appendix B The Linear Growth of the Residual States

### B.1 The minimal model structure to recapitulate residual state peak

We give more details for the claim in Section 2.3, stating that “The residual-state peaks require a three-layer structure.”
Figure 11 presents the difference of residual norms between the $\langle\texttt{s}\rangle$ token and others ($\|{\tt Res}_{\texttt{$\langle\texttt{s}\rangle$}}\|-\mathbb{E}_{v\neq\texttt{$
\langle\texttt{s}\rangle$}}[\|{\tt Res}_{v}\|]$), with different combinations of model structures. The $3\times{\rm TF}$ and $2\times{\rm TF}+{\tt mlp}$ are the architectures that demonstrate clear evidence of residual state peaks.

Figure: Figure 11: Minimal structures to elicit residual state peaks. We use $A+B+C$ to indicate the model with structure $A$, $B$, $C$ in layers 0, 1, and 2, respectively.
Refer to caption: x21.png

### B.2 Additional plots for the three-layer transformer trained on BB task

We provide more results to the three layer transformer model trained on the BB task. They provide supporting evidence for the claim in Section 2.3, stating that “Massive residual states amplify attention sinks and value-state drains in later layers.”
Figures 12, 13, and 14 show the extreme token phenomena in a three-layer transformer. The residual state peaks show different phenomena from those in LLMs, with the last layer output increasing the residual norms of non-$\langle\texttt{s}\rangle$ tokens. Figure 1
demonstrates that the residual state norms of $\langle\texttt{s}\rangle$ drop match the magnitudes of other tokens at the last layer.

Figure: (a) Layer 0
Refer to caption: x22.png

Figure: (a) Layer 0
Refer to caption: x25.png

Figure: (a) Layer 0
Refer to caption: x28.png

### B.3 Potential mechanism for linear growth of the residual state peak in multi-layer models

We give more details for the claim in Section 2.3, stating that “The ReLU attention and changing Adam to SGD eliminates the residual state peaks” We first state Claim 2.3.

###### Claim B.0 (Potential mechanism for the formation of residual-state peaks) .

In the training dynamic of a multi-layer transformer, if the mutual reinforcement mechanism (cf. Claim 2) occurs in upper layers:

- 1.
The gradients of ${\tt Res}_{\texttt{$\langle\texttt{s}\rangle$}}$ have the same direction (aligning with the null space of value matrices in upper layers and the ${\tt Key}_{\texttt{$\langle\texttt{s}\rangle$}}$) along the training dynamics.
- 2.
The layer-norm operations cause the fast decay of the magnitude of the gradients.
- 3.
Adam induces diminishing gradients to be constant updates, leading to the linear growth for the norm of the residual state of the extreme token.

To support the claim, we use the simplified model in Section 2, including the residual state norm.
Denote the layer-norm operation as LayerNorm.
Heuristically, we can split the residual state ${\tt Res}_{\texttt{$\langle\texttt{s}\rangle$}}$ to a summation of two directions.

$$ ${\tt Res}_{\texttt{$\langle\texttt{s}\rangle$}}=m\cdot\bm{\eta}+\bm{ \varepsilon},$ $$

where $\bm{\eta},\bm{\varepsilon}\in\mathbb{R}^{V}$ with $\|\bm{\eta}\|_{2}=\|\bm{\varepsilon}\|_{2}=1$, and $\bm{\eta}^{\top}\bm{\varepsilon}=\rho>0$. The $\bm{\eta}$ corresponds to the direction of ${\tt Key}_{\texttt{$\langle\texttt{s}\rangle$}}$ in the original transformer, and $\bm{\varepsilon}$ corresponds to other directions.
Assume that the attention logit from the token $v$ to the $\langle\texttt{s}\rangle$ token in layer 1 is given by

$$ $\text{logit}_{v,\texttt{$\langle\texttt{s}\rangle$}}=\alpha_{v}=\tilde{\alpha} _{v}\bm{\eta}^{\top}\text{LayerNorm}({\tt Res}_{\texttt{$\langle\texttt{s} \rangle$}})=\tilde{\alpha}_{v}\cdot\frac{m+\rho}{\sqrt{m^{2}+2m\rho+1}}.$ (32) $$

We assume that the scalars $m$ and $\tilde{\alpha}$ are trainable, quantifying the norm of the residual states and magnitude of attention sinks. In the loss function $\mathsf{loss}_{v}$ as defined in Eq. (10), we replace $\alpha_{v}$ by the expression as in Eq. (32), so that the loss function becomes a function of $(\tilde{\alpha}_{v},\bm{\beta},m)$, denoted as

$$ $\widetilde{\mathsf{loss}}_{v}(\tilde{\alpha}_{v},\bm{\beta},m)=\mathsf{loss}_{ v}(\alpha_{v},\bm{\beta}),$ $$

We then consider the total loss as the average of the losses on each non-trigger token, weighted by its proportion in the stable distribution $\{\pi_{v}\}_{v\in\mathcal{V}}$, given by

$$ $\widetilde{\mathsf{loss}}(\tilde{\bm{\alpha}},\bm{\beta},m)=\sum_{v\in\mathcal {V}\setminus\mathcal{T}}\pi_{v}\cdot\widetilde{\mathsf{loss}}_{v}(\tilde{ \alpha}_{v},\bm{\beta},m).$ (33) $$

###### Proposition B.0 .

Assume $\xi_{v}\geq 0$ for any $v$, ${\left\{W_{k}\beta_{k}\right\}}_{k\in\mathcal{V}}$ are not all equal, and $\rho>0$. Fix $\bm{\beta}=\bm{0}$, and consider the gradient flow of $\widetilde{\mathsf{loss}}(\tilde{\bm{\alpha}},\bm{\beta},m)$ over $\tilde{\bm{\alpha}}$ and $m$. With any initial value $\tilde{\alpha}_{v}(0)>0$ for any $v$ and $m(0)>0$, we have that

$$ $\dot{m}(t)=O\Big{(}\frac{\log t}{\sqrt{t}m^{3}}\Big{)}.$ $$

###### Proof of Proposition B.0 .

The chain rule gives that

$$ $\dot{\tilde{\alpha}}_{v}(t)=\dot{\alpha}_{v}\cdot\frac{m+\rho}{\sqrt{m^{2}+2m \rho+1}},$ $$

and

$$ $\dot{m}(t)=\sum_{v=1}^{V}\Big{\{}\dot{\alpha}_{v}\tilde{\alpha}_{v}\cdot\frac{ \mathrm{d}\text{LayerNorm}({\tt Res}_{\texttt{$\langle\texttt{s}\rangle$}})}{ \mathrm{d}t}\Big{\}}.$ $$

With the initial values, $\dot{m}(t)\geq 0$ and $\dot{\tilde{\alpha}}_{v}(t)\geq 0$.
We have $m(t)\geq 0$ for any $t$. Hence,

$$ $\dot{\tilde{\alpha}}_{v}\in[\rho\dot{\alpha}_{v},\dot{\alpha}_{v}].$ $$

Therefore, $\tilde{\bm{\alpha}}=2^{-1}\log t\bm{1}+\tilde{\bm{r}}(t)$ with $\tilde{\bm{r}}(t)$ uniformly bounded over time. Furthermore, we have that

$$ $\displaystyle\dot{m}(t)$ $\displaystyle=\sum_{v=1}^{V}\Big{\{}\dot{\alpha}_{v}\tilde{\alpha}_{v}\cdot \frac{\mathrm{d}\text{LayerNorm}({\tt Res}_{\texttt{$\langle\texttt{s}\rangle$ }})}{\mathrm{d}t}\Big{\}}$ $\displaystyle=O\Big{(}\frac{\log t}{\sqrt{t}}\Big{)}\cdot\frac{1-\rho^{2}}{(m^ {2}+2m\rho+1)^{3/2}}$ $\displaystyle=O\Big{(}\frac{\log t}{\sqrt{t}m^{3}}\Big{)}.$ $$

This proves Proposition B.0.
∎

We use simulation to demonstrate the effect of Adam. We train the scalar $m$ using Adam with gradient $\mathrm{d}m=\log t/[\sqrt{t}m^{3}]$. We set $\beta_{1}=0.9$, $\beta_{2}=0.999$, weight decay$=10^{-8}$, and the learning rate $\text{lr}=0.3$. Figure 15 presents the training dynamics of $m$. We observe the linear growth after a warming-up phase. In contrast, when trained by SGD with learning rate $\text{lr}=0.3$, $m$ remains small. The results match transformer models on BB-task as in Figure 7.

Figure: Figure 15: With the gradient formula in Proposition B.0, Adam causes linear growth of $m$.
Refer to caption: x31.png

###### Claim B.0 (Potential mechanism for the formation of residual-state peaks) .

###### Proposition B.0 .

###### Proof of Proposition B.0 .

## Appendix C Ablations

### C.1 Experimental details

We provide more details for experiments in Section 2. We train transformers with positional embedding, pre-layer norm, ${\sf SoftMax}$ activation in attn, and ReLU activation in mlp.
We use Adam with constant learning rate $0.0003$, $\beta_{1}=0.9$, $\beta_{2}=0.99$, $\varepsilon=10^{-8}$, and a weight decay of $0.01$. We choose a learning rate of $0.03$ for the SGD. In each training step, we resample from the BB task with a batch size of $B=512$ and sequence length $N=256$. Unless otherwise specified, the model is trained for $10,000$ steps. Results are consistent across different random seeds.

### C.2 Additional attention plots of a 1-layer transformer trained on the BB task

We provide more attention plots of the 1-layer transformer on sequences other than those shown in Figure 2.
Figure 16 presents more attention-weight heat maps of the one-layer transformer model trained on the BB task. All attention maps show the attention sink phenomenon. Some non-trigger tokens present attention patterns other than attention sink. For example, trigger tokens serve as attention sinks in some inputs in Figure 16.

Figure: (a) Sequence 0
Refer to caption: x32.png

### C.3 Statics and dynamics of the simplified model in Theorem 2

We provide simulations that justify our model simplifications in Section 2.
We pretrrain the simplified model structure in Figure 5 with several modifications: (1) we use a trainable mlp-layer with random Gaussian initialization; (2) we take ${\tt Val}_{\texttt{$\langle\texttt{s}\rangle$}}={\mathbf{O}}\bm{\beta}$, with ${\mathbf{O}}\in\mathbb{R}^{V\times V}$ and $\bm{\beta}\in\mathbb{R}^{V}$. Both ${\mathbf{O}}$ and $\bm{\beta}$ are trainable. Empirically, with a trainable mlp layer but without the trainable matrix ${\mathbf{O}}$, ${\tt Val}_{\texttt{$\langle\texttt{s}\rangle$}}$ becomes a non-negligible bias term instead of converging to zero. Collectively, we update parameters ${\tt mlp}$, ${\mathbf{O}}$, $\bm{\alpha}$, $\bm{\beta}$, $\lambda$, and $\bm{\xi}$ using Adam with a learning rate of $0.03$.
Figure 17 and 18 present statics and dynamics that match the observations in the one-layer transformer.

Figure: (a) Attention weights
Refer to caption: x35.png

Figure: Figure 18: The dynamics of the simplified model structure trained on the BB task. The horizontal axis is logarithmatically scaled after steps $400$. The excess risk curves match the one-layer transformer. The logit curve is close to the logarithmic growth predicted in Theorem 2.
Refer to caption: x37.png

### C.4 The Bigram-Backcopy task without the ⟨ s ⟩ delimited-⟨⟩ s \langle\texttt{s}\rangle ⟨ s ⟩ token.

We train a one-layer transformer on the BB task without the $\langle\texttt{s}\rangle$ token. Figure 19 shows that the zeroth token is not a sink token. Instead, trigger tokens and delimiter tokens seem to become sink tokens. In particular, the observation that delimiter tokens become extreme matches the observation in LLMs that delimiter tokens may also become extreme tokens (cf. Section G.2).

Figure: (a) Attention weights
Refer to caption: x38.png

## Appendix D More Attention Heads in Dormant and Active Phase

We demonstrate a head with clear active-dormant mechanism in Figure 8.
In this section, we present two more active-dormant heads in Llama 2-7B-Base, in Figures 20 and 21, which are more difficult to interpret than Layer 16 Head 25, but become dormant on some inputs and remain active on others.

Figure: (a) Attention patterns
Refer to caption: extracted/5985254/Figures/L16H20/attn_github_head20.png

Figure: (a) Attention patterns
Refer to caption: extracted/5985254/Figures/L16H28/attn_github_head28.png

## Appendix E Fine-Grained Static Mechanisms for Extreme-Token Phenomena

In this section, we identify more fine-grained static mechanisms for extreme-token phenomena in Llama 3.1-8B-Base. To do this, we identify circuits for the origin of attention sinks and small value states. Then, using ablation studies, we study the origin of massive norms. Again, we use the generic test phrase “$\langle\texttt{s}\rangle$ Summer is warm. Winter is cold.”

Figure: Figure 22: A visualization of attention heads at Layer 0 of Llama 3.1-8B-Base. Notice that many heads have the attention sink property, even at Layer 0 without any cross-token interaction. As usual, the test phrase is “Summer is warm. Winter is cold.” The most clear attention sink is Head 31.
Refer to caption: x42.png

Figure: (a) Correlations between key and query states.
Refer to caption: x43.png

#### Attention sinks and global contextual semantics.

There are many attention heads that exhibit attention sinks at layer 0 0, and the $\langle\texttt{s}\rangle$ token is always the sink token (see Figure 22). From now on until the end of this section, we restrict our attention to Head 31 of Layer 0, which is an attention sink. These attention sinks are caused by two linear-algebraic factors, demonstrated in Figure 23.

- 1.
The key state of the $\langle\texttt{s}\rangle$ token has small dot product with all other key states.
- 2.
The query states of all tokens are nearly orthogonal to the key states of all tokens except the $\langle\texttt{s}\rangle$ token.

These two facts combine to ensure that the key state of the $\langle\texttt{s}\rangle$ token is picked out by each query state, causing the attention sink. Since these query and key states are produced without any cross-token interaction, the alignment of different states is caused purely by the token’s global importance or meaning imparted via pretraining. The $\langle\texttt{s}\rangle$ token has no semantic meaning in the context of prose tokens, so its key state is not aligned with key states of meaningful prose tokens. Also, delimiter tokens, often considered secondary attention sinks (cf. Section G.2), have the most aligned key states to the key state of the $\langle\texttt{s}\rangle$ token, and are also the tokens with the least semantic meaning in the prose context. Thus, we identify that, at least in this restricted example, query state and key state alignment depends heavily on the contextual semantics of the token.

Figure: (a) Value-state drains at Layer 0 Head 31 of Llama 3.1-8B-Base.
Refer to caption: x45.png

#### Value-state drains.

The value states of the $\langle\texttt{s}\rangle$ token at Layer 0 0 Head 31 are already near zero, as demonstrated in Figure 24(a). While the delimiter tokens, which are less semantically meaningful in the prose context, have smaller value states than the rest, they are not as small as the value state of the $\langle\texttt{s}\rangle$ token which is guaranteed to not have any semantics.

#### Residual state peaks.

Residual state peaks are caused by the first two layers’ MLPs. In particular, we perform several ablations, comparing between the residual state norms in a later layer ($24$) of an un-edited forward pass versus forward passes where we force the output of either multiple layers, a single layer, an attention block, or an MLP to be zero (and hence remove its contribution from the residual stream). As shown in Figure 24(b), ablating either Layer 0’s or Layer 1’s MLP is sufficient to remove the residual state peak. In particular, the second-largest token at Layer 24 in each ablation (including the original setup) has norm between $29$ and $38$, so the interventions ensure that all tokens have similar size.

## Appendix F Extreme-Token Phenomena Over Many Samples

In this section we show that the extreme-token phenomena, and our predictions from the BB model, exhibit in prompts other than “Summer is warm. Winter is cold.” To this end, we use 128 samples from the Wikipedia dataset, each truncated to 8 tokens. Figure 25 provides aggregate statistics of extreme-token phenomena in Llama 3.1-8B, which are similar to the fine-grained statistics over a single prompt from Figure 1. Figure 26 provides aggregate statistics of the development of extreme-token phenomena over the training dynamics of OLMo, which are similar to the fine-grained statistics over a single prompt from Figure 9 and Figure 10.

Figure: (a) Attention weights (L24).
Refer to caption: x47.png

Figure: (a) Attention weights (L24).
Refer to caption: x50.png

## Appendix G Assorted Caveats

### G.1 Multiple attention sinks vs. one attention sink

As we have seen, attention heads in the BB task (Section 2), Llama 2-7B-Base (Section 3.1), and OLMo (Section 3.2) exhibit multiple sink tokens. That is, when heads in these models are dormant, they tend to have two sink tokens. For the LLMs in this group, at least on prose data, the $\langle\texttt{s}\rangle$ token as well as the first delimiter token (e.g., representing . or ;) are sink tokens. Meanwhile, Llama-3.1-8B-Base (Section 3) only ever has one attention sink on prose data, and the $\langle\texttt{s}\rangle$ token is always the sink token. Here, we offer a possible explanation of this phenomenon. For the BB task, multiple sink tokens are necessary to solve the task. For LLMs, we believe this distinction may be explained by the relative proportion of coding data, in which delimiters have a greater semantic meaning than prose, within the training set. For instance, OLMo was trained on DOLMA (Soldaini et al., ), which has around 411B coding tokens. Meanwhile, Llama 2 used at most (2T $\times$ 0.08 =) 0.16T coding tokens. Finally, Llama 3.1 used around (15.6T $\times$ 0.17 =) 2.6T coding tokens (Dubey et al., ). On top of the raw count being larger, coding tokens are a larger proportion of the whole pretrtraining dataset for Llama 3.1 compared to other model families. Thus, during training, the presence of delimiters would not be considered unhelpful towards next-token prediction, since such delimiters carry plenty of semantics in a wide variety of cases. Our earlier hypothesis in Section 3.1 proposes that only tokens which lack semantics in almost all cases are made to be sink tokens. This could be a reason for the distinction.

### G.2 The role of a fixed ⟨ s ⟩ delimited-⟨⟩ s \langle\texttt{s}\rangle ⟨ s ⟩ token in the Active-Dormant mechanism

Some models, such as OLMo, are not trained with a $\langle\texttt{s}\rangle$ token. Despite this, the first token of the input still frequently develops into a sink token. We can study the effect of positional encoding of the tokens on the attention sink phenomenon by shuffling the tokens before inputting them into the transformer, and observing how and why attention sinks form. If we do this with the phrase “Summer is warm. Winter is cold.” with OLMo, we observe that at Layer 24, there are many attention sink heads where the first token and first delimiter token share attention mass, even if the sentence is jumbled up and makes no grammatical sense. This points towards the observation that without a $\langle\texttt{s}\rangle$ token, the attention sink formation uses both positional data and, to a greater degree, the semantic data of each token. We leave studying this effect in greater detail to future work.

Figure: Figure 27: Attention sinks with shuffled input in Layer 24 of OLMo. In order to understand the impact of positional encodings when there is no $\langle\texttt{s}\rangle$ token, we shuffle the input of the test string “Summer is warm. Winter is cold.” in OLMo. We observe that there is still an attention sink on token 0 0, despite it being a random token that does not usually start sentences or phrases (since it is uncapitalized). This shows that the positional embedding, say via RoPE, has a large impact on the formation of attention sinks — when the semantics of each token have switched positions, the attention sink still forms on the zeroth token.
Refer to caption: extracted/5985254/Figures/bos_study/attn_heads_shuffled.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: -	In section 2.3, it is mentioned that no residual peaks appear in a one-layer transformer(BB model) or even with different combinations (2xTF, MLP+TF+ML, Attn+TF+MLP); it would be interesting to know why it seems to have residual peaks in certain instances. The authors claim that with a BB model, “residual state peaks contribute to the attention sink and value state drain phenomena.” However, it still remains unclear between the dependency between residual state peaks and, attention sink and value state drain. It would be helpful to see the comparison of residual state peaks and attention sink and value state drain phenomena for the 1-layer BB model even when residual state peaks do not appear.
soundness: 3
strengths: -	The paper provides valuable insights into LLMs' underlying mechanisms of extreme-token phenomena (attention sinks, value-state drains, and residual-state peaks). By identifying the active-dormant mechanism in attention heads and a mutual reinforcement process, the paper contributes to a deeper understanding of transformer model behavior, enhancing interpretability in LLMs.
-	The paper provides a comprehensive theoretical and empirical analysis of the mutual reinforcement mechanism in a toy one-layer transformer model to explain extreme-token phenomena. The identification of mutual reinforcement dynamics is robustly supported by theoretical analysis and validated through experiments on pre-trained LLMs using the Bigram-Backcopy task, offering strong evidence that this mechanism underlies attention sinks, value-state drains, and residual state peaks in both simplified and complex models.
-	The paper’s application of findings from toy models to analyze real-world LLMs, such as Llama2 and OLMo, significantly strengthens the study's impact. Validating the mutual reinforcement mechanism and active-dormant attention head dynamics in large-scale models highlights these identified mechanisms' broader applicability and relevance. This cross-model validation not only supports the robustness of the findings but also enhances the study's generalizability and potential utility for advancing transformer-based architectures
summary: The paper investigates the underlying mechanisms of "extreme-token phenomena" in large language models (LLMs), characterized by three main effects: attention sinks, value-state drains, and residual-state peaks. These effects result in specific tokens persistently attracting high attention weights (attention sinks), displaying reduced value states (value-state drains), and exhibiting notably large norms in intermediate representations (residual-state peaks). The authors explore this by analyzing transformer models in a controlled setting, the Bigram-Backcopy (BB) task, identifying an active-dormant mechanism where attention heads behave differently based on input tokens. They find that mutual reinforcement between attention sinks and value-state drains drives these phenomena, with residual state peaks emerging in deeper architectures. They also find that attention sinks and value state draining are related, and interaction between the mutual reinforcement mechanism and Adam is responsible for residual state peaks. Expanding from a toy model's theoretical and empirical analysis, the paper shows that an active-dormant mechanism controls attention heads. Specifically, it highlights the active-dormant mechanism in existing pre-trained LLMs like Llama2 and OLMo-7B-0704 by identifying an interpretable active-dormant head through causal intervention analyses. By intervening in attention mechanisms and optimization strategies, such as replacing SoftMax with ReLU activations, they demonstrate methods to mitigate these extreme-token phenomena.
weaknesses: -	The motivation behind the mitigation of extreme-token phenomena is not clear. It is not clearly established that extreme token adversely affects the model’s performance. Although the paper explores interventions like replacing SoftMax with ReLU, these modifications are limited and could unintentionally impact model performance. The study doesn’t extensively examine trade-offs associated with these changes, such as how they might affect other aspects of model performance or generalization ability. This leaves questions about whether mitigating these phenomena would lead to meaningful improvements in model usability or performance in practical applications. Could the authors provide a more comprehensive explanation or analysis of why or what specific effects reduction of the phenomena on the downstream task?
-	The paper claims that the formation of extreme tokens in LLMs follows a similar mutual reinforcement mechanism as observed in a simplified BB model, with theoretical results that can be generalized to more complex models. However, this justification relies on specific assumptions, including zeroing out attention heads for intervention outcomes, which will not impact token prediction accuracy. This assumption may oversimplify the interactions in real-world LLMs, where dependencies between heads and layers are often more intricate. Additionally, while some empirical evidence in pre-trained models aligns with the theory (active and dormant attention heads within Llama2), the scope of models evaluated is limited, raising questions about the robustness of the claim across diverse architectures and tasks. Would it be possible for authors to show whether their claims about specific attention heads being active and dormant but under the premise that context are not significantly different? For example, mathematics-based context vs language comprehension-based context. It would help gauge the sensitivity of the claims when the contexts are not vastly different. In particular, would we still find attention heads with simple and interpretable active behavior for weaker contextual differences?

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: See weaknesses part.
soundness: 2
strengths: - The paper is very well written and easy to follow.

- Extreme-Token Phenomena are very interesting and quite common across LLMs. The paper provides a simplified setting to theoretically describe their emergence. In particular, the authors show that these phenomena manifest themselves in synergy through training, via "mutual reinforcement dynamics".

- The paper shows that the analysis developed on the BB task extends to some extent to LLMs. In particular, by manually deleting a specific attention head on the LLAMA model, the authors show that the resulting model's performance depends on whether the head is dormant or not, which in turn depends on the specific input context (Fig. 6). In addition, Fig. 7 and 8 corroborate the theoretical results on the simple BB model.
summary: The paper studies the emergence of extreme-token phenomena in modern transformer neural networks, namely: Attention sinks, Value state drains and Residual state peaks, all previously identified in recent works. The first refers to the observation that in some attention heads, special tokens named sink tokens attract a large proportion of the attention weights. The second highlights that values states associated with sink tokens are usually significantly smaller than those of other tokens. The third refers to the observation that the intermediate representations (except for initial and final layers) of sink tokens are characterised by larger norms compared to those of other tokens. 


The authors first show that attention sinks and value state drains can be identified in a simple settings: one layer transformer models trained on the Bigram-Backcopy task which involves sequences generated either through bigram transitions or by copying previous tokens in the presence of special trigger tokens. Through an in-depth empirical analysis, the authors show that all non-trigger tokens exhibit attention sinks, while the attention for trigger tokens is concentrated on their preceding positions. In particular, the BOS tokens attracts a significant proportion of the attention weights for the non-trigger tokens, and its value state decreases in norm accordingly. The authors highlight that the aforementioned extreme token phenomena develop reinforcing each other during training. This observation is corroborated by a theoretical analysis showing that indeed "Attention logits grow logarithmically reinforced by small value states" and "Value state shrinks to a small constant vector reinforced by large attention logits". 

The final part of the paper is devoted to an investigation on the emergence of these phenomena in LLMs, with a particular focus on Llama 2-7B-Base and  OLMo-7B, both open-source. On the first model, a specific head is described in terms of its behaviour when the model is prompted with two different data sources. A very similar mechanism is shown to emerge, with the head being dormant on Wikipedia data (prose) as opposed to Github (code) data. Thanks to the availability of intermediate checkpoints for OLMo-7B, the authors are able to draw a parallel between their theoretical analysis of the learning dynamics of their toy model and that of OLMO, resulting in a good agreement.
weaknesses: - In section 2.1 (see Figure. 4), the considered one-layer transformer does not alternate the attention layer and the MLP along the depth dimension. Rather it uses these two components in parallel. The reason for this choice is not clear and should be better motivated. In addition, the trained models use pre-layer norm and it is not clear if residual connections are used and how.

- In the derivation of the loss(alpha, beta), it is not clear to me how the stable distribution is defined. Could you clarify?

- **Replacing SoftMax by ReLU attention removes extreme-token phenomena**: Can you provide more intuition on the role of the softmax? Its utilisation appears to be correlated with the emergence of extreme-token phenomena. The natural question then is whether this component is actually needed in practice and why modern LLMs use it. Could you elaborate on this?

- A question on Theorem 3, first point: I would have expected the logits to grow logarithmically as beta decreases (i.e. small value states). Instead, no constraint on beta seems to be needed for the logits to grow. I would appreciate it if the authors could clarify this point.

- As for the emergence of residual peaks, the experiments show that 3 layers are needed for the residual peaks to appear. First, are the authors now using a standard transformer architecture with attention layers interleaved with mlps? Second, do the authors have any intuition on why at least 3 layers are needed?

- In the construction of the token embeddings, the authors are assuming them to have dimension V, i.e. they are represented as one-hot vectors. This implies different tokens are orthogonal to one another. How important is this assumption?

-  In the BB example, it is shown that a head can have sinks yet being non-dormant (due to trigger tokens). Is this case contemplated by the theory? This seems to be quite general and is reflected in real LLMs (fig. 18 and fig. 19 in the appendix). To this extent, fig. 6 seems to be a pretty special case, where, for wikipedia data, the head is actually dormant as there are only two sink tokens and nothing else.

- How did you practically identify the heads exhibiting sink tokens in real LLMs? 

- Are the results obtained for OLMO also valid for Pythia (for which intermediate checkpoints are also available)?

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: - What is the motivation for using the BB task and designing the trigger and non-trigger tokens?
- On Lines 80 and 87, the word "mechanism" appears twice consecutively, which seems to be a typo.
- In Claim 4, could you clarify the case where "adding any value state from previous tokens worsens the prediction"?
- Why does the causal mask cause training dynamics to favor the 
 token? In Figure 6, attention is also given to token 3, which is described as semantically unmeaningful. So, is the cause of the attention sink truly the causal mask, or is it the lack of semantic meaning in certain tokens?
soundness: 2
strengths: - The paper addresses the attention sink phenomenon in large language models, an important and timely issue.
- The analysis using a toy example is original, providing intuitive insights and interpretations of the problem.
- The motivation is clear and well-articulated.
- The paper is well-organized, and the figures are highly illustrative.
summary: This paper examines the "attention sink" phenomenon in large language models, focusing on three related concepts: attention sinks, value-state drains, and residual-state peaks. Through a toy Bigram-Backcopy task, the paper reveal an active-dormant mechanism in which the attention head assigns dominant weight to the initial token in the dormant phase and shifts focus to relevant context tokens in the active phase. They attribute this behavior to a mutual reinforcement mechanism. Extending this analysis to real large language models, the authors show that this mechanism also emerges during pre-training, suggesting a understanding of the issue and potential mitigation strategies.
weaknesses: My primary concern is the statistical significance of the observed results, which raises questions about the validity of the conjecture regarding the active-dormant mechanism. For instance, in the BB task toy experiments, a simplistic generation rule (bigram transitions) is assumed, making the behavior of trigger and non-trigger tokens very distinct. This makes it difficult to accept that the theory and conjecture are applicable to real large language models, where such clear distinctions between tokens (trigger vs. non-trigger) do not exist. Section 3, a specific attention head in a specific layer is shown to exhibit behavior similar to the conjectured mechanism from Section 2. However, I find this unconvincing, as demonstrating a single head is insufficient to establish that the active and dormant phases align with what was observed in Section 2 for the same reasons.

My second concern is the lack of demonstrated practical benefits from this analysis. Although the authors suggest that this study could offer potential mitigation strategies, the paper does not validate the effectiveness of these mitigations, which limits the overall impact and applicability of the work.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: See weaknesses.
soundness: 3
strengths: 1. The study makes clear contributions by identifying the active-dormant mechanism and mutual reinforcement dynamic, helping to explain complex behaviors in transformer-based models.

2. The paper rigorously combines theoretical characterizations with empirical validation, making its claims robust and well-supported.
summary: This paper examines the mechanisms behind the "extreme-token phenomena" in large language models (LLMs), such as attention sinks, value-state drains, and residual-state peaks. These phenomena are collectively attributed to the "active-dormant" mechanism in transformer-based LLMs. The study first demonstrates these phenomena on simpler models trained on the Bigram-Backcopy (BB) task, identifying a mutual reinforcement dynamic driving these behaviors. The authors then extend their analysis to pre-trained LLMs (e.g., Llama and OLMo) and suggest mitigation techniques, such as replacing SoftMax with ReLU or switching optimization methods to control these extreme-token phenomena.
weaknesses: 1. There remains a gap between the BB Task and real language scenarios. Although I recognize the potential practical value of the authors' intuitive explanation of these phenomena, I would still expect authors to validate their explanations in real scenarios. For example, it would be helpful to confirm whether deactivating attention heads with attention sinks truly has no impact on the final output.

2. The paper’s explanation for why Adam tends to cause residual state peaks is insufficient, as it mainly describes the phenomenon without deeper analysis. Furthermore, I did not see an evaluation of the impact of applying SGD and ReLU to real language models. In my opinion, both SGD and ReLU come with certain drawbacks, which should be discussed.

3. Here is a potentially relevant study [1] where the "broadcasting" mechanism mentioned has a formation process similar to the attention sink.

[1] Anchor function: a type of benchmark functions for studying language models. arXiv:2401.08309.

Meta-review: This paper investigated the phenomena called "extreme-token phenomena" observed in transformer-based large language models that include attention sinks, value-state drains, and residual-state peaks. Based on a toy task with a small model, the authors identify an active-dormant mechanism. They develped theoretical analysis and empirical verification on both toy model and open-sourced pretrained LLMs.

Reviewers appreciate the valuable insights into the extreme-token phenomena and praise their theoretical and empirical analysis. However, there are common concerns on the simplified assumptions in the BB task, the limited robustness of the findings in real-world tasks, and the lack of practical benefits in the mitigating those phenomena.

All reviewers remain at borderline on assessing the contribution of this study at its current form after rebuttal. It is not sufficient for publication at ICLR yet. I would encourage authors to further develop their methodology to fill the gap between the simplified setup and the real-world application setting.

Decision: Reject