PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Tianyu Guo, Druv Pai, Yu Bai, Jiantao Jiao, Michael I. Jordan, Song Mei

arXiv:2410.13835v2cs.LG

TL;DR

Extreme-token phenomena in LLMs create challenges for inference, quantization, and interpretability, yet their mechanisms are not satisfactorily explained. The paper studies simplified transformers trained on the Bigram-Backcopy task and pretrained Llama and OLMo models to identify those mechanisms. It finds active-dormant attention heads and mutual reinforcement underlying aligned static and dynamic behaviors, and shows that ReLU attention and SGD can mitigate them in the toy setting.

  • Problem

    Extreme-token phenomena challenge inference, quantization, and interpretability, while their underlying mechanisms remain insufficiently explained.

    Source: §1 · p. 2

  • Method

    The paper combines theoretical and empirical analysis of one- to three-layer transformers trained on the Bigram-Backcopy task with experiments on pretrained Llama and OLMo models.

    Source: §1 · pp. 2–3, §2 · p. 5

  • Results

    The analyses identify active-dormant attention heads and mutual reinforcement as mechanisms governing extreme-token phenomena, with BB predictions aligning with observations in pretrained LLMs.

    Source: §1 · pp. 2–3, §4 · p. 15

  • Takeaways & Limitations

    Replacing SoftMax with ReLU eliminates extreme-token phenomena in the BB task, while switching from Adam to SGD removes residual-state peaks.

    Source: §1 · p. 3

    It remains unclear whether eliminating extreme-token phenomena improves downstream inference and quantization, and the paper does not analyze Llama training dynamics because intermediate checkpoints are unavailable.

    Source: §3.2 · p. 13, §4 · p. 15

Abstract

from arXiv · show

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.

1 Introduction

Extreme-token phenomena—attention sinks, value-state drains, and residual-state peaks—create inference, quantization, and interpretability challenges, while their mechanisms remain insufficiently explained. This work attributes them to active-dormant attention heads and mutual reinforcement, validates these mechanisms in toy transformers and pretrained LLMs, and proposes mitigation strategies.

1 Introduction

  • Extreme-token phenomena comprise sink tokens with disproportionately high attention, smaller value-state norms, and larger residual-state norms, and they recur across pretrained LLMs.These behaviors are collectively termed extreme-token phenomena.Source: §1 · pp. 1–2
  • Prior work links these phenomena to challenges in long-context inference, quantization, and attention interpretability, but does not satisfactorily explain their mechanisms.Adding a special sink token prevents other tokens from becoming sinks but does not eliminate extreme-token phenomena.Source: §1 · p. 2
  • The paper identifies an active-dormant mechanism in attention heads and a mutual-reinforcement mechanism that jointly drive extreme-token phenomena during training.In the BB task, attention sinks and value-state drains reinforce one another until non-extreme queries produce nearly identical logits for extreme-token keys.Source: §1 · p. 2
  • Pretrained LLMs exhibit the same active-dormant behavior and mutual-reinforcement dynamics identified in the Bigram-Backcopy task.A Llama 2 head activates on GitHub data but remains dormant on Wikipedia data, while OLMo training dynamics show the corresponding stable phase.Source: §1 · p. 3
  • Theoretical BB results and empirical LLM results consistently show sink-logit concentration, with extreme-token-key logits nearly identical across non-extreme queries.The paper summarizes these aligned quantitative properties in Table 1.Source: §1 · p. 3, Table 1 · p. 3
  • Replacing SoftMax with ReLU eliminates extreme-token phenomena in the BB task, while replacing Adam with SGD removes residual-state peaks.The authors discuss whether similar modifications could mitigate these phenomena in LLMs.Source: §1 · p. 3

1.1 Related work

  • The paper builds on prior studies of attention sinks, transformer dynamics, mechanistic interpretability, and quantization while distinguishing its active-dormant and mutual-reinforcement account.It presents this account as a mechanistic explanation supported by toy-model training dynamics and LLM evidence.Source: §1.1 · pp. 3–4

1.2 Preliminaries and notations

  • The paper formalizes transformer components, attention states, residual states, and logits under a pre-norm causal-attention architecture with residual connections.It uses zero-indexed layer and head notation and treats the Beginning-of-Sequence token as a common extreme token.Source: §1.2 · pp. 4–5

2 Extreme-token Phenomena in the Bigram-Backcopy Task

This section examines extreme-token phenomena in the Bigram-Backcopy task, focusing on active-dormant attention heads, attention sinks, value-state drains, and residual-state peaks. It analyzes their mutual-reinforcement mechanism and how their emergence or mitigation depends on attention type, optimization method, and model depth.

2 Extreme-token Phenomena in the Bigram-Backcopy Task

  • The Bigram-Backcopy task combines Markov bigram transitions with backcopying after three fixed trigger tokens in a vocabulary of 64 tokens.The remaining 61 tokens are non-trigger tokens, and sequences begin with the ⟨s⟩ token.Source: §2 · p. 5

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

  • A one-layer transformer learns attention sinks at ⟨s⟩ for non-trigger tokens, while trigger-token attention concentrates on preceding positions and ⟨s⟩ exhibits a value-state drain.The model reaches near-optimal accuracy after 10,000 Adam steps; the attention head handles Backcopy while the MLP handles Bigram prediction.Source: §2.1 · pp. 5–6, Figure 2 · p. 6
  • The active-dormant mechanism makes the attention head active on trigger tokens but dormant on non-trigger tokens, where it assigns dominant weight to ⟨s⟩ while adding minimal residual value.On active tokens, the head attends to relevant context and substantially affects the model output.Source: §2.1 · p. 6

2.2 Analysis of a minimally-sufficient transformer architecture

  • Theorem 2 shows mutual reinforcement: small ⟨s⟩ value states drive logarithmically growing sink logits, while large sink logits shrink ⟨s⟩ value states toward a small constant vector.The dynamics eventually stabilize with large, nearly identical sink logits and small value states; causal masking favors ⟨s⟩ as the extreme token.Source: §2.2 · p. 10
  • Training experiments support the theory: ⟨s⟩ attention logits increase while its value-state norm decreases, and a reparameterized model reproduces the one-layer dynamics.The theoretical predictions are expected to extend when a head can remain dormant for prediction but previous-token values remain task-relevant; Section 3 tests this in LLMs.Source: §2.2 · pp. 10–11
  • ReLU attention makes Bigram and Backcopy risks reach Bayes risk after 200 steps without increasing ⟨s⟩ attention logits or shrinking its value state.This supports the prediction that removing SoftMax prevents mutual reinforcement, attention sinks, and value-state drains.Source: §2.2 · p. 11

2.3 The emergence of residual-state peaks

  • Residual-state peaks occur in three-or-more-layer transformers but not single- or two-layer transformers trained on the Bigram-Backcopy task.The experiments track residual-state norms after layer 0; additional results are provided in Appendix B.1 and B.2.Source: §2.3 · p. 11
  • Zeroing layer 0 removes the massive residual state and returns middle-layer attention logits and value-state norms to normal levels.The intervention confirms that the residual-state peak contributes to attention sinks and value-state drains in the middle layer.Source: §2.3 · p. 11
  • Adam training produces linear growth of the ⟨s⟩ residual-state norm at layer 0 in three-layer transformers.Figure 7c reports this growth during pretraining on the Bigram-Backcopy task.Source: §2.3 · p. 11
  • Switching from Adam to SGD eliminates residual-state peaks while attention sinks remain, and switching to ReLU attention also eliminates the peaks.These results indicate that residual-state peaks depend on both SoftMax attention and the Adam optimization algorithm.Source: §2.3 · p. 12

3 Extreme-token Phenomena in pretrained LLMs

Pretrained LLMs exhibit context-dependent active-dormant attention heads and training dynamics matching the BB model’s extreme-token predictions. In OLMo, attention sinks, shrinking value states, concentrated sink logits, and growing residual norms emerge together during pretraining.

3.1 Active-dormant mechanism in LLMs

  • The active-dormant mechanism predicts that sink formation depends on context, with attention sinks and value-state drains indicating a head’s dormant phase.A head may be irrelevant for next-token selection in some contexts or tasks but useful in others.Source: §3.1 · p. 12

3.2 Extreme-token phenomena along training dynamics of LLMs

  • In OLMo Layer 24, attention on extreme tokens rapidly rises and stabilizes near [0.9, 1], while their value-state norms shrink below those of other tokens.These synchronized trends match the BB task’s mutual-reinforcement prediction.Source: Figure 9 · p. 13, §3.2 · p. 14
  • Residual norms of extreme tokens increase linearly during pretraining, whereas other tokens’ residual norms converge to small values.This behavior matches the residual-state dynamics predicted by the BB model.Source: Figure 9 · p. 13, §3.2 · p. 13

3.1 Active-dormant mechanism in LLMs

  • Attention heads can be dormant on prose but active on code: removing Llama 2’s L16H25 hurts GitHub performance while barely affecting Wikipedia performance.The head is dormant when it forms an attention sink and active when it does not.Source: §3.1 · p. 13

3.2 Extreme-token phenomena along training dynamics of LLMs

  • OLMo forms attention sinks despite lacking a ⟨s⟩ token, with the first token always acting as a sink and the period delimiter also becoming a sink in the analyzed prompt.The analysis used OLMo checkpoints because Llama does not provide intermediate pretraining checkpoints.Source: §3.2 · pp. 13–14
  • During OLMo training, sink-logit differences increase as value states decrease, then non-sink queries converge to nearly identical logits for the sink key.The logit concentration is supported by low variance for logits corresponding to ⟨s⟩ and Delim.Source: §3.2 · p. 14

4 Conclusions

The paper argues that extreme-token phenomena arise from active-dormant attention heads and mutual reinforcement, based on BB-task analyses and pretrained-LLM observations. It also proposes simple mitigation strategies, while leaving their downstream benefits unresolved.

4 Conclusions

  • The study identifies active-dormant attention heads and mutual reinforcement as mechanisms underlying attention sinks, value-state drains, and residual-state peaks.Evidence comes from one- to three-layer transformers trained on the BB task and experiments on pretrained LLMs.Source: §4 · pp. 14–15
  • Replacing SoftMax with ReLU eliminates extreme-token phenomena in the BB task, while replacing Adam with SGD removes residual-state peaks.The paper presents these changes as potential pathways for mitigating the phenomena during LLM pretraining.Source: §4 · p. 15
  • The paper does not establish whether eliminating extreme-token phenomena improves downstream inference or quantization performance.The authors identify this as an open question requiring costly large-scale pretraining experiments.Source: §4 · p. 15

A Proofs of Theorem 1 and 2

The appendix proves that a one-layer transformer can represent the Bigram-Backcopy task and recover its ground-truth transitions in the limit. The construction uses sufficient embedding dimension, explicit attention and MLP weights, and asymptotic parameter scaling.

A Proofs of Theorem 1 and 2

  • The appendix establishes positivity and spectral properties of the matrices used in the theoretical analysis.Several lemmas show diagonal dominance, positive semidefiniteness, and lower bounds needed for subsequent arguments.Source: A Proofs of Theorem 1 and 2 · p. 22

A.1 Proof of Theorem 1

  • The proof relies on a large embedding dimension d ≥ V N + 1 to provide enough orthonormal basis vectors, although the authors note this condition may be relaxed.The stronger dimension condition is adopted to simplify the construction and is unrelated to Theorem 2.Source: §A.1 · p. 23
  • A one-layer transformer can generate the BB task’s ground-truth transitions in the limit under Assumption A.The construction handles Markov transitions for non-trigger tokens and copies the preceding token for trigger tokens.Source: §A.1 · p. 23
  • The proof constructs orthonormal embeddings and attention queries that separately encode trigger-token copying and non-trigger-token transitions.The query matrix assigns trigger tokens position-dependent directions, while non-trigger tokens receive token-dependent query directions.Source: §A.1 · p. 23
  • The MLP maps token embeddings to logarithmic transition probabilities, yielding the required output distribution after the attention construction.For non-trigger tokens, the MLP output is log pv, while trigger-token attention retrieves the preceding token.Source: §A.1 · p. 24

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

The stable-phase proof characterizes global minimizers of the loss and establishes their second-order stability. It derives the gradient-flow equations used to analyze convergence dynamics.

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

  • The Hessian at these minimizers is positive semidefinite, establishing second-order stability of the stable phase.The proof reduces the Hessian analysis to a matrix B and shows B is positive semidefinite.Source: §A.2 · pp. 24, 27, 28
  • The appendix derives the gradient-flow system for α and β from the loss gradients.These ordinary differential equations provide the dynamical basis for the attention-sink and value-state-drain analyses.Source: §A.2 · p. 25
  • The stable-phase parameters α⋆ = α · 1 and β⋆ = c · 1 − e^−α · M ◦ ξ are stationary points and global minimizers.The result assumes nonnegative ξv, positive πv, and non-identical values of Mi · ξi.Source: §A.2 · p. 26

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

The appendix analyzes attention-sink and value-state-drain dynamics under gradient flow, then documents the minimal structures associated with residual-state peaks. These results connect asymptotic parameter behavior to the observed extreme-token phenomena.

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

  • Under fixed β, the attention parameter grows asymptotically as 2 log t plus a uniformly bounded remainder.The proof also shows the relevant scalar is monotone increasing under the stated assumptions.Source: §A.3 · pp. 28–29

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

  • With fixed α, β converges toward a value determined by the weighted quantities Mvξv, yielding the value-state-drain limit.The proof uses positive semidefiniteness and a Polyak–Łojasiewicz condition to establish asymptotic convergence.Source: §A.4 · pp. 29–30

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

  • Clear residual-state peaks occur in the three-layer 3 × TF and 2 × TF + mlp architectures.Figure 11 compares the residual-norm difference between the ⟨s⟩ token and other tokens across model-structure combinations.Source: §B.1 · p. 31, Figure 11 · p. 31

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

Additional BB-task experiments characterize three-layer extreme-token behavior and propose a mechanism for residual-state growth. The appendix links that growth to mutual reinforcement, layer normalization, and optimizer dynamics.

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

  • Three-layer BB transformers exhibit attention sinks, value-state drains, and residual-state peaks, although their final-layer residual pattern differs from that observed in LLMs.The last layer increases non-⟨s⟩ residual norms, causing ⟨s⟩ norms to match other tokens at the output.Source: §B.2 · p. 31, Figure 12 · p. 31, Figure 13 · p. 32

B.3 Potential mechanism for linear growth of the residual state peak in multilayer models

  • The proposed mechanism attributes residual-state peaks to aligned residual gradients, gradient decay from layer normalization, and Adam converting diminishing gradients into constant updates.The resulting updates can produce linear growth in the extreme token’s residual-state norm.Source: §B.3 · p. 32
  • In the scalar simulation, Adam produces linear growth of m after warm-up, whereas SGD with the same learning rate keeps m small.The simulation uses Adam β1 = 0.9, β2 = 0.999, weight decay = 10−8, and learning rate 0.3.Source: §B.3 · p. 33, Figure 15 · p. 33

C.1 Experimental details

  • The BB experiments use positional embeddings, prelayer normalization, SoftMax attention, ReLU MLPs, and separate Adam and SGD learning rates.Unless otherwise specified, training uses batch size B = 512, sequence length N = 256, and 10,000 steps.Source: §C.1 · p. 34

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

Additional BB-task experiments show that attention sinks persist across inputs, simplified dynamics match the one-layer transformer, and removing the start token shifts sink behavior to trigger or delimiter tokens.

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

  • Attention sinks appear across additional one-layer BB-task attention maps, although some non-trigger tokens exhibit different patterns and triggers can become sinks.Source: §C.2 · p. 34, Figure 16 · p. 34

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

  • The simplified model reproduces the one-layer transformer's observed statics and dynamics when trained with trainable MLP, O, α, β, λ, and ξ parameters.The simulations use Adam with learning rate 0.03 and show matching behavior in Figures 17 and 18.Source: §C.3 · p. 34, Figure 17 · p. 35, Figure 18 · p. 35

C.4 The Bigram-Backcopy task without the ⟨s⟩token.

  • Without the ⟨s⟩ token, the zeroth token is not a sink; trigger and delimiter tokens instead become sinks, matching delimiter extremes observed in LLMs.Source: §C.4 · p. 34, Figure 19 · p. 35

D More Attention Heads in Dormant and Active Phase

Additional Llama 2-7B-Base heads exhibit the active-dormant mechanism, becoming dormant on some inputs while remaining active on others.

D More Attention Heads in Dormant and Active Phase

E Fine-Grained Static Mechanisms for Extreme-Token Phenomena

Fine-grained Llama 3.1-8B-Base analyses attribute extreme-token phenomena to token-dependent query–key geometry, semantic differences in value states, and early-layer MLP contributions to residual peaks.

Attention sinks and global contextual semantics

Value-state drains

Residual state peaks

F Extreme-Token Phenomena Over Many Samples

Across many samples, Llama 3.1-8B and OLMo-7B reproduce the paper's extreme-token patterns: attention concentrates on sinks, value norms drain, and residual norms selectively grow during training.

F Extreme-Token Phenomena Over Many Samples

  • Across 128 Wikipedia samples, nearly all attention-sink heads concentrate almost all attention mass on the ⟨s⟩ token.The aggregate distribution of A_{b,:,j} shows large mass at j=0 and small mass for j≥1.Source: F Extreme-Token Phenomena Over Many… · p. 39, Figure 25 · p. 39
  • In OLMo-7B layer 24, attention sinks form in nearly all heads and samples early in training, with the top tokens' attention mass nearly summing to 1.Source: Figure 26 · p. 40
  • During OLMo-7B training, non-sink logits stabilize after sink formation, the first sink's value norm remains much lower, and token 0's residual norm grows linearly.Source: Figure 26 · p. 40

G.1 Multiple attention sinks vs. one attention sink

Attention-sink multiplicity differs across models and data domains: some heads develop multiple sinks, while others retain only the first token. Shuffling experiments further indicate that positional information contributes to sink formation, especially without a fixed ⟨s⟩ token.

G.1 Multiple attention sinks vs. one attention sink

  • Multiple sinks appear in the BB task, Llama 2-7B-Base, and OLMo, whereas Llama-3.1-8B-Base shows only one sink on prose data.For the multi-sink group, ⟨s⟩ and the first delimiter are typically sinks; in Llama 3.1, ⟨s⟩ remains the sink.Source: §G.1 · p. 41
  • The authors hypothesize that coding-data prevalence helps explain sink multiplicity because delimiters carry more semantics in coding than in prose.They report approximately 411B coding tokens for OLMo, at most 0.16T for Llama 2, and around 2.6T for Llama 3.1.Source: §G.1 · p. 41

G.2 The role of a fixed ⟨s⟩token in the Active-Dormant mechanism

  • When OLMo inputs are shuffled, attention sinks still form on early positions, indicating that positional and token-semantic information both contribute without a fixed ⟨s⟩ token.At Layer 24, the first token and first delimiter share attention mass even when the sentence is ungrammatical; the figure shows a sink on token 0 despite changed token semantics.Source: §G.2 · p. 41, Figure 27 · p. 41
Loading 2410.13835v2…