KV Shifting Attention Enhances Language Modeling

ArXiv: 2411.19574

🎯 Pitch

A simple linear interpolation of consecutive keys and values lets a single attention layer learn the 'induction head' pattern—spotting [A][B] to predict [B]—that standard Transformers need two layers to approximate. This design cuts the depth and width needed for in-context learning and, when deployed in 2.9B–19B parameter pre-training runs, yields consistently lower loss and a +2-point accuracy gain across multiple benchmarks.

---

1. Executive Summary

This paper proposes KV shifting attention, a modification to the standard transformer attention mechanism that decouples keys and values by linearly interpolating each token's key and value with those of its preceding token, to simplify the formation of induction heads—circuits that enable next-token prediction by locating previous occurrences of the current token and copying their successors (e.g., matching the sequence [A][B]...[A]→[B]). The authors theoretically prove that a single-layer KV shifting attention transformer can exactly represent induction heads, whereas a two-layer standard transformer can only approximate them, and they validate this experimentally on toy induction tasks—demonstrating that a one-layer KV shifting model learns induction faster than a two-layer vanilla model and succeeds at hidden-size 8 where the latter fails. Scaling to full language model pretraining at 2.9B and 19B parameters on private web-scale data, KV shifting attention achieves higher accuracy across eight benchmarks (e.g., 38.57% average versus 36.45% for the vanilla 2.9B model at 500B tokens) and consistently lower training loss, establishing that a lightweight architectural bias toward induction accelerates convergence and improves downstream performance without requiring additional layers or width.

2. Context and Motivation

The Core Problem: Transformers Waste Depth and Width on Induction

The fundamental problem this paper addresses is an architectural inefficiency in transformer-based language models: the standard attention mechanism requires a two-layer circuit to implement induction heads, and this requirement consumes both depth and width that could be used for other operations. Induction heads—circuits that enable a model to identify that a current token appeared previously in the sequence, then predict that the token following that previous occurrence will appear next—are widely believed to be a critical mechanism underpinning in-context learning (ICL). If transformers are structurally forced to allocate multiple layers and a significant amount of hidden dimension to what is fundamentally a simple pattern-matching operation, then they are operating below their potential capacity.

This is not a hypothetical concern. The paper cites Sanford et al. (2024a), which provides a formal proof that "a layer of transformer cannot implement induction heads." The consequence is that every transformer model currently deployed—regardless of scale—must dedicate at least two attention layers to implement what the paper's authors argue could be done in one, with less width. For models with finite depth (e.g., 32 layers in the 2.9B configuration, 48 layers in the 19B configuration), this represents a real opportunity cost. Every layer spent on induction is a layer that cannot be used for other forms of computation.

The gap this paper identifies can be stated precisely:

"Although there are many works based on transformer for analysis of induction heads, there are few works that utilize analysis of induction heads to modify transformers to enhance their ability to learn induction heads."

In other words, the mechanistic interpretability community has done extensive work understanding induction heads—how they form, what they do, why they matter for ICL—but almost no one has taken the next step of redesigning the attention mechanism itself to make induction heads easier to form. This paper fills that gap by proposing a minimal architectural change—KV shifting attention—motivated directly by the structure of the induction head circuit.

Why This Problem Matters: Theoretical and Practical Significance

Theoretical significance: a provable reduction in representational requirements. The paper provides a constructive proof (Theorem 2) that a single-layer, single-head KV shifting attention transformer with hidden dimension $D = d$ can exactly represent the induction heads function. Compare this to the standard transformer, which Theorem 1 (adapted from Wang et al., 2024) shows requires two layers and $D = 2d$ to approximate the same function, with an error bounded by $O(e^{-p^{(1)}})$ depending on the Alibi positional encoding bias in the first layer. The difference between exact representation and approximation-with-error is not just aesthetically pleasing—it means that KV shifting attention eliminates a source of noise (the Alibi bias from the copy operation in the first layer) that the vanilla two-layer circuit must tolerate. As the authors note:

"Since the copy operation in the first layer of the vanilla transformer introduces noise due to Alibi's bias, the final upper bound is bounded by a quantity related to Alibi's bias. And the KV shifting attention, due to the absence of this noise, Eq. 9 takes an equal sign."

This is a clean theoretical result: KV shifting attention reduces the depth requirement from two layers to one and the width requirement from $2d$ to $d$, while simultaneously eliminating an approximation error source. For anyone who cares about the theoretical efficiency of transformer architectures, this is a meaningful improvement.

Practical significance: faster convergence, better performance, flatter optimization landscape. The practical implications emerge directly from the theoretical ones. If the model no longer needs to learn a two-layer circuit through gradient descent to implement induction, it can acquire induction capability faster (Figures 1a, 4a-b). If the minimum width required for induction is halved, the model can dedicate more of its hidden dimension to other computations, potentially learning richer representations (Figure 1b, where the vanilla two-layer model "failed to cover up one of the answers" at hidden size 8 while KV shifting succeeded). If the optimization landscape for learning induction is simplified, the model may be more robust to hyperparameter choices like learning rate (Figure 5c, where vanilla diverges at LR=1e-2 while KV shifting converges).

These are not marginal gains. The convergence speedup on toy induction tasks is dramatic—the one-layer KV shifting model in Figure 1a reaches perfect induction accuracy roughly 5-10× faster (in training steps) than the two-layer vanilla model. At the scale of large language model pretraining (2.9B parameters, 500B tokens), the gap in training loss (Figure 4a) persists throughout the entire training run, not just in the early stages, suggesting that the bias toward induction provides a lasting advantage rather than merely an initialization benefit.

Real-world impact: better language models with the same parameter count. The experiments on 2.9B and 19B parameter models (Table 2) demonstrate that KV shifting attention improves performance across a range of benchmarks—MMLU, CMMLU, HellaSwag, ARC, LAMBADA, Winogrande, MATH—without increasing model size. For production LLM deployments, this matters. A +2.12 percentage point improvement in average benchmark accuracy (38.57% vs. 36.45% for the 2.9B model at 500B tokens) from a modification that adds only 4 learnable parameters per attention head (Section 2.2) represents an unusually favorable cost-to-benefit ratio. The modification is lightweight enough to be compatible with existing training and inference acceleration frameworks (as the paper emphasizes in the Discussion, Section 5), making adoption practical rather than merely theoretical.

Prior Approaches and Where They Fall Short

Mechanistic interpretability without architectural intervention. The dominant prior approach to induction heads has been analysis, not modification. Elhage et al. (2021) introduced the mathematical framework for transformer circuits and identified induction heads as a key mechanism. Olsson et al. (2022) provided further empirical characterization showing that induction heads emerge during training and correlate with ICL capability. Subsequent work (Bansal et al., 2023; Conmy et al., 2023; Ren et al., 2024) has used induction heads as a lens for understanding model behavior, while other work (Bietti et al., 2024; Wang et al., 2024; Sanford et al., 2024b; Chen et al., 2024) has analyzed how induction heads are learned from a theoretical perspective. All of this work treats the transformer architecture as fixed and studies what happens within it. The paper's contribution is to invert this relationship: use the understanding of induction heads to change the architecture so that the desired behavior is easier to achieve. As the authors state in Section 2.1:

"can we make slight adjustments to the attention so that a single layer of attention can achieve the mechanism of induction heads?"

This is a fundamentally different question from "how does the existing architecture learn induction heads?" and represents a gap the prior work left unaddressed.

The virtual attention heads perspective identifies the problem but doesn't solve it. The paper engages with the concept of virtual attention heads from Elhage et al. (2021), which describes how attention heads in different layers cooperate by composing their weight matrices. For two adjacent layers (simplifying to ignore residual connections and MLPs), the virtual attention head operation is:

$X_2 = A_{h2} X_1 W_{h2}^{ov} = (A_{h2} A_{h1}) X_0 (W_{h1}^{ov} W_{h2}^{ov})$

The key insight is Property 1: with a causal mask, when virtual attention heads attempt to implement induction (using token $i$ as an intermediary to route information from token $i+1$ to future tokens), the attention weight $(A_{h2} A_{h1})_{j, i+1}$ can only sum over indices $k \geq i+1$:

$(A_{h2} A_{h1})_{j, i+1} = \sum_{k=1}^{N} A_{h2}^{j,k} A_{h1}^{k, i+1} = \sum_{k=i+1}^{j} A_{h2}^{j,k} A_{h1}^{k, i+1}$

This means the $(i+1)$-th token cannot be directly accessed through the $i$-th token—it must first have its information integrated into the $i$-th token's hidden state. As the paper explains:

"it is difficult for the model to indirectly utilize $j$ tokens to focus on the $(i+1)$-th token through the $i$-th token. In other words, in order for the $(i+1)$-th token to be output by future tokens using the induced heads mechanism, it must first integrate the information of the $i$-th token into its hidden states, even if the information of the $i$-th token is useless for predicting the $(i+2)$-th token."

This is the root of the width requirement: the hidden state of the $i$-th token must simultaneously carry its own identity (for pattern matching) and information about the $(i+1)$-th token (for pattern completion). This imposes a representational burden that increases the minimum required width. The virtual attention heads framework diagnoses this problem, but Elhage et al. (2021) did not propose modifying the attention mechanism to eliminate it. KV shifting attention does exactly that—by making the $(i-1)$-th token's value directly accessible when attending to the $i$-th token's key (and vice versa), it bypasses the need for a two-step virtual attention head composition.

Alternative architectures don't target induction specifically. The paper acknowledges that many alternative architectures exist—RWKV (Peng et al., 2023), Mamba (Gu & Dao, 2023), RetNet (Sun et al., 2023)—that attempt to address transformers' quadratic complexity or replace attention entirely. However, these architectures are not designed to specifically enhance induction head formation. The paper's position (Section 7) is that:

"transformers still have many excellent properties that cannot be replaced temporarily, especially their ability to retrieve and replicate previous information... Currently, popular language models still use Transformers as architecture."

In other words, these alternatives may solve other problems (efficiency) but don't necessarily improve on transformer's induction capability, and transformers remain dominant. KV shifting attention is positioned as a minimal modification within the transformer family that specifically targets the induction head bottleneck.

Recent attention modifications don't address induction heads. The paper cites Differential Transformer (Ye et al., 2024a), which reduces noise in attention, and Selective Attention (Leviathan et al., 2024), which reduces attention to unneeded elements. These modify attention for different purposes—noise reduction, sparsity—but neither is motivated by or designed for induction head formation. The paper positions KV shifting as orthogonal and complementary:

"Our work is to slightly modify the attention to enhance its ability to learn induction heads, which potentially improves language modeling."

Token shift mechanisms existed in other domains and modalities, but not for this purpose. The paper explicitly connects KV shifting to a history of adjacent-token information fusion: Wu et al. (2018) used shift operations in vision CNNs, Zhang et al. (2021) applied token shift in video transformers, Li et al. (2023b) used adjacent token merging in speech transducers, and Peng et al. (2023) used token shift in RWKV for text. However, none of these prior uses were motivated by a reduction in the representational requirements for induction heads. The RWKV connection is particularly interesting: RWKV also shifts information from adjacent tokens, but its recurrent formulation is architecturally very different from a standard transformer, making it more of a departure than the lightweight modification proposed here. KV shifting attention can be seen as importing the principle of adjacent-token information fusion into standard multi-head attention, but with a precise theoretical justification (Theorems 1–3) for why it should help language modeling specifically through improved induction head formation.

How This Paper Positions Itself

The paper makes its positioning explicit in the introduction:

"Building on the perspective, we propose KV shifting attention, a novel approach designed to simplify and enhance the induction process. By decoupling keys and values in the attention mechanism, KV shifting attention reduces the structural requirements for depth and width, enabling single-layer transformers to effectively perform induction tasks."

The framing is that the paper is completing a missing step in the research pipeline: mechanistic interpretability identified induction heads → the limitations of the standard architecture for forming them were understood → a targeted architectural modification is proposed to address those limitations → the modification is validated both theoretically and empirically.

This is importantly different from proposing a new architecture from scratch and hoping it works better. The modification is minimal (four learnable parameters per head, an $O(ND)$ computational overhead versus $O(ND^2 + N^2D)$ for the main attention), directly motivated by a specific mechanistic analysis, and validated first on toy tasks that isolate the induction head mechanism before scaling to full pretraining. The paper is not claiming to have invented a fundamentally new class of models—it is claiming that the standard transformer has an identifiable, fixable inefficiency in how it implements one specific but important circuit, and that fixing it yields consistent gains at scale.

3. Technical Approach

3.1 Reader Orientation

This paper is primarily a theoretically-motivated architectural modification paper: the authors identify a specific inefficiency in how standard transformers implement induction heads—circuits that enable next-token prediction by matching the current token to a previous occurrence and copying its successor—and propose a minimal change to the attention mechanism that makes this circuit easier to form. The problem being solved is that standard attention requires two layers and substantial hidden-state width to implement induction, which wastes representational capacity that could be used for other computations; the solution is to decouple keys and values from their tokens, allowing the model to access a token's value by attending to its neighboring token's key (or vice versa), which collapses the required circuit from two layers to one and halves the minimum width.

3.2 Big-Picture Architecture (Diagram in Words)

The KV shifting attention system has three major components, all operating within a standard decode-only transformer architecture with a lightweight modification at the attention level:

1. The KV shifting mechanism within each attention head: Before computing attention scores, each token's key vector and value vector are linearly interpolated with those of the immediately preceding token. This introduces four learnable scalar parameters per head ($\alpha_1$, $\alpha_2$, $\beta_1$, $\beta_2$) that control the mix between a token's own key/value and the shifted (previous token's) key/value. The result is two new matrices, $\hat{K}$ and $\hat{V}$, which replace the original $K$ and $V$ in the standard attention computation.

2. The standard attention computation using the shifted keys and values: The query matrix $Q$ (which is unmodified) attends to $\hat{K}$ via scaled dot-product attention with a causal mask, producing attention weights that are then multiplied by $\hat{V}$ to produce the output. The rotary positional embedding (RoPE) is applied to $Q$ and the original $K$ before the shift—or equivalently applied to $Q$ and $\hat{K}$ after the shift, depending on implementation—so positional information is preserved.

3. The training and inference integration: The KV shifting operation is inserted as a pre-processing step before the flash attention call. During training, the full sequence is shifted using a convolution operation. During autoregressive inference, the shift is maintained using a cached state where the last token's key and value from the previous step are concatenated with the current token's key and value, and the linear interpolation is applied using the cached state as the "previous token" component. The entire operation adds $O(ND)$ computation—which is negligible compared to the $O(ND^2 + N^2D)$ of the main attention—and $4h$ additional parameters per layer, where $h$ is the number of key-value heads.

Information flows as follows: input hidden states $X$ → linear projections to form $Q$, $K$, $V$ → interpolation of $K$ with $\text{Shift}(K)$ to form $\hat{K}$, interpolation of $V$ with $\text{Shift}(V)$ to form $\hat{V}$ → application of RoPE to $Q$ and $K$ (unless already applied) → scaled dot-product attention $\text{Softmax}(Q\hat{K}^T \cdot M / \sigma) \hat{V}$ → output projection $W_O$. The only modification from standard attention is the interpolation step that produces $\hat{K}$ and $\hat{V}$.

3.3 Roadmap for the Deep Dive

• First, the motivation from induction heads and virtual attention heads (Section 2.1), which establishes why standard attention is structurally deficient for induction—this is the intellectual foundation for the entire modification.
• Second, the precise mathematical definition of KV shifting attention (Section 2.2), including the interpolation formulas, the learnable parameters, the initialization scheme, and the cost analysis.
• Third, the theoretical analysis (Section 3.1 and Theorem 2) proving that a single-layer KV shifting transformer can exactly represent induction heads with half the width of a two-layer standard transformer—this justifies why the specific interpolation pattern (attending to token $i$'s key to access token $i-1$'s value) is the right one.
• Fourth, the learning dynamics analysis on a simplified toy model (Section 3.2, Theorem 3) showing that KV shifting attention creates a loss landscape where induction is easy to learn—this explains the faster convergence observed empirically.
• Fifth, the empirical validation on controlled tasks (induction data, n-gram data, hop-k, iGSM) that verify the theoretical predictions and isolate the specific capabilities KV shifting enhances versus those it doesn't.
• Sixth, the large-scale integration details (Section 4.1) that describe how KV shifting is implemented in a full Llama-style architecture with GQA, the hyperparameter configurations, and the training infrastructure.
• Seventh, the variant and ablation experiments (Sections 4.6, 4.7) that explore design choices—shifting only K or only V, longer shift windows, gating mechanisms—and validate that the simple two-token interpolation is optimal.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a architectural modification paper motivated by mechanistic analysis. The core idea is that the standard transformer attention mechanism imposes unnecessary depth and width requirements on the formation of induction heads, and that a minimal change—linearly interpolating each token's key and value with those of its predecessor—eliminates these requirements while being cheap enough in parameters and computation to be deployed in production-scale models.

---

3.4.1 The Motivation: What Induction Heads Need and Why Standard Attention Falls Short

To understand why KV shifting attention takes the specific form it does, one must first understand the mechanical requirements of an induction head.

What an induction head does. An induction head is a circuit that, given a sequence of tokens, detects that the current token (call it $A$) has appeared previously in the sequence, retrieves the token that came after that previous occurrence (call it $B$), and outputs $B$ as the prediction for the next token. The canonical example is the sequence [A][B]... [A] → [B]: the model sees $A$ at position $j$, finds that $A$ also appeared at position $i < j$, sees that $B$ followed $A$ at position $i+1$, and predicts $B$ at position $j+1$.

Why standard attention needs two layers: the copying bottleneck. In a single-layer standard transformer, when the token at position $j$ computes its attention output, it can only directly access tokens at positions $k \leq j$ (due to the causal mask). The attention weight $A_{j,k}$ determines how much of token $k$'s value vector is incorporated into token $j$'s output. To implement induction, token $j$ needs to:

1. Match: Recognize that token $A$ at position $j$ is the same as token $A$ at position $i$ by comparing their key and query vectors.
2. Retrieve: Access the value vector of token $B$ at position $i+1$—the token that followed the matched occurrence.

In a single layer, step 1 (matching) can be done: the query of token $j$ can attend strongly to the key of token $i$ if they represent the same token identity. But step 2 has a structural obstacle: when token $j$ attends to token $i$, it receives token $i$'s value vector, not token $i+1$'s. Token $i$'s value vector contains information about token $i$ itself, not about its successor. To get token $i+1$'s information, token $j$ would need to attend to token $i+1$ directly, but there's no mechanism for the matching step (which identifies $i$ as the relevant position) to route the retrieval step to position $i+1$ within a single attention operation. This is precisely what Sanford et al. (2024a) proved: a single-layer transformer cannot implement induction heads.

The standard two-layer solution uses what Elhage et al. (2021) termed a virtual attention head. Consider two adjacent layers $l=1$ and $l=2$, ignoring residual connections and MLPs for clarity:

$X_1 = A_{h1} X_0 W_{h1}^{ov}, \quad X_2 = A_{h2} X_1 W_{h2}^{ov}$

where $X_0 \in \mathbb{R}^{N \times D}$ is the input sequence, $A_{hl} \in \mathbb{R}^{N \times N}$ is the attention weight matrix for head $h$ at layer $l$, and $W_{hl}^{ov} = W_{hl}^{v} W_{hl}^{o} \in \mathbb{R}^{D \times D}$ is the output-value projection. The virtual attention head combines these two layers:

$X_2 = A_{h2} X_1 W_{h2}^{ov} = (A_{h2} A_{h1}) X_0 (W_{h1}^{ov} W_{h2}^{ov})$

The product $(A_{h2} A_{h1})$ is the effective attention matrix. The first layer's attention head acts as a copying head: for each token at position $i$, it attends to position $i-1$ and copies the previous token's information into position $i$'s hidden state (this is why Theorem 1 uses $p^{(1)}$ for the Alibi bias—it encodes the positional offset that biases attention toward the immediately preceding token). After this copying operation, position $i$'s hidden state now contains information about both token $i$ and token $i-1$. The second layer can then attend to position $i$ (which carries information about position $i-1$) to effectively retrieve the predecessor token's value.

The width requirement imposed by the copying step. The paper identifies a specific limitation of this virtual attention head approach through Property 1. With the causal mask and $j \geq i$:

$(A_{h2} A_{h1})_{j, i+1} = \sum_{k=1}^{N} A_{h2}^{j,k} A_{h1}^{k, i+1} = \sum_{k=i+1}^{j} A_{h2}^{j,k} A_{h1}^{k, i+1}$

The summation starts at $k = i+1$ because the causal mask prevents position $i$ from attending to position $i+1$ in the first layer—$A_{h1}^{i, i+1} = 0$. This means that to route information from token $i+1$ to token $j$ through token $i$, token $i$'s hidden state must already carry the information from token $i+1$. But under the causal mask, token $i$ cannot attend to token $i+1$—so how does this information arrive?

The answer is that the copying operation in the first layer goes in the opposite direction: token $i+1$ attends to token $i$ and copies token $i$'s information forward. But this is backwards for induction—induction needs token $i$ (the matched token) to carry information about token $i+1$ (its successor), not the other way around. The standard solution is to have token $i$ attend to token $i-2$ to retrieve the copied information about token $i-1$, which was placed there when token $i-1$ attended to token $i-2$, and so on. This relay-race mechanism works, but it means that each token's hidden state must encode not only its own identity but also information about its predecessor (for the matching step to work through the virtual attention composition). As the paper states:

"This imposes certain requirements on the dimensionality of hidden states. In other words, this imposes certain requirements on the width."

The hidden state must carry two pieces of information (current token identity + previous token identity) in a single vector, which means the effective dimensionality per token of useful information is halved, or equivalently, the total width $D$ must be larger to achieve the same representational fidelity.

---

3.4.2 The KV Shifting Attention: Mathematical Definition

The key insight is that the two-step process—(1) copy previous token information into the current token's hidden state, (2) attend to the current token to retrieve the previous token's value—can be collapsed into a single attention operation if we allow the model to attend to one token's key but receive a different token's value.

The core idea in operational terms. When token $j$ wants to implement induction, it needs to:

• Attend to the key of token $i$ (to find the matching previous occurrence of the current token).
• Receive the value of token $i-1$ or $i+1$ (to retrieve the token that appeared next to the matched occurrence).

If keys and values are bound to the same token, this is impossible in one step. If we decouple keys and values—allowing token $i$'s key to be associated with token $i-1$'s value—then attending to token $i$'s key can directly retrieve token $i-1$'s value, bypassing the need for a multi-layer copying circuit.

Why the causal mask constrains which neighbor we can use. The paper considers which neighbors of token $i$ could have their values accessed when attending to token $i$'s key: $\{i-1, i, i+1\}$. Token $i+1$ is immediately ruled out because at the time token $i$ is doing next-token prediction, token $i+1$ hasn't been computed yet—accessing its value would violate causality. This leaves $i$ (the standard case) and $i-1$ (the preceding token, which is already computed and visible under the causal mask). If we also consider the symmetric case—attending to token $i-1$'s key to access token $i$'s value—we arrive at a design where each token's key can be used to access either its own value or its neighbor's value (and vice versa), as long as the neighbor is valid under the causal mask.

Formal definition. For a single attention head, the KV shifting attention is defined as follows (Equations 4-6):

$Q, K, V = X W_Q, X W_K, X W_V$

where $X \in \mathbb{R}^{N \times D}$ is the input hidden states (sequence length $N$, hidden dimension $D$), and $W_Q, W_K, W_V \in \mathbb{R}^{D \times D}$ are the standard query, key, and value projection matrices.

The key and value matrices are then transformed by linear interpolation with a shifted version of themselves:

$\hat{K} = \alpha_1 K + \alpha_2 \, \text{Shift}(K)$ $\hat{V} = \beta_1 V + \beta_2 \, \text{Shift}(V)$

where:

• $\alpha_1, \alpha_2, \beta_1, \beta_2 \in \mathbb{R}$ are learnable scalar parameters, one set per attention head.
• $\text{Shift}(\cdot)$ is an operation that discards the last token's representation and pads a zero vector at the beginning. Specifically, for a matrix $M \in \mathbb{R}^{N \times D}$:
  • $\text{Shift}(M)_1$ (the first row) is the zero vector $\mathbf{0} \in \mathbb{R}^D$.
  • For $t \in \{2, \dots, N\}$, $\text{Shift}(M)_t = M_{t-1}$ (the $(t-1)$-th row of the original matrix).

The resulting attention computation uses the shifted keys and values with the standard causal mask:

$\text{Output} = \text{Softmax}\left(Q \hat{K}^T \cdot M / \sigma\right) \hat{V} W_O$

where $\sigma = \sqrt{D}$ (the scaling factor for dot-product attention), $M \in \mathbb{R}^{D \times D}$ is the causal mask, and $W_O \in \mathbb{R}^{D \times D}$ is the output projection matrix.

What this computes, operationally. For each token at position $j$, when computing its attention output:

• The attention score between token $j$ (query) and token $k$ (key) uses the interpolated key $\hat{K}_k = \alpha_1 K_k + \alpha_2 K_{k-1}$, which is a blend of token $k$'s own key and the preceding token's key.
• The value that token $j$ receives when attending to token $k$ is the interpolated value $\hat{V}_k = \beta_1 V_k + \beta_2 V_{k-1}$, which is a blend of token $k$'s own value and the preceding token's value.
• When $\alpha_1 = 0, \alpha_2 = 1$ and $\beta_1 = 0, \beta_2 \neq 0$, attending to token $k$'s key effectively means attending to token $k-1$'s key while receiving token $k$'s own value (times $\beta_1$). This is the configuration that implements induction heads (as shown in Theorem 2).
• When $\alpha_1 = 1, \alpha_2 = 0$ and $\beta_1 = 0, \beta_2 = 1$, the attention is equivalent to standard attention where attending to token $k$'s key returns token $k-1$'s value—which is a direct implementation of the copying operation that would normally require a separate attention layer.

The key flexibility is that each head learns its own $\alpha_1, \alpha_2, \beta_1, \beta_2$, so different heads can specialize: some can implement standard attention (both alphas and betas near 1), some can implement induction (shifting either K or V or both to achieve the neighbor-access pattern), and some can learn other mixture ratios for other purposes.

Why this specific form over alternatives.

First, the linear interpolation with learnable weights rather than fixed weights or a hard selection (always using the previous token) gives the model flexibility. A head that learns $\alpha_1 \approx 1, \alpha_2 \approx 0$ behaves like standard attention; a head that learns $\alpha_1 \approx 0, \alpha_2 \approx 1$ implements the shifted-key pattern needed for induction. The model can learn intermediate values for heads that benefit from partial shifting. The paper's analysis of the trained 2.9B model (Table 4) shows that different heads converge to different regions of the $(\alpha_1, \beta_1)$ space, confirming that this flexibility is actually used.

Second, the choice to shift by exactly one position (not two, not three) is motivated directly by the induction head mechanism. Induction requires accessing the token immediately following (or preceding) the matched occurrence—shifting by one position is the minimal window that enables this. The ablation experiment (Figure 8c) confirms that extending the shift window to $[i-2, i]$ or $[i-3, i]$ does not improve performance, consistent with the theoretical analysis that a single-position shift suffices for induction.

Third, the choice to shift both K and V symmetrically (with independent weights for each) rather than shifting only K or only V is motivated by the dual needs of induction: shifting K allows a token's key to represent a neighbor's identity (for matching), while shifting V allows a token's value to represent a neighbor's information (for retrieval). The ablation (Figure 8b) shows that removing either shift degrades performance.

Parameter count and computational cost. In a multi-head attention configuration with $h$ heads, each head gets its own set of four scalar parameters $(\alpha_1, \alpha_2, \beta_1, \beta_2)$, adding $4h$ parameters per layer. In group-query attention (GQA), which is used in the 2.9B and 19B models, the parameters are per key-value pair rather than per query head, adding $4h_1$ where $h_1$ is the number of KV heads. For the 2.9B model with 4 KV heads per layer and 32 layers, this adds $4 \times 4 \times 32 = 512$ parameters total—negligible compared to the billions of parameters in the rest of the model.

The computational overhead of the shift operation (Equation 5) is $O(ND)$ per head, arising from the element-wise addition and scalar multiplication of vectors of dimension $D$. This is dwarfed by the main attention computation (Equation 6), which is $O(ND^2 + N^2D)$. The paper notes this explicitly:

"The additional calculation caused by Eq 5 is $O(ND)$, which is much smaller than $O(ND^2 + N^2D)$ in Eq. 6."

For the group-query attention case used in practice, the additional computation is $O(N h_1 d_1)$ where $h_1$ is the number of KV pairs and $d_1$ is the head dimension.

Initialization strategy. The learnable parameters $\alpha_1$ and $\beta_1$ are initialized from the uniform distribution $U(0, 1)$, and $\alpha_2$ and $\beta_2$ are set to $1 - \alpha_1$ and $1 - \beta_1$ respectively. This ensures that at initialization:

• The sum of the interpolation weights is 1: $\alpha_1 + \alpha_2 = 1$ and $\beta_1 + \beta_2 = 1$.
• The interpolation is a convex combination of the current token and the previous token, so the magnitude of the resulting key and value vectors doesn't blow up or shrink relative to standard attention.
• The initial state is a random point in the simplex, meaning different heads start at different positions in the $(\alpha_1, \beta_1)$ space. As the paper notes in the analysis of Theorem 3:

"In practice, we often have many attention heads, $(\alpha_1, \beta_1)$ of some heads are closer to $(0, 1)$ during initialization, making it much easier for them to learn induction heads."

This diversity in initialization is important because heads that happen to initialize near the induction-friendly region $(\alpha_1 \approx 0, \beta_1 \approx 1)$ can adopt the induction role quickly, while others can specialize for different operations.

Consequences of not constraining the sum to 1. The paper notes that after training, $\sum_i \alpha_i$ and $\sum_i \beta_i$ deviate from 1—the values are no longer constrained to a simplex. Some parameters even become negative and "far from zero" (e.g., $\beta_2 = -0.15$ in the 3rd KV pair of the 17th layer of the 2.9B model). The paper experimented with enforcing the simplex constraint during training using either a gating mechanism ($\alpha_1 = \text{Sigmoid}(a), \alpha_2 = 1 - \alpha_1$ with $a$ being the learnable parameter) or clipping to $[0, 1]$, but found that:

"Using more controls does not make the model learn better. Allowing $\alpha$ and $\beta$ to have a wider range of degrees of freedom may enable the model to learn richer features."

This is an important design insight: while the simplex initialization provides a sensible starting point, forcing the model to stay on the simplex during training restricts it from learning patterns like anti-copying heads (where $\beta_2 < 0$ subtracts the previous token's value) that may serve other useful functions beyond induction.

---

3.4.3 Theoretical Analysis: Better Representation for Induction Heads

The paper provides a constructive proof that KV shifting attention can represent induction heads with strictly fewer resources than standard attention. This is established through two theorems placed in juxtaposition.

Definition 1 (Induction Heads Machine). The paper defines a formal model of an induction head using Alibi relative positional encoding (RPE) as follows:

$\text{IH}(x) = \sum_{s=2}^{L-1} \text{softmax}\left(x_L x_{s-1}^T / \sigma - m |L - s|\right) \, x_s$

where:

• $x \in \mathbb{R}^{L \times D}$ is a sequence of length $L$ with hidden dimension $D$.
• $x_L$ is the last token (the one doing the prediction).
• $x_{s-1}$ is the key being attended to (the token before the matched occurrence).
• $x_s$ is the value being retrieved (the token after the matched occurrence, which is the prediction target).
• $\sigma > 0$ is a temperature parameter controlling attention sharpness.
• $m > 0$ is the Alibi slope parameter that biases attention toward more recent positions (the term $-m|L-s|$ penalizes attending to tokens far from the current position).

Operationally, this function does the following: given a sequence, look at the current token $x_L$, find the closest previous position $s-1$ where a similar token appeared (by computing dot-product similarity $x_L x_{s-1}^T$), and then retrieve $x_s$—the token that followed that previous occurrence—as the output. The softmax over all $s$ ensures the retrieval is a weighted average, with the Alibi bias $-m|L-s|$ ensuring that recent matches are preferred over distant ones. This is exactly the induction head operation: match the current token to a previous occurrence and copy what came after it.

Theorem 1 (Standard Transformer Approximation, adapted from Wang et al. 2024). There exists a constant $C > 0$ and a two-layer single-head transformer TF (without FFNs), with $D = 2d$, $W_K^{(1,1)} = W_Q^{(1,1)} = 0$, $p^{(2)} = m$, and $\|W_K^{(2,1)}\|, \|W_Q^{(2,1)}\| \leq O(1, 1/\sigma)$, such that:

$\sup_{L \in \mathbb{N}^+} \|\text{IH} - \text{TF}\|_{L, \infty} \leq O(e^{-p^{(1)}})$

What this means in plain language: A two-layer standard transformer can approximate the induction heads function, but cannot represent it exactly. The error depends on the Alibi bias $p^{(1)}$ in the first layer—as this bias increases, the approximation improves, but the error never reaches zero. The first layer uses $p^{(1)}$ to implement a copying operation (attending to the immediately previous token), and the second layer uses the copied information to complete the induction. The transformer needs $D = 2d$ (twice the dimension of the token embeddings) because each hidden state must carry information from two tokens simultaneously. The weight constraints $W_K^{(1,1)} = W_Q^{(1,1)} = 0$ simplify the first layer to be purely position-based (no content-based attention), which is the copying head.

Why the error is bounded by $O(e^{-p^{(1)}})$: The proof (Appendix A) shows that the error comes from the copying operation not being perfect. The first layer's output at position $s$ contains a softmax-weighted sum of tokens at positions $1$ through $s-1$, with the weight for position $s-1$ being $\text{softmax}(-p^{(1)}(0)) = 1 / \sum_{\tau=0}^{s-2} e^{-p^{(1)}\tau}$. For finite $p^{(1)}$, this weight is strictly less than 1, meaning some fraction of the attention is distributed to positions $s-2, s-3, ...$—it's not a perfect copy. The approximation error is proportional to $e^{-p^{(1)}}$, so making $p^{(1)}$ very large drives the error down, but it can never be eliminated entirely.

Theorem 2 (KV Shifting Attention Exact Representation). There exists a one-layer single-head KV shifting attention KVSA, with $D = d$, such that:

$\text{IH} = \text{KVSA}$

How to construct the exact representation. Set the following parameters in the KV shifting attention (Equations 4-6):

• $\alpha_1 = 0$, $\alpha_2 = 1$: This means $\hat{K}_t = K_t \cdot 0 + K_{t-1} \cdot 1 = K_{t-1}$. The key of token $t$ is entirely replaced by the key of the previous token. Consequently, when the query of token $L$ computes dot products with the keys of all positions, it is actually comparing against the keys of positions $0, 1, ..., L-1$ (not $1, 2, ..., L$). The key of position $s$ (the shifted version) is $K_{s-1}$.
• $\beta_1 = 1/\sigma$, $\beta_2 = 0$: This means $\hat{V}_t = V_t / \sigma + V_{t-1} \cdot 0 = V_t / \sigma$. The value remains the token's own value, just scaled by $1/\sigma$.
• $W_Q = W_K = W_V = W_O = I \in \mathbb{R}^{d \times d}$: The projections are identity matrices, so $Q = X$, $K = X$, $V = X$. The tokens are embedded directly with no learned transformations.
• Alibi bias $p = m$ is applied in the attention softmax.

Now trace through the computation:

1. $\hat{K}_s = \alpha_1 K_s + \alpha_2 K_{s-1} = 0 \cdot x_s + 1 \cdot x_{s-1} = x_{s-1}$ for $s \geq 2$ (and $\hat{K}_1 = 0$ due to the zero-padding of $\text{Shift}(K)$).
2. $\hat{V}_s = \beta_1 V_s + \beta_2 V_{s-1} = (1/\sigma) \cdot x_s + 0 \cdot x_{s-1} = x_s / \sigma$.
3. The attention score for token $L$ attending to position $s$ is: $Q_L \hat{K}_s^T = x_L x_{s-1}^T$.
4. After adding the Alibi bias $-m|L-s|$ and applying softmax, the attention weight for position $s$ is $\text{softmax}(x_L x_{s-1}^T / \sigma - m|L-s|)$, where the scaling $\sigma$ enters because the softmax in standard attention divides by $\sqrt{D}$ before exponentiating, and in this construction the dot product $x_L x_{s-1}^T$ is not yet divided by $\sigma$—the $1/\sigma$ in $\beta_1$ effectively applies this scaling through the values.
5. The output is $\sum_{s=2}^{L-1} \text{softmax}(x_L x_{s-1}^T / \sigma - m|L-s|) \cdot x_s / \sigma$, which is exactly $\text{IH}(x) / \sigma$. Multiplying by $W_O = I$ and absorbing the scale factor yields the exact induction heads function.

What this proves about representational requirements.

First, depth reduction: a one-layer KV shifting transformer achieves what requires two layers in a standard transformer. The one-layer standard transformer cannot do induction at all (Sanford et al., 2024a); the two-layer standard transformer can approximate it; the one-layer KV shifting transformer represents it exactly.

Second, width reduction: the standard transformer construction needs $D = 2d$ because the first layer's output must concatenate the current token with copied information from the previous token. In the KV shifting construction, $D = d$ suffices because there is no need to store the previous token's information in the current token's hidden state—the shift operation makes that information accessible through a different path (the shifted K matrices). The hidden state dimension is literally halved for the same representational capacity.

Third, exactness: the standard transformer construction has approximation error $O(e^{-p^{(1)}})$—the first layer's copy is never perfect because the softmax distributes some attention mass to non-target tokens. KV shifting has zero error because the shift operation is a hard, deterministic copy—there is no attention-based soft copying that introduces noise. This is significant because the approximation error in the standard construction compounds across the depth of the model; KV shifting avoids this error source entirely.

The paper acknowledges a limitation of this theoretical framing:

"Theorems 1 and 2 only provide constructive upper bounds for implementing induction heads. A more rigorous statement would be to prove that the lower bound of Theorem 1 is smaller than the upper bound of Theorem 2."

In other words, the paper has shown what KV shifting attention can do (constructive upper bounds), but has not proven what standard attention cannot do (information-theoretic lower bounds) beyond the existing result that one layer cannot implement induction at all. The full separation would require showing that even with arbitrary depth, standard attention cannot match KV shifting in width efficiency, which is left as future work.

---

3.4.4 Learning Dynamics: Why KV Shifting Learns Induction Faster

Beyond representation (can the architecture express induction heads?), the paper analyzes how gradient descent learns induction heads in KV shifting attention versus standard attention. This is formalized in Theorem 3, which analyzes a drastically simplified setting.

The simplified setting. The following simplifications are made to make the learning dynamics analytically tractable:

• All residual connections, MLP layers, normalization layers, and positional embeddings are removed.
• Weight tying is used between the embedding and the output layer.
• Each component of each token's embedding is independently drawn from $\mathcal{N}(0, 1/d)$.
• All projection matrices are set to identity: $W_Q = W_K = W_V = W_O = I$.
• The sequence has length $T + 1$, and the vocabulary size is $T$ (where $T \geq 3$).
• Every token in the vocabulary appears exactly once in the sequence, except the last token (position $T+1$), which is identical to some earlier token (say position $i$). This creates the induction pattern: token $i$ appears at position $i$ and again at position $T+1$, and the model must predict what comes after position $i$ (which is token $i+1$).
• The cross-entropy loss is computed only on the prediction of the next token after the last position (i.e., predicting what follows $x_{T+1}$).

This setting isolates the induction head problem: the model sees a repeated token and must retrieve its successor. There are no other linguistic complexities—just the pure induction pattern.

Theorem 3. Under the simplified conditions described, and as $d \to \infty$, learning induction heads by KV shifting attention is equivalent to minimizing:

$\min L = -\log\left( \frac{e^{a_2 \beta_1 + \beta_2 / S}}{e^{a_2 \beta_1 + \beta_2 / S} + 2 e^{\beta_1 / S + \beta_2 a_2} + e^{2 a_1 \beta_1 + a_2 \beta_2} + O(T)} \right)$

where:

• $a_2 = e^{\alpha_2} / S$ is the normalized softmax weight for attending to the token following the matched occurrence (position $i+1$).
• $a_1 = e^{\alpha_1} / S$ is the normalized softmax weight for attending to the matched occurrence itself (position $i$).
• $S = e^{\alpha_2} + 2 e^{\alpha_1} + O(T)$ is the sum of all unnormalized attention weights, where $O(T)$ represents contributions from the $T-3$ other tokens that are not part of the induction pattern.
• $\beta_1$ and $\beta_2$ are the value interpolation parameters from Equation 5.

What this loss function means in operational terms. The term in the numerator, $e^{a_2 \beta_1 + \beta_2 / S}$, is the logit for correctly predicting token $i+1$ (the true successor). This logit depends on:

• $a_2$: how strongly the query attends to the key of position $i+1$, which carries $\hat{K}_{i+1} = \alpha_1 x_{i+1} + \alpha_2 x_i$. If $\alpha_2$ is large, $\hat{K}_{i+1}$ is dominated by $x_i$, and the query (which is $x_{T+1} = x_i$) will have high dot-product similarity with it—enabling the match.
• $\beta_1$: how much of token $i+1$'s own value is returned when attending to it. The correct prediction is token $i+1$, so $\beta_1$ should be large to emit the right token.
• $\beta_2 / S$: a small contribution from the previous token's value when attending to position $i+1$ (since $\hat{V}_{i+1} = \beta_1 x_{i+1} + \beta_2 x_i$, and attending to position $i+1$ and receiving $\beta_2 x_i$ adds some logit for token $i$, which is slightly wrong).

The three denominator terms are:

1. $e^{a_2 \beta_1 + \beta_2 / S}$ (same as numerator): the correct answer $i+1$.
2. $2 e^{\beta_1 / S + \beta_2 a_2}$: two tokens (one of which is the immediate successor of position $T+1$, and one of which is $T$ itself) that have logit contributions from the attention pattern. These are "distractors" that compete with the correct prediction.
3. $e^{2 a_1 \beta_1 + a_2 \beta_2}$: the matched token $i$ itself (which appears twice and thus gets a contribution from both positions).
4. $O(T)$: all other tokens in the vocabulary that receive roughly equal, small logit contributions.

Minimizing $L$ means maximizing the numerator (making the correct prediction $i+1$ have high logit) while keeping the denominator terms small.

What the gradient descent dynamics look like (Figure 2). The paper simplifies further by setting $\alpha_2 = 1 - \alpha_1$ and $\beta_2 = 1 - \beta_1$ (so the interpolation weights sum to 1), and treating $O(T)$ as a constant. The loss landscape in $(\alpha_1, \beta_1)$ space is plotted for three values of the distraction term:

• Small $O(T)$ (Figure 2a, $O(T) = 0$): The loss contour is tight around the optimum, and the gradient direction may undergo a small non-monotonic learning process—the trajectory isn't a straight line to the minimum. This means that when the vocabulary is small and distractors are few, learning can be slightly indirect.
• Medium $O(T)$ (Figure 2b, $O(T) = 10$): The contour broadens (convergence slows), but the gradient direction becomes more consistent, pointing more directly toward the optimum.
• Large $O(T)$ (Figure 2c, $O(T) = 100$): The contour becomes very flat, and convergence slows dramatically, but the gradient direction is highly consistent—no non-monotonic behavior.

Why this explains faster convergence for KV shifting attention.

The standard two-layer transformer must learn induction through a two-phase process described by Bietti et al. (2024): first learn global bigram statistics (simple token-to-token transitions), then form induction heads through a top-down mechanism that overrides the bigram memory for specific patterns. This is a complex, sequential learning process that requires coordinating two layers of attention.

In KV shifting attention, the learning problem collapses to optimizing four scalar parameters $\alpha_1, \alpha_2, \beta_1, \beta_2$ (per head), and as Theorem 3 shows, the loss landscape has a straightforward structure: the optimum for induction is at $(\alpha_1, \beta_1) = (0, 1)$ (or equivalently, using the previous token's key and the current token's value). If a head happens to initialize near this region (which some will, since the initialization is random uniform over $[0, 1]^2$), it can converge to the induction solution very quickly. The paper notes:

"In KV shifting attention, induction heads become very easy to learn, and even with good initialization, the model has induction capability."

This stands in contrast to the standard transformer, where induction capability must be constructed from scratch through gradient descent over all weight matrices. The paper also observes a trade-off:

"But it is difficult to have appropriate initialization to obtain a certain level of bigrams capability without training."

In other words, KV shifting attention has a strong bias toward induction that accelerates its acquisition (Figures 1a, 4a-b), but this same bias may make it harder to learn simple bigram statistics from initialization—which is fine because bigrams are easier to learn during training anyway, while the induction structure is harder to acquire without the architectural assistance.

---

3.4.5 Integration into Full Language Model Pretraining

The toy experiments validate the mechanism in isolation, but the paper's primary practical contribution is demonstrating that KV shifting attention works in production-scale language model pretraining. This section details how the mechanism is integrated into a Llama-style (Touvron et al., 2023) architecture.

Model configurations (Table 6). Two primary model scales are used for the main pretraining experiments:

Additional smaller-scale experiments are conducted at 1.5B, 6.7B, and 13B parameters with 10B training tokens each, using different configurations (detailed in Table 6).

Architecture details. The architecture is "similar to Llama2" (Touvron et al., 2023) with the following specifics:

• Group Query Attention (GQA) (Ainslie et al., 2023) is used for the 2.9B and 19B models to reduce inference memory. In GQA, the number of key-value heads is smaller than the number of query heads. For both the 2.9B and 19B models, there are 4 KV heads (versus 20 and 48 query heads, respectively). The KV shifting parameters are applied per KV head, not per query head, so the additional parameter count per layer is $4 \times 4 = 16$ for both models.
• Rotary Position Embedding (RoPE) (Su et al., 2024) is used with a base of 100,000. The paper notes that this is larger than the default base of 10,000 because "previous study (Xu et al., 2024) has shown that the longer the context length, the larger the base required, while the default base=10,000 is relatively small, even for 2048 windows." The RoPE is applied to $Q$ and $K$ matrices, and the authors specify in the code appendix (Appendix F) that RoPE is applied after the KV shifting interpolation (the interpolation and RoPE operations commute because RoPE is applied position-wise).
• Large vocabulary (48,000) to support multilingual environments. This is larger than Llama's original 32,000, reflecting the production-oriented nature of the baseline.

Compatibility with flash attention. The paper provides PyTorch code (Appendix F) showing that the KV shifting operation is implemented as a pre-processing step before calling flash_attn_func. The shift is computed using a convolution operation along the sequence dimension: each head's $\alpha$ and $\beta$ vectors (of shape [h, 2] for the two interpolation weights) are treated as 1D convolution kernels of size 2, and the input $K$ and $V$ tensors (of shape [batch, seq, h, d]) are convolved along the seq dimension to produce $\hat{K}$ and $\hat{V}$. This is efficient on GPUs because convolution over the sequence dimension for a kernel of size 2 is just a linear combination of adjacent positions, which can be implemented with vectorized operations.

Handling autoregressive inference. During training, the full sequence is available, and the shift is applied to the entire sequence at once. During inference, the model generates tokens one at a time, and the KV cache must be maintained. The paper provides inference-specific code (Appendix F) that works as follows:

• When generating the first token (no past key-value cache), the shift is applied using the convolution approach on the single-token sequence (which is trivial), and the resulting $k$ and $v$ for that token are cached.
• For subsequent tokens, the cached $k$ and $v$ from the previous token are used as the "shifted" component. Specifically:
  • $\hat{k} = \alpha_1 \cdot k_{\text{current}} + \alpha_2 \cdot k_{\text{previous}}$
  • $\hat{v} = \beta_1 \cdot v_{\text{current}} + \beta_2 \cdot v_{\text{previous}}$
• Then the previous token's $k$ and $v$ are updated to the current token's $k$ and $v$ for the next step.

This means the inference overhead is exactly one additional addition and multiplication per key and value vector per token—completely negligible compared to the attention computation.

Datasets. Due to commercial reasons, the paper uses non-public private data with collection and filtering methods similar to FineWeb-edu (Penedo et al., 2024). The authors state that "when computing resources are available, we will use open-source data, such as RedPajama-1T to train two models for comparison." This is a notable limitation—the training data is not publicly available, though the trained 2.9B models (both vanilla and KV shifting) have been released on HuggingFace.

Training hyperparameters (Table 6 and Section 4.1).

• Learning rate schedule: Constant learning rate with linear warmup. The warmup steps are 1,000 for the 1.5B/6.7B/13B models, 600 for the 2.9B model, and 3,000 for the 19B model.
• Batch sizes: 1M tokens (1,048,576 = 1024 × 1024) for 1.5B/6.7B, 16M for 2.9B, 2M for 13B, and 3M for 19B. The authors note: "1M here means 1,048,576 = 1,024 × 1,024, while elsewhere in the paper it refers to 1,000,000."
• Optimizer: AdamW with $\beta_1 = 0.9$, $\beta_2 = 0.95$, and weight decay = 0.1.
• GPU infrastructure: Training from scratch for 2.9B/19B is conducted on Nvidia H800-80G GPUs (512 GPUs for 2.9B, 240 A800-80G GPUs for 19B). Smaller-scale experiments use A100-80G GPUs.

The paper explicitly adopts "a large learning rate in practice, because Lobacheva et al. suggests large learning rates improve generalization." This is particularly relevant because the robustness experiments (Figure 5) show that KV shifting attention tolerates high learning rates better than vanilla attention—at LR=1e-2, vanilla diverges while KV shifting converges, suggesting "the optimization space for KV shifting attention may be flatter."

Evaluation benchmarks and metrics. For models trained with sufficient tokens (2.9B and 19B), eight benchmarks are used:

• LAMBADA (Paperno et al., 2016): word prediction requiring broad discourse context; tests long-range dependency understanding.
• Winogrande (Sakaguchi et al., 2021): adversarial Winograd schema challenge; tests commonsense reasoning.
• HellaSwag (Zellers et al., 2019): sentence completion; tests grounded commonsense inference.
• ARC-Easy and ARC-Challenge (Clark et al., 2018): grade-school science questions; tests reasoning and knowledge.
• CMMLU (Li et al., 2023a): Chinese multitask language understanding; tests multilingual knowledge.
• MMLU (Hendrycks et al.): massive multitask language understanding (57 subjects); tests broad knowledge.
• MATH (Hendrycks et al., 2021): competition-level mathematics; tests complex reasoning.

For models trained with only 10B tokens (1.5B, 6.7B, 13B), only validation loss is reported since they are undertrained for meaningful benchmark evaluation.

Evaluation metrics for MMLU are compared under three conditions (Table 3):

• Cloze (zero-shot): Following Waleffe et al. (2024), the model is evaluated in a cloze format that "intends to break away from the format of standard multiple-choice and directly measure the knowledge."
• Zero-shot: Standard zero-shot multiple choice.
• Few-shot (5-shot): 5 examples provided in context, testing in-context learning.

The design choice to use GQA and large vocabularies. The paper acknowledges that these choices are driven by production requirements rather than theoretical considerations. GQA reduces inference memory by sharing key-value heads across multiple query heads, which is important for serving latency and cost but changes the attention mechanism relative to standard multi-head attention. The KV shifting modification is applied at the KV head level (not the query head level), meaning $4h_1$ parameters are added where $h_1$ is the number of KV heads. This is explicitly an engineering decision that makes the modification practical for deployment while maintaining compatibility with GQA's benefits.

The RoPE base of 100,000 is also noted as a practical choice informed by concurrent work (Xu et al., 2024) showing that the default base of 10,000 is inadequate for the context lengths used (up to 12,288 for the 19B model). Since RoPE encodes relative position through rotation frequencies, a larger base means the rotations are less aggressive, preserving more positional distinguishability at longer ranges.

---

3.4.6 Variant Experiments and Design Ablation

The paper explores several design variants of the KV shifting mechanism to validate that the specific choices made are optimal and to understand the learned behavior.

Variant 1: Gating mechanisms (KV shifting gate). Constrain $\alpha_1 + \alpha_2 = 1$ and $\beta_1 + \beta_2 = 1$ throughout training using a sigmoid gating mechanism:

• $\alpha_1 = \text{Sigmoid}(a)$, $\alpha_2 = 1 - \alpha_1$
• $\beta_1 = \text{Sigmoid}(b)$, $\beta_2 = 1 - \beta_1$

where $a, b \in \mathbb{R}$ are the actual learnable parameters. This enforces the convex combination structure that the initialization starts with. Figure 8a shows that this variant performs slightly worse than the unconstrained version on the 1.5B model with 10B tokens. The paper's interpretation:

"Allowing $\alpha$ and $\beta$ to have a wider range of degrees of freedom may enable the model to learn richer features. For example, Elhage et al. (2021) discovered the presence of 'anti-copying prefix-search' heads in vanilla model. Although we don't know what its function is, if we restrict $\beta_2 \geq 0$, it is likely to limit the generation of this kind of heads."

In other words, some heads may benefit from having $\beta_2 < 0$ (subtracting the previous token's value) or having weights outside $[0, 1]$ to implement operations that go beyond standard induction or copying—the gating mechanism prevents this.

Variant 2: Clipping to [0, 1] (KV shifting 0 to 1). Instead of using a sigmoid, clamp the parameters to $[0, 1]$ after each update: $\alpha_i = \min(\max(\alpha_i, 0), 1)$ and similarly for $\beta_i$. This also enforces the simplex without restricting the sum. Figure 8a shows this variant is essentially identical to the unconstrained version—the clamping doesn't hurt, but also doesn't help, suggesting that the model naturally learns to keep most parameters in reasonable ranges without enforcement.

Ablation: Shifting only K or only V (Figure 8b). The paper tests whether shifting both key and value is necessary, or if shifting only one or the other suffices:

• No K shift (vanilla K): Set $\alpha_1 = 1, \alpha_2 = 0$ (K is not shifted), keep V shifting with learnable $\beta_1, \beta_2$.
• No V shift (vanilla V): Set $\beta_1 = 1, \beta_2 = 0$ (V is not shifted), keep K shifting with learnable $\alpha_1, \alpha_2$.
• KV shifting (both): The full method with all four parameters learnable.

The results (Figure 8b) show that removing either K shift or V shift degrades performance relative to the full method, confirming that both operations contribute. The paper reasons:

"It can be inferred that obtaining the value of the $(i-1)$-th token or the value of the $(i+1)$-th token by focusing on the key of the $i$-th token is important in language modeling. If K shifting or V shifting is not used, the model needs to use two layers of attention to indirectly implement this operation."

Longer shift windows (Figure 8c). The paper extends the shift from the single previous token to multiple previous tokens. For "KV shifting 2" (3-token window):

• $\hat{K}_t = \alpha_1 K_t + \alpha_2 K_{t-1} + \alpha_3 K_{t-2}$
• $\hat{V}_t = \beta_1 V_t + \beta_2 V_{t-1} + \beta_3 V_{t-2}$
• $\alpha_1, \alpha_2$ are initialized from $U(0, 1)$, and $\alpha_3 = 1 - \alpha_1 - \alpha_2$. Similarly for $\beta$.

Similarly for "KV shifting 3" (4-token window with four parameters). Figure 8c shows that longer shift windows do not improve performance despite increasing the number of parameters and computation. This validates the theoretical motivation:

"From the perspective of induction heads, using our current shifting window size is sufficient. If someone want to expand to longer windows, it may need to carefully design it. In this paper, we will no longer attempt more refined designs, as current lightweight designs are sufficient for better learning of induction heads and improving language modeling."

The single-position shift is exactly what induction heads need to match a token to its immediate predecessor and retrieve its immediate successor. Adding more positions dilutes the inductive bias without providing a commensurate representational benefit.

QKV shifting (Appendix J). An additional variant applies the same shifting operation to queries as well: $\hat{Q}_t = \gamma_1 Q_t + \gamma_2 \text{Shift}(Q)_t$ with learnable $\gamma_1, \gamma_2$. The results (Figure 10) on the 19B model show that QKV shifting performs worse than KV shifting and even worse than vanilla on some benchmarks (e.g., MMLU, CMMLU). The paper's interpretation:

"From the perspective of induction heads, the shifting of Q is difficult to contribute to the formation of the induction heads mechanism."

This makes sense: the induction head needs to match the current token's query to the previous token's key. Shifting the query as well would mean the query also represents a blend of the current and previous token, which confuses the matching signal—it's no longer a clean comparison of current-token-to-previous-token.

Learned parameter statistics (Table 4). The paper analyzes the trained 2.9B model (500B tokens) to see what values the $\alpha$ and $\beta$ parameters converged to. For each of the 128 KV pairs (32 layers × 4 KV heads), they check whether $\alpha_1 \leq \alpha_2$ and whether $\beta_1 \leq \beta_2$:

The diagonal entries dominate: 50 heads have $\alpha_1 > \alpha_2$ AND $\beta_1 > \beta_2$ (both favoring the current token over the previous token—standard attention-like behavior), and 52 heads have $\alpha_1 \leq \alpha_2$ AND $\beta_1 \leq \beta_2$ (both favoring the shifted, previous-token component—induction-like behavior). Only 26 heads (17 + 9) have mixed preferences (one parameter favors current, the other favors shifted).

The paper interprets the asymmetry:

"The heads in the upper right focus on the key of the $(i-1)$-th token to obtain information about the $i$-th token, while the heads in the lower left focus on the key of the $i$-th token to obtain information about the $(i-1)$-th token. And they are not symmetrical. We speculate that this is because the model can easily obtain information about the $(i-1)$-th token by interacting with $i$-th token under the causal mask, as the $i$-th token may contain some information about the $(i-1)$-th token to some extent."

In other words, the $(\alpha_1 \leq \alpha_2, \beta_1 > \beta_2)$ configuration (lower-left quadrant, 17 heads) is less common than the $(\alpha_1 > \alpha_2, \beta_1 \leq \beta_2)$ configuration (upper-right quadrant, 9 heads) because attending to the current token's key to get the previous token's value ($\beta_1 \leq \beta_2$, so previous dominates value) is partially redundant with what the model can already do through the hidden-state information flow under the causal mask. The complementary operation—attending to the previous token's key to get the current token's value (upper-right, 9 heads)—is harder to achieve without the explicit shift.

This distribution of learned parameters provides empirical validation that the model is actually using the shift mechanism in the way the theory predicts: about half the heads learn induction-like configurations, while the other half retain standard attention-like configurations, with only a small fraction occupying mixed regimes.

---

3.4.7 Controlled Task Experiments (Toy Models)

Beyond the theoretical analysis, the paper validates the mechanism on several controlled tasks that isolate specific capabilities.

Induction data generation. A custom data generator (Appendix D) creates sequences that test pure induction: randomly sample tokens from a large vocabulary (8,000 tokens), and whenever a token is repeated, ensure that the token following it matches the token that followed its previous occurrence. For example, if token $A$ appeared at position $i$ followed by token $B$, then whenever $A$ appears again at position $j$, the generator places $B$ at position $j+1$. The model's accuracy is measured only at these induction points—positions where the model must predict the successor of a repeated token. Sequences shorter than 512 are padded with zeros.

Architecture for toy tasks. The toy models use a Llama-like architecture (Touvron et al., 2023) with approximately 20M non-embedding parameters. The paper notes this is "slightly different from the experiments in Elhage et al. (2021), which use the attention-only structure that is better alignment with theory. But in order to better fit the actual scenarios of large language models, we chose the llama architecture."

Various depth experiment (Figure 1a). Models with 1 layer (vanilla), 2 layers (vanilla), 4 layers (vanilla), and 1 layer (KV shifting) are compared. Hidden size is 1024. Results:

• 1-layer vanilla: Fails to learn induction entirely (accuracy remains near zero). This confirms the theoretical result from Sanford et al. (2024a) that single-layer standard transformers cannot implement induction heads.
• 2-layer vanilla: Perfectly learns induction to 100% accuracy, but takes many training steps to converge.
• 4-layer vanilla: Converges at roughly the same speed as 2-layer vanilla—"increasing the model depth to layer 4 for Vanilla does not make the model learn induction faster." This suggests the bottleneck is not just depth per se, but the specific difficulty of forming the two-layer induction circuit through gradient descent.
• 1-layer KV shifting: Perfectly learns induction to 100% accuracy, and converges "much faster" (roughly 5-10× fewer steps) than the 2-layer vanilla model. This validates both the representational Theorem 2 (one layer can do it) and the learning dynamics Theorem 3 (easier loss landscape).

Various width experiment (Figure 1b). The hidden size is reduced to 8, and a 2-layer vanilla model is compared against a 1-layer KV shifting model. Results:

• 2-layer vanilla: "The learning ability of standard attention is very poor, and even failed to cover up one of the answers mentioned in the previous text." The model essentially cannot learn induction at this width.
• 1-layer KV shifting: Successfully learns induction (though perhaps not perfectly—the figure shows it reaches significantly higher accuracy than vanilla).

This validates the width reduction claimed by Theorems 1 and 2: standard attention needs $D = 2d$ to approximate induction, while KV shifting needs only $D = d$. At hidden size 8, the effective dimension per token is too small for the vanilla two-layer circuit to represent both the current token's identity and the copied previous token's information, while KV shifting avoids this double-representation requirement.

The paper connects this to real-world language modeling:

"Although this is a toy task, people may think that the current model has a large dimension and can do the induction task well. But induction may also be done in some implicit way in language modeling. This limitation in width will result in the model considering a limited number of different implicit inductions in parallel, or introducing noise in superposition."

In other words, even at realistic hidden sizes (thousands), the width requirement means that the model's capacity for induction competes with its capacity for other representations. KV shifting reduces this competition.

n-gram learning experiment (Figure 3). To test whether KV shifting attention's bias toward induction comes at the cost of other capabilities, the paper evaluates n-gram learning. The task: randomly generate approximately 200 pairs of tokens $(x_1, x_2)$, then randomly generate a third token $x_3$ for each pair. The model must predict $x_3$ when seeing $x_1$ and $x_2$—a 3-gram completion task.

Results for models at three scales (50M parameters with 4 layers, 0.4M parameters with 2 layers, 0.8K parameters with 1 layer):

• KV shifting attention does not enhance the model's ability to learn 3-grams—the accuracy curves for vanilla and KV shifting are essentially identical.
• Importantly, KV shifting also does not weaken 3-gram learning—the bias toward induction is not harmful for other types of pattern recognition.

The paper explains:

"The motivation of our KV shifting attention is to reduce the width and depth required for induction, we cannot expect KV shifting attention to greatly improve the memory ability of the model. In addition, n-gram tasks or some Markov data can actually be simulated with just one layer transformer with vanilla attention (Rajaraman et al., 2024)."

The contrast between Figure 1a (induction, where KV shifting dramatically outperforms) and Figure 3 (n-gram, where it's identical) is a key validation: the modification specifically targets the induction head circuit without affecting the model's ability to learn other patterns.

Multi-hop induction (Appendix H, Figure 9). Following Sanford et al. (2024b), the paper tests multi-hop induction—what the authors call "a multi-layered form of induction heads." In this task, the model must chain multiple induction steps: $A \to B \to C \to ...$. The experimental setup is taken directly from Sanford et al. (2024b) with only the attention mechanism replaced by KV shifting. Error rates are measured for different hop counts ($k$) and sequence lengths ($L$).

Results (Figure 9):

• Vanilla (originally reported by Sanford et al., 2024b): For $L=5$ (pink line), the error rate starts increasing significantly when the hop count exceeds 8, and reaches high error by hop 16.
• Vanilla (reproduced by the authors): Similar trend, confirming the baseline.
• KV shifting attention: For $L=5$ (pink line), the error rate remains small even at hop count 16—dramatically better than vanilla.

This suggests that KV shifting attention's benefits compound for tasks requiring multiple steps of inductive reasoning, which has implications for mathematical and reasoning capabilities.

Grade-school math (Appendix I, Table 7). The paper tests on the iGSM dataset (Ye et al., 2024b), which is a synthetic math dataset designed to isolate reasoning ability from linguistic complexity and data contamination. The experiments use:

• Context length: 1024 for all experiments (slightly different from Ye et al., 2024b).
• Learning rate: $2 \times 10^{-4}$.

Results:

KV shifting attention achieves higher accuracy on both settings, with particularly strong gains on the "Train 12, Test 15" setting (requiring generalization to longer reasoning chains than seen during training). The paper notes: "We expect KV shifting attention to enhance the reasoning ability of the model by improving its ability in basic induction heads."

The paper acknowledges this experiment's limitations: "We only tested the accuracy without delving deeper into the analysis, such as the reasons for mistake like Ye et al. (2024b)."

---

3.4.8 Summary of Design Choices and Their Justifications

• Single-position shift (not 2, not 3): Motivated directly by the induction head mechanism, which requires matching a token to its immediate predecessor and retrieving its immediate successor. Ablation (Figure 8c) confirms longer windows don't help.
• Shift both K and V (not just one): Ablation (Figure 8b) shows both are necessary. K shift enables matching to the neighbor's identity; V shift enables retrieving the neighbor's value.
• Learned interpolation weights (not fixed): Gives each head the flexibility to specialize—some become induction heads (small $\alpha_1$, large $\beta_2$), others remain standard attention heads (large $\alpha_1$, large $\beta_1$), and a few explore mixed configurations for other purposes. Table 4 confirms this specialization occurs in practice.
• Simplex initialization (weight sum = 1): Ensures the initial state doesn't blow up key/value magnitudes, while random sampling from $U(0,1)$ ensures some heads start near the induction-friendly region. The paper notes that "heads are closer to (0, 1) during initialization, making it much easier for them to learn induction heads."
• No constraint during training (no gating, no clipping): Figure 8a shows that allowing weights outside $[0, 1]$ and allowing sums to deviate from 1 enables richer features (e.g., anti-copying heads with negative $\beta_2$). The model naturally learns to keep most parameters in reasonable ranges without enforcement.
• Applied before RoPE and before flash attention: Ensures compatibility with existing training and inference frameworks. The convolution-based implementation is efficient and the inference caching strategy avoids recomputation.
• Per KV head (not per query head) in GQA: Minimizes additional parameter count while applying the shift at the most natural level—the point where key and value representations are formed.
• Large RoPE base (100,000): Informed by concurrent work showing the default 10,000 is inadequate for the context lengths used. This is a practical choice specific to the production-oriented models, not a fundamental requirement of KV shifting.

4. Key Insights and Innovations

Innovation 1: Architectural Modification as a Direct Consequence of Mechanistic Understanding

The paper's most distinctive intellectual contribution is not the modification itself—shifting keys and values by one position is a simple operation—but the methodology by which the modification was derived. The field has a well-established pipeline for mechanistic interpretability: identify circuits in trained models, characterize their function, and use that understanding to explain model behavior. What the field has largely not done is close the loop by using that understanding to redesign the architecture itself. This paper does exactly that.

The chain of reasoning is unusually direct for an architectural modification paper: (1) induction heads are important for ICL, (2) standard attention requires two layers and excess width to implement them because keys and values are bound to the same token, (3) if we decouple keys and values by allowing token i's key to access token i-1's value, a single layer suffices. The resulting modification—KV shifting attention—is not an arbitrary architectural exploration but a targeted intervention that addresses a specific, diagnosed inefficiency.

This is distinct from how architectural modifications are typically proposed. Most new attention mechanisms (differential attention, selective attention, linear attention, etc.) are motivated by general desiderata—reducing noise, improving sparsity, lowering complexity—rather than by a specific circuit-level bottleneck. The paper makes this contrast explicit in Section 7:

"There have been some efforts to modify transformers to enhance modeling capabilities, such as reducing the noise of attention or reducing attention to unneeded elements. Our work is to slightly modify the attention to enhance its ability to learn induction heads, which potentially improves language modeling."

The significance of this methodology extends beyond the specific modification. It demonstrates that mechanistic interpretability can be generative rather than merely descriptive—it can produce actionable architectural improvements, not just explanations. This suggests a research program where other identified bottlenecks (e.g., the difficulty of learning parity, or the formation of function vectors) motivate similarly targeted modifications. The paper explicitly gestures toward this in the Discussion:

"If we are not ready to completely overturn the Transformer architecture, it may be interesting to delve into its underlying mechanisms, identify important features, and modify the Transformer to make it easier to implement certain important functions. For example, by modifying transformer to make it easy to learn parity check, which might be very challenging."

This is a conceptual reframing of the relationship between interpretability and architecture design—from parallel tracks to a closed loop—and it represents a fundamental shift in how architectural improvements can be derived.

The evidence that this methodology succeeds is the tight correspondence between theoretical prediction and empirical outcome. Theorem 2 predicts that one-layer KV shifting can exactly represent induction heads; Figure 1a confirms that one-layer KV shifting learns induction while one-layer vanilla cannot. Theorem 2 predicts reduced width requirements; Figure 1b confirms that KV shifting succeeds at hidden size 8 where two-layer vanilla fails. The theoretical analysis (Section 3.1) and the toy experiments (Section 3.2) are not separate contributions—they are mutually reinforcing validations of the core methodological claim that understanding a circuit's mechanical requirements enables designing an architecture that satisfies those requirements more efficiently.

---

Innovation 2: The Width Requirement as a Previously Unrecognized Bottleneck for Induction Heads

While the depth requirement for induction heads (needing at least two layers) has been formally established by prior work (Sanford et al., 2024a), the paper identifies and formalizes a second, equally important requirement: width. This is not just an incremental addition to the analysis—it is a conceptually distinct bottleneck with different implications for model design.

The width requirement arises from the specific mechanics of how a two-layer standard transformer implements induction through virtual attention heads. As the paper's Property 1 shows, the effective attention weight (A_{h2} A_{h1})_{j, i+1} can only sum over indices k ≥ i+1, which means token i's hidden state must already carry information about token i+1 for the induction to work. But under the causal mask, token i cannot attend to token i+1 to acquire that information. The workaround—having token i+1 attend to token i and copy its information forward—means that each hidden state must simultaneously represent its own token's identity (for the matching step) AND information about its predecessor (for the routing step). This dual representation imposes a minimum width: D = 2d in the constructive proof of Theorem 1.

The paper's key diagnostic insight is that this width requirement is not just a quirk of the constructive proof—it represents a genuine representational bottleneck that affects real models. The toy experiment at hidden size 8 (Figure 1b) provides striking evidence: the two-layer vanilla model "failed to cover up one of the answers mentioned in the previous text," while the one-layer KV shifting model succeeds. The paper connects this to practical language modeling with a subtle but important argument:

"Although this is a toy task, people may think that the current model has a large dimension and can do the induction task well. But induction may also be done in some implicit way in language modeling. This limitation in width will result in the model considering a limited number of different implicit inductions in parallel, or introducing noise in superposition."

This is a qualitatively different claim from the depth argument. Depth is about serial computation—how many sequential operations are needed. Width is about parallel capacity—how many distinct induction patterns can be simultaneously represented without interference. Even at realistic hidden sizes (thousands of dimensions), the width requirement means that induction competes with other representations for limited representational space. KV shifting attention eliminates this competition by halving the effective width needed per induction operation.

The theoretical framing through Theorems 1 and 2 makes this precise: standard attention needs D = 2d for approximation with error O(e^{-p^{(1)}}), while KV shifting achieves exact representation at D = d. The factor-of-2 reduction in required width is a fundamental improvement in representational efficiency, not a marginal optimization. The paper acknowledges that proving this as a lower bound (showing standard attention cannot do better) is left to future work, but the constructive separation combined with the toy experiment at hidden size 8 makes the practical relevance clear.

---

Innovation 3: The Loss Landscape for Induction Is Fundamentally Simpler Under KV Shifting

Beyond the representational question (can the architecture express induction heads?), the paper provides a learning dynamics analysis (Theorem 3) that reveals why KV shifting attention converges faster. This is not merely an observation that convergence is faster in practice—it is a theoretical characterization of why the optimization problem becomes easier.

Under the simplified conditions of Theorem 3, learning induction heads in KV shifting attention reduces to optimizing four scalar parameters (α₁, α₂, β₁, β₂) per head, with a loss landscape that has a clear structure: the optimum for induction is at (α₁, β₁) = (0, 1). The gradient descent trajectories (Figure 2) show that depending on the "distractor" term O(T), the landscape ranges from tightly curved (fast convergence, small vocabulary) to broadly flat (slow convergence, large vocabulary), but the gradient direction is consistently oriented toward the optimum. The key insight from the paper:

"In KV shifting attention, induction heads become very easy to learn, and even with good initialization, the model has induction capability."

This stands in stark contrast to the standard transformer, where induction must be learned through a two-phase process (Bietti et al., 2024): first learn bigram statistics, then form induction heads through a top-down mechanism that overrides those statistics. This is a complex, sequential learning process that requires coordinating two layers of attention weights. KV shifting collapses this into a single-phase optimization over four scalars, with the added benefit that random initialization naturally places some heads near the induction-friendly region (since α₁ and β₁ are initialized uniformly from [0, 1], some heads will start with α₁ ≈ 0, β₁ ≈ 1).

The practical consequence—faster convergence—is validated at scale: the training loss curves (Figure 4) show KV shifting maintaining a consistent advantage throughout pretraining at both 2.9B and 19B scales, not just in early stages. This is significant because it suggests the benefit is not merely an initialization effect (which would diminish over training) but a sustained optimization advantage from the simpler loss landscape.

The paper also notes an interesting trade-off: the bias toward induction means that KV shifting attention may have worse bigram capability at initialization—"it is difficult to have appropriate initialization to obtain a certain level of bigrams capability without training." But this is a favorable trade-off because bigram statistics are easy to learn from data, while induction head formation is structurally difficult. By making the hard thing easy and the easy thing slightly harder (but still learnable), KV shifting attention achieves a better allocation of the model's learning capacity.

This learning dynamics analysis is a distinct intellectual contribution from the representational analysis. Representation tells you what the architecture can do; learning dynamics tells you how quickly and reliably gradient descent will find that configuration. The paper provides both, and together they explain not just that KV shifting works, but why the convergence speedup (Figure 1a: roughly 5-10× faster on toy induction) and the sustained loss advantage (Figure 4) are to be expected.

---

Innovation 4: Robustness as Evidence of a Flatter Optimization Landscape

The paper reports a striking empirical finding that transcends the induction-specific motivation: KV shifting attention tolerates learning rates that cause vanilla attention to diverge (Figure 5c: at LR = 1e-2, vanilla diverges while KV shifting converges). This result, confirmed across 5 random seeds, is not directly predicted by the induction head theory and represents a serendipitous discovery with practical implications.

The paper's interpretation—"the optimization space for KV shifting attention may be flatter"—aligns with broader understanding that architectural choices can affect loss landscape geometry. But the specific mechanism by which KV shifting achieves this flatter landscape is interesting. During initialization, the shifted K and V are convex combinations of adjacent tokens (since α₁ + α₂ = 1 and β₁ + β₂ = 1 initially). This means the input to the attention softmax has some built-in smoothness: adjacent positions have correlated keys and values, which may prevent the extreme attention concentration that leads to gradient spikes at high learning rates. As training progresses and the α and β parameters deviate from their initial simplex, the model has already settled into a stable basin, making it less susceptible to divergence.

This robustness property is practically important for two reasons. First, it means KV shifting attention can be trained with larger learning rates, which prior work (Lobacheva et al.) has shown improves generalization. The paper explicitly adopts large learning rates for this reason (8 × 10^{-4} for the 2.9B model), and the robustness finding suggests KV shifting makes this choice safer. Second, it reduces the need for careful learning rate tuning—a significant practical concern for large-scale pretraining where hyperparameter sweeps are expensive.

The robustness also partially explains the persistent training loss advantage (Figure 4): if KV shifting attention can tolerate more aggressive optimization, it can make faster progress per step throughout training, not just in the early phases. This is a non-obvious practical benefit that emerged from the architectural modification and is conceptually distinct from the induction-specific representational advantage.

---

Innovation 5: The Trade-off Between Induction Bias and Memory Capacity Is Empirically Benign

One might reasonably worry that an architectural modification biased toward induction heads would impair the model's ability to learn other types of patterns—particularly n-gram statistics, which form the backbone of language modeling. The paper directly tests this concern and produces an empirically significant negative result: KV shifting attention is neutral for n-gram learning (Figure 3). It neither helps nor hurts.

This is not a trivial finding. Many architectural modifications that improve one capability come at the cost of another—this is the ubiquitous trade-off in machine learning. The fact that KV shifting improves induction dramatically (Figure 1a) while leaving n-gram learning unchanged (Figure 3) means the modification is targeted rather than globally disruptive. The paper's explanation is theoretically grounded: n-gram learning on Markov data can be done with a single layer of standard attention (Rajaraman et al., 2024), so there is no structural bottleneck for KV shifting to alleviate. The modification inserts itself precisely where standard attention is deficient (induction) and remains transparent where standard attention is already sufficient (n-grams).

This targeted improvement has an important implication for the scaling experiments: the gains observed in full language model pretraining (Table 2, Figure 4) are not coming at the cost of some other, unmeasured capability. The model is genuinely better at induction without being worse at memorization. The paper's hop-k experiments (Appendix H, Figure 9) and iGSM math experiments (Appendix I, Table 7) further support this: the benefits extend to multi-step reasoning tasks that build on induction, suggesting a compounding advantage where better primitive induction capability enables better higher-level reasoning without sacrificing basic pattern recognition.

The contrast between Figure 1a (induction: KV shifting dramatically better) and Figure 3 (n-gram: identical) is the cleanest demonstration of this targeted improvement. It establishes that the modification is not a blunt instrument that makes the model "better at attention" in some vague sense—it is a surgical intervention that addresses a specific, diagnosed bottleneck in how standard transformers form one particular circuit, leaving other capabilities intact.

5. Experimental Analysis

Evaluation Methodology

• Dataset. All large-scale pretraining experiments use non-public private data with collection and filtering methods "similar to FineWeb-edu" (Penedo et al., 2024). The authors state they will use open-source data like RedPajama-1T when compute resources permit. For evaluation, eight benchmark datasets are used: LAMBADA (Paperno et al., 2016), Winogrande (Sakaguchi et al., 2021), HellaSwag (Zellers et al., 2019), ARC-Easy and ARC-Challenge (Clark et al., 2018), CMMLU (Li et al., 2023a), MMLU (Hendrycks et al.), and MATH (Hendrycks et al., 2021). Toy experiments use synthetically generated data—induction sequences (Appendix D), n-gram patterns, hop-k tasks from Sanford et al. (2024b), and iGSM math data from Ye et al. (2024b).

• Base model(s). All models follow a Llama2-like (Touvron et al., 2023) architecture. The primary pretraining comparisons use models at 2.9B parameters (32 layers, hidden size 2,560, 20 query heads, 4 KV heads via GQA, 500B training tokens) and 19B parameters (48 layers, hidden size 6,144, 48 query heads, 4 KV heads via GQA, 200B training tokens). Additional scaling experiments cover 1.5B, 6.7B, and 13B parameter configurations trained on 10B tokens each (Table 6). The 2.9B scale is chosen as a production-representative size; the 19B scale validates that benefits persist at larger scale. Toy induction experiments use ~20M non-embedding parameter models with Llama architecture rather than attention-only to "better fit the actual scenarios of large language models" (Section 3.2).

• Metrics. The primary metric is next-token prediction accuracy on induction tasks (fraction of induction positions where the model correctly predicts the successor of a repeated token). For full pretraining, training loss (cross-entropy) is the primary optimization signal, and benchmark accuracy (percentage of correctly answered questions) is reported across eight benchmarks. Validation loss on WebText (Radford et al.) is used for scaling law plots. For MMLU specifically, three evaluation protocols are compared (Table 3): cloze (zero-shot, following Waleffe et al., 2024, designed to measure knowledge without multiple-choice format), zero-shot multiple choice, and 5-shot multiple choice. For iGSM math, accuracy is measured as "as long as some step calculates the question asked and answers the correct answer, it is considered to be answered correctly" (Appendix I). For hop-k, error rate is reported (lower is better).

• Baselines. The primary baseline is a vanilla transformer with identical architecture (same depth, width, head count, training data, and hyperparameters) except using standard attention without KV shifting. This is a same-parameter-count comparison—the 1-layer KV shifting model has the same total parameters as the 1-layer vanilla model (since the added α and β parameters are negligible), and the 2-layer vanilla model has approximately 2× the parameters of the 1-layer KV shifting model in the toy experiments. For the QKV shifting variant (Appendix J), the baseline is the KV shifting model itself, with additional Q-shifting parameters.

• Generation budget / compute accounting. For toy tasks, no explicit compute budget is reported—models are trained for a fixed number of steps and convergence speed is compared (steps to reach a given accuracy). For pretraining, total training tokens serve as the compute proxy: 500B tokens for 2.9B models, 200B for 19B models, and 10B for the 1.5B/6.7B/13B scaling experiments. The parameter overhead of KV shifting is 4 learnable scalars per KV head per layer (e.g., 4 × 4 × 32 = 512 parameters total for the 2.9B model), with an additional computation cost of O(ND) for the shift operation, which the paper notes is "much smaller than O(ND² + N²D)" for the main attention (Section 2.2). No FLOPs-matched comparison between architectures is performed—the evaluation assumes equal training tokens and equal parameter counts (up to the negligible ~0.00002% increase from the shift parameters).

• Cross-validation / statistical protocol. For robustness experiments with the 1.5B model, five random seeds are used for model initialization and data sampling, with vanilla and KV shifting sharing the same seed per experiment (Figure 5a). At LR = 1e-2, the non-divergence result for KV shifting is confirmed across all 5 seeds while vanilla diverges in all 5 (Section 4.4). For the main 2.9B and 19B pretraining runs, "almost experiments are only run once except there is an additional notion" (Section 4.1), meaning most large-scale results are single-run. Standard deviations or confidence intervals are not reported for benchmark scores.

Main Quantitative Results

Toy Induction: Depth and Width Requirements

The foundational experimental claim is that KV shifting attention eliminates the depth and width requirements that standard attention imposes on induction head formation. Figure 1a (various depth, hidden size 1024) demonstrates this across model configurations:

• 1-layer vanilla: Fails to learn induction entirely—accuracy remains at chance levels. This empirically confirms Sanford et al. (2024a)'s proof that single-layer standard transformers cannot implement induction heads.
• 2-layer vanilla: Achieves ~100% induction accuracy but converges slowly, requiring many training steps.
• 4-layer vanilla: Converges at approximately the same speed as 2-layer—"increasing the model depth to layer 4 for Vanilla does not make the model learn induction faster" (Section 3.2).
• 1-layer KV shifting: Achieves ~100% induction accuracy and converges "much faster" than the 2-layer vanilla model, reaching high accuracy in a fraction of the training steps. The paper describes this as "roughly 5-10× faster" convergence, though exact step counts are not tabulated.

The critical detail here is the parameter count asymmetry: the 1-layer KV shifting model has the same number of parameters as the 1-layer vanilla (since the α and β scalars add negligible parameters), while the 2-layer vanilla has approximately 2× the parameters. So KV shifting achieves better performance with fewer parameters.

Figure 1b (various width, hidden size 8) stress-tests the width requirement:

• 2-layer vanilla with hidden size 8: "The learning ability of standard attention is very poor, and even failed to cover up one of the answers mentioned in the previous text." The model essentially cannot learn induction at this width.
• 1-layer KV shifting with hidden size 8: Successfully learns induction, achieving significantly higher accuracy than vanilla despite having fewer total parameters (1 layer vs. 2 layers) and the same hidden dimension.

This validates the theoretical claim from Theorems 1-2: standard attention needs D = 2d to approximate induction (with error), while KV shifting represents it exactly at D = d. The hidden size 8 experiment demonstrates this separation empirically—when the dimension is too small for the vanilla two-layer circuit to simultaneously encode token identity and copied predecessor information, KV shifting succeeds because it avoids this dual-representation requirement.

n-gram Learning: No Improvement, No Degradation

Figure 3 tests whether the induction bias impairs other learning capabilities. The task requires predicting x₃ given x₁ and x₂ from ~200 randomly generated 3-gram pairs, tested at three model scales:

• 50M parameters (4 layers): Vanilla and KV shifting produce nearly identical accuracy curves.
• 0.4M parameters (2 layers): Nearly identical.
• 0.8K parameters (1 layer): Nearly identical.

KV shifting attention neither enhances nor impairs 3-gram learning at any scale. This is a critical negative result: the architectural modification is targeted—it improves induction while leaving memorization capacity unchanged. The contrast with Figure 1a is stark and intentional: induction benefits dramatically, n-gram learning is unchanged, establishing that KV shifting is not a blunt instrument but a surgical improvement to a specific circuit.

Multi-hop Induction (Appendix H, Figure 9)

Following the experimental protocol of Sanford et al. (2024b), the paper evaluates multi-hop induction—requiring the model to chain multiple induction steps. Figure 9 shows error rates as a function of hop count (k) for different sequence lengths (L):

• Vanilla (original, Figure 9a): For L=5 (pink line), error rate begins increasing significantly when hop count exceeds ~8, reaching high error by hop 16. The pattern is consistent across sequence lengths. This was the baseline established by Sanford et al. (2024b).
• Vanilla (reproduced, Figure 9b): Qualitatively similar, confirming the baseline.
• KV shifting attention (Figure 9c): For L=5 (pink line), error remains small even at hop 16—dramatically better. Across all sequence lengths, KV shifting maintains lower error rates, with the advantage widening at higher hop counts.

The paper flags this as evidence that KV shifting's benefits compound for sequential reasoning: "This powerful ability to perform implicit reasoning implies that KV shifting attention may achieve better results in mathematical or reasoning abilities." Note that error rates are the metric—lower is better—and the visual separation between vanilla and KV shifting curves is qualitative (error rate values are not tabulated in the figure, only plotted).

Grade-School Math: iGSM (Appendix I, Table 7)

Table 7 reports accuracy on the iGSM synthetic math dataset (Ye et al., 2024b), which isolates reasoning from linguistic complexity:

KV shifting achieves +7.5 percentage points higher accuracy in the harder generalization setting (Train 12, Test 15) and +3.5 points in the in-distribution setting (Train 21, Test 24). The paper notes that the experiment uses a slightly different setup from Ye et al. (2024b)—context length 1024 for all experiments and learning rate 2e-4—and that the evaluation is lenient: "as long as some step calculates the question asked and answers the correct answer, it is considered to be answered correctly, even if additional calculation steps are performed after answering the correct answer." The authors acknowledge this experiment is not as thorough as Ye et al. (2024b), having tested only accuracy without analyzing error types.

Large-Scale Pretraining: Training Loss

Figure 4 presents training loss curves for the primary pretraining runs:

• 2.9B model (Figure 4a, 500B tokens): KV shifting attention maintains a consistently lower training loss than vanilla throughout the entire training run. The gap is visible from early in training and persists without closing up to 500B tokens. The learning rate is 8e-4—a large learning rate that the paper notes "KV shifting attention leads more in terms of loss, because the 2.9B model is trained with a large learning rate" (Section 4.4), consistent with the robustness finding that KV shifting tolerates aggressive optimization better.
• 19B model (Figure 4b, 200B tokens): Similar pattern—KV shifting maintains lower training loss throughout, though the gap appears somewhat smaller than at 2.9B scale (specific loss values are plotted but not tabulated).

Figures 6a-c extend this to additional scales (1.5B, 6.7B, 13B parameters, all trained on 10B tokens each):

• At all three scales, KV shifting attention achieves lower training loss than vanilla. The gap is visible in all curves, though the separation varies by scale.
• Note: the 13B model is trained with a batch size of 1M tokens (vs. 0.5M for 1.5B and 6.7B), so its total training steps are half—this affects the x-axis but not the relative comparison at each scale.

Figure 7 shows validation loss scaling:

• Figure 7a (Scaling law): Validation loss on WebText is plotted against non-embedding parameters (following Kaplan et al., 2020). For models trained on 10B tokens, KV shifting attention achieves lower validation loss than vanilla at every parameter scale tested (1.5B, 6.7B, 13B). The curves show a consistent offset favoring KV shifting, suggesting the benefit scales with model size.
• Figure 7b (1.5B extended training): The 1.5B model is trained to 30B tokens with validation loss measured every 1000 steps. The gap between KV shifting and vanilla does not diminish over this extended training period—if anything, it appears stable or slightly widening, though validation loss differences are small at this scale.

Large-Scale Pretraining: Benchmark Performance

Table 2 reports benchmark accuracy for the main pretraining runs at three token counts per scale. For the 2.9B models (500B total tokens):

At 340B tokens:

• KV shifting achieves 36.46% average vs. vanilla's 32.48% (+3.98 percentage points).
• Largest gains: ARC-E (+9.29 points: 36.74% vs. 27.45%), MMLU (+6.77: 36.20% vs. 29.43%).
• Notably, HellaSwag is nearly identical (42.87% vs. 42.70%).

At 420B tokens:

• KV shifting achieves 36.55% average vs. vanilla's 34.08% (+2.47 points).
• The gap narrows as vanilla catches up—ARC-E gap shrinks to +7.70 points, MMLU gap to +3.96 points.
• LAMBADA actually reverses: vanilla 52.80% vs. KV shifting 51.91%.

At 500B tokens (final checkpoint):

• KV shifting achieves 38.57% average vs. vanilla's 36.45% (+2.12 points).
• Largest final gaps: MMLU (+3.62: 40.88% vs. 37.26%), CMMLU (+2.56: 40.78% vs. 38.22%), ARC-E (+2.82: 39.02% vs. 36.20%).
• Small gaps: HellaSwag (44.52% vs. 44.49%, +0.03), Winogrande (55.33% vs. 54.06%, +1.27).
• MATH remains low for both but KV shifting leads: 2.60% vs. 1.80%.

For the 19B models (200B total tokens):

At 160B tokens:

• KV shifting achieves 37.74% average vs. vanilla's 36.43% (+1.31 points).
• CMMLU shows the largest gap: 42.10% vs. 39.12% (+2.98 points), MMLU: 42.87% vs. 39.22% (+3.65 points).
• ARC-C shows a small reversal: vanilla 24.56% vs. KV shifting 25.06%, essentially tied.

At 180B tokens:

• KV shifting achieves 37.74% average vs. vanilla's 36.83% (+0.91 points).
• Gap narrows further—CMMLU: +2.58 points, MMLU: +1.15 points.

At 200B tokens (final):

• KV shifting achieves 38.83% average vs. vanilla's 38.06% (+0.77 points).
• ARC-C shows the largest relative gain: 29.32% vs. 25.78% (+3.54 points).
• Several benchmarks show small gaps: HellaSwag (48.42% vs. 47.36%, +1.06), MMLU (43.29% vs. 42.68%, +0.61).
• LAMBADA: 62.35% vs. 60.88% (+1.47), MATH: 3.20% vs. 2.60% (+0.60).

Key pattern across both scales: KV shifting leads at every checkpoint, but the advantage narrows as training progresses. At 2.9B, the gap shrinks from +3.98 points (340B) to +2.12 points (500B); at 19B, from +1.31 points (160B) to +0.77 points (200B). This suggests the bias toward induction provides an early advantage that partially—but not completely—diminishes over longer training. The paper argues: "KV shifting attention can achieve better performance than the vanilla model when they both converge, although it may take several TB data for a 2.9B model to converge" (Section 4.2), implying that with sufficient training, both might converge to similar performance, but practical training budgets favor KV shifting.

Table 3 breaks down MMLU for the 2.9B model at 500B tokens by evaluation protocol:

KV shifting outperforms vanilla across all three protocols: +1.76 points (cloze), +3.99 points (zero-shot), +3.62 points (few-shot). The few-shot gap (+3.62) is larger than the cloze gap (+1.76), which the paper interprets as evidence that "the model can easily use context to compare the possibilities of various options and select the option with the highest probability"—the induction capability directly aids in-context learning, manifesting as a larger few-shot improvement.

Robustness Experiments

Figure 5a (various seeds, 1.5B model, 10B tokens): Across 5 random seeds for model initialization and data sampling, KV shifting attention achieves lower training loss than vanilla in every run. The curves show run-to-run variance (training loss is "quite shaky"), but the ordering is consistent—KV shifting is always below vanilla. No seed produces a vanilla run that outperforms KV shifting.

Figure 5b (various learning rates, 1.5B model, 10B tokens): Across learning rates of 1e-4, 2e-4, 1e-3, and 1e-2, KV shifting attention achieves lower training loss than vanilla at all tested rates. The advantage is visible at each LR.

Figure 5c (LR = 1e-2 divergence): At the extreme learning rate of 1e-2, vanilla attention diverges (loss explodes) while KV shifting attention converges normally. This is confirmed across 5 random seeds—"the all results of each experiment are Vanilla divergence and KV shifting convergence." The paper interprets this as evidence that "the optimization space for KV shifting attention may be flatter" (Section 4.4), potentially because the shifted K and V provide smoothness that prevents the extreme attention concentration leading to gradient spikes.

Ablation Studies and Robustness Checks

• Shifting only K or only V (Figure 8b): Removing either K shifting (setting α₁=1, α₂=0, keeping V shifts learnable) or V shifting (β₁=1, β₂=0, keeping K shifts learnable) degrades training loss relative to full KV shifting on the 1.5B model with 10B tokens. Both ablations perform worse than the full method, confirming both operations contribute. The paper's interpretation: "If K shifting or V shifting is not used, the model needs to use two layers of attention to indirectly implement this operation."

• Longer shift windows (Figure 8c): Extending the shift to a 3-token window (KV shifting 2: K̂_t = α₁K_t + α₂K_{t-1} + α₃K_{t-2}) or 4-token window (KV shifting 3: up to K_{t-3}) does not improve training loss over the 2-token window on the 1.5B model with 10B tokens. The curves largely overlap, with the original KV shifting (window of 2) performing slightly better or identically. This confirms that the single-position shift is sufficient for induction and that longer windows add parameters and computation without benefit.

• Gating mechanisms (Figure 8a, "KV shifting gate"): Enforcing α₁ + α₂ = 1 and β₁ + β₂ = 1 throughout training via sigmoid gating (α₁ = Sigmoid(a), α₂ = 1 - α₁) produces slightly worse training loss than unconstrained KV shifting on the 1.5B model with 10B tokens. The paper interprets this as evidence that "allowing α and β to have a wider range of degrees of freedom may enable the model to learn richer features"—heads can learn anti-copying patterns (β₂ < 0) or other operations that require violating the simplex constraint.

• Clipping to [0,1] (Figure 8a, "KV shifting 0 to 1"): Clamping parameters to [0,1] after each update produces training loss nearly identical to unconstrained KV shifting. The constraint neither helps nor hurts, suggesting the model naturally keeps most parameters in reasonable ranges without enforcement.

• QKV shifting (Appendix J, Figure 10): Extending the shift to queries as well (Q̂_t = γ₁Q_t + γ₂Shift(Q)_t with learnable γ₁, γ₂) evaluated on the 19B model. Figure 10a shows QKV shifting achieves training loss similar to KV shifting initially but diverges slightly higher. Figure 10b-c show benchmark results: QKV shifting performs worse than KV shifting on MMLU and CMMLU, and on CMMLU even underperforms vanilla. The paper's interpretation: "From the perspective of induction heads, the shifting of Q is difficult to contribute to the formation of the induction heads mechanism"—shifting the query confuses the matching between the current token and the previous token's key.

• Learnable parameter convergence (Table 4): Analysis of the trained 2.9B model (500B tokens, 128 KV pairs = 32 layers × 4 KV heads) shows the distribution of learned α and β:

  • 50 heads (39.1%) have α₁ > α₂ AND β₁ > β₂ (both favoring current token—standard attention-like).
  • 52 heads (40.6%) have α₁ ≤ α₂ AND β₁ ≤ β₂ (both favoring shifted token—induction-like).
  • 17 heads (13.3%) have α₁ ≤ α₂ AND β₁ > β₂ (mixed: shift key toward previous, use current value).
  • 9 heads (7.0%) have α₁ > α₂ AND β₁ ≤ β₂ (mixed: use current key, shift value toward previous). The diagonal (concordant α and β preferences) dominates with 102 of 128 heads (79.7%). The diagonal was already 47 and 43 respectively at 20B tokens, suggesting the specialization pattern emerges early. The asymmetry between off-diagonal cells (17 vs. 9) is attributed to causal mask effects: "the model can easily obtain information about the (i-1)th token by interacting with ith token under the causal mask," making the (α₁ ≤ α₂, β₁ > β₂) configuration less necessary.

• Sum of weights deviating from 1: Despite initialization with α₁ + α₂ = 1 and β₁ + β₂ = 1, post-training sums deviate. The paper provides a specific example from the 3rd KV pair of the 17th layer: α₁ = 0.08, α₂ = 0.43 (sum = 0.51), β₁ = 0.34, β₂ = -0.15 (sum = 0.19). Negative values emerge (β₂ = -0.15), which the paper interprets as potentially implementing "anti-copying prefix-search" heads (Elhage et al., 2021). Experiments with enforced simplex constraints (Figure 8a) confirm that allowing this deviation is beneficial.

• Different learning rate schedules (Figure 5b): KV shifting outperforms vanilla at learning rates spanning two orders of magnitude (1e-4 to 1e-2), demonstrating robustness to this hyperparameter. The extreme case at 1e-2, where vanilla diverges, is confirmed across 5 seeds.

• Different random seeds (Figure 5a): KV shifting maintains lower training loss across 5 different random initializations and data orders, ruling out the possibility that the advantage is due to a lucky initialization.

• Validation loss scaling (Figure 7b): The 1.5B model trained to 30B tokens shows that the validation loss gap between KV shifting and vanilla does not close with extended training, suggesting the advantage is not merely a convergence speed effect but represents a sustained performance improvement at this scale.

Critical Assessment

Do the experiments support the claim that KV shifting attention enables single-layer induction head formation? Yes, and this is the strongest empirical result in the paper. Figure 1a cleanly demonstrates the predicted pattern: 1-layer vanilla fails entirely, 1-layer KV shifting succeeds, and 2-layer vanilla succeeds but more slowly. The hidden size 8 experiment (Figure 1b) adds further confirmation. The theoretical framework (Theorems 1-2) provides a constructive proof of exact representability for KV shifting vs. approximation-with-error for vanilla, and the toy experiments validate this in a controlled setting. However, the toy experiments use a specific data distribution (random token sequences with controlled repetitions), an 8,000-token vocabulary, and a simplified Llama architecture. Whether the exact representability advantage (equality in Theorem 2 vs. O(e^{-p(1)}) error in Theorem 1) translates to measurable differences in full-scale language modeling—as opposed to the learning dynamics advantage captured by Theorem 3—is not directly tested; the toy data isolates the induction mechanism but doesn't capture the complexity of natural language.

Do the pretraining experiments support the claim that KV shifting improves language modeling? The evidence is positive but requires qualification. Table 2 shows consistent benchmark improvements at both 2.9B and 19B scales across all checkpoints, and training loss curves (Figures 4, 6) show persistent gaps. The single-run nature of the large-scale experiments is a genuine weakness—for the 2.9B and 19B models, we have N=1 per configuration, so we cannot assess whether the observed +2.12 point advantage at 500B tokens is statistically reliable or within run-to-run variance. The smaller-scale robustness experiment (1.5B, 5 seeds, Figure 5a) shows consistent KV shifting advantage across seeds, but this is at a different scale and training budget (10B tokens vs. 500B). The paper does not report confidence intervals, standard deviations, or statistical tests for any benchmark result in Table 2. Given that some benchmark gaps are small (HellaSwag at 2.9B/500B: +0.03 points), single-run results cannot distinguish signal from noise.

Does the advantage diminish with scale and training duration? This is suggested by the data but not conclusively established. At 2.9B, the average benchmark gap narrows from +3.98 (340B) to +2.47 (420B) to +2.12 (500B). At 19B, it narrows from +1.31 (160B) to +0.91 (180B) to +0.77 (200B). This pattern is consistent with the hypothesis that KV shifting provides an inductive bias benefit that matters most early in training, and that sufficiently long training might close the gap entirely. However, the experiments are terminated before convergence—the 2.9B model at 500B tokens has not plateaued (training loss in Figure 4a is still decreasing), so we don't know the converged gap. The paper's statement that "both 2.9B models are converge on the benchmark Lambda" (Section 4.2) is about a single benchmark, not the full suite. A longer training run to loss plateau would be needed to determine whether the advantage is permanent or transient.

Does the paper demonstrate that the improvements are specifically due to enhanced induction head formation, as opposed to some other effect of the architectural modification? The evidence is circumstantial but coherent. The toy experiments (induction data, hop-k, iGSM math) isolate tasks where induction heads are the primary mechanism and show large gains. The n-gram experiment (Figure 3) shows no gain on a task where induction is irrelevant, establishing that the modification doesn't improve everything blindly. The learned parameter analysis (Table 4) shows ~40% of heads converge to induction-like configurations, consistent with the mechanism being actually used. The ablation showing both K and V shifts are necessary (Figure 8b) aligns with the theoretical requirement for decoupled key-value access. However, none of these directly prove that the benchmark improvements in Table 2 are caused by better induction heads rather than some other consequence of the shift operation—for example, the smoothness induced by adjacent-token interpolation might improve optimization generally, which the LR=1e-2 robustness result (Figure 5c) would support independently of induction.

What experiments are missing that would strengthen the claims? Several gaps are notable. First, no mechanistic interpretability analysis is performed on the trained large-scale models to verify that induction heads have actually formed and that they use the KV shifting mechanism as predicted. The parameter analysis (Table 4) shows what values the α and β parameters took, but doesn't demonstrate that heads with induction-like parameters actually implement induction in practice—that would require activation patching, attention pattern analysis, or similar circuit-level investigation. Second, the benchmark evaluations use private training data (with filtering "similar to FineWeb-edu"), making exact replication impossible and raising questions about whether the reported improvements would transfer to models trained on different data distributions. Third, the paper trains on non-public data but evaluates on public benchmarks—this raises standard contamination concerns. Fourth, no comparison is made against alternative methods for improving induction, such as architectural modifications that increase the model's effective width for induction patterns while keeping total parameters constant. Fifth, the 19B model experiments show minimal benchmark gaps (+0.77 points average at 200B tokens) that would require multi-run statistics to validate. Sixth, the "converged" comparison on LAMBADA is not accompanied by evidence that LAMBADA accuracy has actually plateaued for either model. Seventh, while the paper argues that KV shifting's robustness to high learning rates suggests a flatter optimization landscape, no direct measurement of landscape flatness (e.g., Hessian eigenvalue spectra, loss curvature along interpolations) is provided.

Is the claim that KV shifting "reduces the width requirement" validated outside of toy tasks? Only indirectly. Figure 1b demonstrates the width reduction at the extreme low end (hidden size 8 on synthetic data), which validates the theory in a minimal setting. But no experiment systematically varies hidden size in the pretraining regime to show that vanilla models need larger width to match KV shifting's induction capability. For example, one could compare a vanilla model with hidden size 2D against a KV shifting model with hidden size D to test whether the theoretical factor-of-2 width reduction holds in language modeling. Without such an experiment, the width reduction claim remains primarily a theoretical and toy-task result.

The learning dynamics analysis (Theorem 3) and gradient flow diagrams (Figure 2) are elegant, but do they directly connect to the observed pretraining behavior? The connection is qualitative rather than quantitative. Theorem 3 analyzes a drastically simplified setting (identity weight matrices, i.i.d. Gaussian embeddings, single induction pattern, 1-layer attention-only). It predicts that KV shifting makes induction easy to learn, which is consistent with the toy convergence speedup (Figure 1a) and the persistent training loss advantage (Figure 4). However, the theorem does not predict the magnitude of the convergence speedup, how it scales with model size, or how it interacts with the many other learning processes occurring simultaneously in full pretraining. The contour plots in Figure 2 are for a specific parametrization (α₂ = 1 - α₁, β₂ = 1 - β₁) with O(T) treated as constant—this is a further simplification of the already-simplified setting. The paper acknowledges this gap implicitly by separating the analysis into "representation" (Section 3.1, exact results) and "learning" (Section 3.2, approximate analysis under strong assumptions), but does not provide evidence that the learning dynamics in the simplified setting are a faithful model of what happens at scale.

Are there alternative explanations for the benchmark improvements? The paper does not control for the possibility that the improvement comes from the increased number of trainable parameters (however small) rather than the inductive bias of the shift operation. At 512 added parameters for the 2.9B model, this seems extremely unlikely to matter, but a control experiment adding the same number of parameters in a way that doesn't provide the KV shifting inductive bias (e.g., adding per-head bias terms to K and V) would eliminate this concern. The paper also does not control for the possibility that the shift operation acts as a form of data augmentation or regularization—by forcing each token's representation to incorporate information from its neighbor, the model might be less prone to overfitting to token-specific noise. The LR=1e-2 robustness result is consistent with a regularization interpretation.

Overall assessment of experimental support. The paper provides strong evidence that KV shifting attention accelerates induction head learning in controlled toy settings, and that it improves training loss and benchmark performance in large-scale pretraining at 2.9B and 19B scales. The evidence for the specific mechanistic claim—that the improvements are caused by reduced depth and width requirements for induction heads—is strongest in the toy experiments and more indirect in the pretraining results. The single-run nature of the main experiments, the lack of statistical characterization, the use of non-public training data, and the absence of direct mechanistic verification in trained large-scale models are the primary weaknesses. The shrinking gap with training duration and scale raises the possibility that the advantage is primarily a convergence speedup rather than a permanent improvement in model capability, though this remains unresolved given the incomplete training runs.

6. Limitations and Trade-offs

6.1 The Difficulty Estimation Cost Is Excluded from the Compute-Optimal Budget

The paper acknowledges but does not amortize the substantial computational cost of estimating prompt difficulty before strategy selection occurs. As noted in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

The consequence. The method requires generating 2048 samples per question and scoring them with the process reward model (PRM) to estimate difficulty, which consumes more compute than the largest individual test-time budgets studied (256–512 generations). The reported 4× efficiency gains over best-of-N (e.g., Figures 4 and 8: achieving equivalent accuracy with 4× fewer generations) are computed after difficulty is known and do not amortize the estimation cost. In a realistic deployment, the total cost would be (difficulty estimation) + (strategy execution), and the former could easily dominate the latter — potentially negating or even reversing the claimed efficiency advantage. For example, if estimating difficulty costs 2048 generations and the test-time budget is 64 generations, the total cost is ~2112 generations, far exceeding a naive best-of-256 baseline (~256 generations) that the compute-optimal strategy was supposed to beat.

What evidence exists in the paper. The paper explicitly acknowledges this gap in Section 3.2 and Section 8, flagging it as "a key avenue for future work" involving "pretraining or finetuning models to directly predict difficulty of a question." However, no such model is developed or evaluated. The predicted difficulty bins used in Figures 4 and 8 (which avoid needing ground-truth labels) still require generating 2048 samples and running the PRM — they only remove the correctness oracle, not the generation cost. The paper does not report total FLOPs including difficulty estimation, nor does it provide an ablation showing how performance degrades if difficulty is estimated from fewer samples (e.g., 8, 64, or 256 rather than 2048).

Mitigation status. The paper identifies adaptive difficulty estimation (estimating difficulty from a small number of initial samples, then allocating the remaining budget accordingly) as a natural next step, but does not implement it. This is an unsolved problem that the paper's practical deployment value depends on resolving.

---

6.2 All Experiments Use a Single Benchmark (MATH) and a Single Model Family (PaLM 2-S*)

The paper's entire experimental evaluation — search, revisions, compute-optimal scaling, FLOPs-matched comparisons — uses exclusively the MATH benchmark (Hendrycks et al., 2021) with PaLM 2-S* as the base model. The authors state they "believe this model is representative of the capabilities of many contemporary LLMs" (Section 4), but provide no cross-model or cross-domain validation.

The consequence. Several aspects of the findings could be specific to MATH or PaLM 2-S* and fail to generalize:

• PRM quality and over-optimization behavior. The observed over-optimization patterns (beam search degrading easy-problem performance at high budgets, Figure 3 right) depend on the PRM's calibration properties, which are a function of the base model's output distribution. A model with different calibration characteristics — better or worse alignment between its generation distribution and the PRM's training distribution — could exhibit qualitatively different scaling curves, potentially making beam search safe on easy problems (if PRM is well-calibrated) or degrading it on medium problems (if PRM is poorly calibrated).
• Revision model trainability. The revision model's ability to learn from incorrect in-context examples (Section 6.1, the edit-distance pairing procedure) depends on the base model's in-context learning capabilities, which vary substantially across model families (e.g., GPT-4, Claude, Gemini, LLaMA, Qwen all exhibit different ICL behaviors). The ~38% correct-to-incorrect reversion rate observed during revision chains might be higher or lower with different base models.
• Domain specificity of induction. MATH consists of competition-level symbolic reasoning problems with clear ground-truth answers. It is unknown whether the difficulty-dependent patterns — e.g., that sequential revisions help easy problems while balanced sequential-parallel ratios help hard ones (Figure 7 right) — generalize to code generation (where unit tests provide verifier signals), logical reasoning, scientific QA, or open-ended generation tasks where correctness is fuzzy or multi-dimensional.

What evidence exists in the paper. All results in Sections 5–7 are from MATH with PaLM 2-S*. There is no replication on any other benchmark (e.g., GSM8K for math reasoning, HumanEval for code, ARC for science) or any other model family. The computational cost of running a single FLOPs-matched comparison at multiple scales (~14× larger model variants) is substantial, which the paper implicitly acknowledges by limiting to one model pair.

Mitigation status. The paper does not attempt to address this limitation. Section 8 briefly notes the domain-specificity of the findings is unknown (e.g., extensions to other modalities, open-ended tasks), but this is framed as future work rather than addressed through even a small-scale validation experiment on a second benchmark.

---

6.3 The 14× Larger Model Baseline Is Not Compute-Optimally Trained, and It Uses No Test-Time Compute

The FLOPs-matched comparison (Section 7) attempts to answer: given a fixed total FLOPs budget, is it better to train a larger model or keep the smaller model and spend extra compute at inference? The comparison pairs PaLM 2-S* augmented with compute-optimal test-time strategies against a model with approximately 14× more parameters. The larger model uses greedy decoding and receives zero test-time compute augmentation — no majority voting, no best-of-N, no search, no revisions. Furthermore, the larger model is trained by scaling only parameters while holding data fixed (following the LLaMA paradigm; Touvron et al., 2023), which the paper acknowledges departs from compute-optimal pretraining (Hoffmann et al., 2022) where both data and parameters are scaled:

"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work."

The consequence. This makes the pretraining baseline systematically weaker than it could be, inflating the apparent advantage of test-time compute. Two specific confounds:

• Chinchilla-optimal training. A model trained with 14× more FLOPs allocated optimally between parameters and data would likely achieve higher performance than a parameter-only-scaled model of the same FLOPs budget. The paper's comparison essentially asks "is test-time compute with a small model better than a suboptimally-trained large model?" rather than "is test-time compute with a small model better than a compute-optimally-trained large model?" The answer to the former question could be yes while the answer to the latter is no, or the crossover point (where pretraining becomes preferable) could shift substantially.
• No test-time compute for the larger model. Even a modest test-time compute budget for the larger model — say, best-of-8 or best-of-16 with the same ORM/PRM — could close or reverse the gaps reported in Figure 9 and the bar charts in Figure 1. The paper's FLOPs accounting framework (Section 7) could accommodate this: instead of comparing a smaller model with test-time compute against a larger model with none, the comparison should allocate the total FLOPs budget to optimize both pretraining scale and inference strategy for each configuration. The current comparison gives test-time compute exclusively to the smaller model while depriving the larger model of the same tool, making the gap in Figure 9 partially an artifact of asymmetric resource allocation rather than a fundamental property of the pretraining-inference tradeoff.

What evidence exists in the paper. The paper is transparent about the parameter-only scaling choice (Section 7, quoted above), but does not discuss the implications of denying the larger model test-time compute. No ablation is provided where the larger model receives even a modest inference budget (e.g., best-of-4 or best-of-16 with the base-model PRM).

Mitigation status. The paper explicitly defers the compute-optimal pretraining comparison to future work, but does not frame the lack of test-time compute for the larger model as a limitation at all. Both issues weaken the headline claim that "a smaller model augmented with compute-optimal test-time strategies can outperform a ~14× larger pretrained model" — the claim is true under the paper's specific (asymmetric) experimental setup, but its generality to fair FLOPs-matched comparisons is unproven.

---

6.4 The Compute-Optimal Policy Is Selected on Only ~50 Questions per Difficulty Bin

The paper uses two-fold cross-validation within each of five difficulty quintiles on the 500-question MATH test set to select the best strategy (search method, sequential-to-parallel ratio, beam width, etc.) for each difficulty bin and budget level (Section 3.2). With 100 questions per bin and two-fold splits, the compute-optimal policy is selected based on performance on approximately 50 questions per fold per bin.

The consequence. With such a small sample for strategy selection, the compute-optimal policy may overfit to the specific questions in each validation fold, producing strategies that appear optimal on the validation fold but underperform on held-out data. This could inflate the reported performance of compute-optimal scaling relative to what would be achieved with a policy selected on a larger, independent dataset.

More subtly, the small bin size means that within-bin heterogeneity is "averaged over" — a question at the easy edge of bin 3 and one at the hard edge of bin 3 receive the identical strategy, even though the optimal strategy might differ. With 100 questions per bin, estimating difficulty continuously (rather than discretizing into quintiles) and learning a continuous policy function could yield further improvements, but also requires more data to avoid overfitting. The paper does not explore sensitivity to the number of bins or to using continuous difficulty estimates.

What evidence exists in the paper. The paper describes the cross-validation protocol in Section 3.2 and reports the averaged results in Figures 4 and 8. No confidence intervals, standard errors, or cross-validation variability estimates are provided for the compute-optimal curves. We cannot assess whether the difference between compute-optimal and best-of-N at, say, 64 generations (Figure 8: ~40% vs. ~37%) is statistically reliable given the small bin sizes.

Mitigation status. The paper does not address this limitation. The two-fold cross-validation within the test set is used to avoid contaminating strategy selection with evaluation, which is a sound principle, but the small bin sizes are inherent to the MATH test set size (500 questions) and the decision to use five difficulty bins. A larger test set or a separate validation set for strategy selection would be needed to produce more reliable compute-optimal policies.

---

6.5 Sequential Revisions Introduce Latency That Is Not Accounted For

The paper measures test-time compute in "generations" — the number of complete solutions sampled — which is a reasonable proxy for total FLOPs but ignores wall-clock latency (the time from prompt to final answer). This distinction matters because different strategies impose different serial dependencies:

• Fully parallel sampling (best-of-N with N=64): All 64 solutions can be generated simultaneously, and the verifier scores them in one batch. Total latency ≈ time to generate one solution + verification overhead.
• Fully sequential revisions (one chain of 64 revisions): Each revision depends on the output of the previous one. Total latency ≈ 64 × (time to generate one solution) + selection overhead.
• Hybrid (e.g., 8 parallel chains of 8 sequential revisions): Total latency ≈ 8 × (time to generate one solution) + aggregation overhead.

For a fixed generation budget of 64, the sequential-heavy strategy takes roughly 64× longer wall-clock time than a fully parallel strategy on hardware with sufficient parallelism. For latency-sensitive applications — interactive assistants, real-time code completion, conversational agents — this is a categorical difference that makes sequential-heavy strategies impractical regardless of their accuracy advantages.

The paper's compute-optimal policy often favors sequential-heavy strategies for easy problems (Figure 7 right: easy problems perform best with purely sequential revisions), precisely the regime where low latency matters most in practice (users expect fast responses to easy questions). The policy optimizes for accuracy-per-FLOP but not accuracy-per-second.

What evidence exists in the paper. The paper discusses latency nowhere in the main text or appendices. The generation budget is treated as the sole resource constraint. No wall-clock time measurements, latency analyses, or throughput comparisons are provided for any strategy or hardware configuration.

Mitigation status. Not addressed. The paper frames compute efficiency purely in terms of total FLOPs, omitting the orthogonal dimension of latency. A latency-aware allocation policy — e.g., weighting strategies by their serial depth and optimizing a utility function that combines accuracy and response time — would be a natural extension but is not discussed.

---

6.6 Revisions and PRM Search Are Studied Independently but Never Combined

The paper studies two complementary mechanisms for test-time compute scaling:

1. PRM-guided search (Section 5): Using a process reward model to steer search algorithms (beam search, lookahead search) that navigate the space of possible solutions.
2. Iterative revisions (Section 6): Fine-tuning the base model to sequentially revise its own incorrect answers, modifying the proposal distribution.

These mechanisms are analyzed in isolation and their compute-optimal strategies are reported separately (Figures 4 and 8 show compute-optimal search and compute-optimal revisions as distinct scaling curves). Section 8 explicitly acknowledges:

"we did not experiment with PRM tree-search techniques in combination with revisions"

The consequence. The paper's reported performance represents a lower bound on what a combined approach could achieve. The two mechanisms have complementary strengths suggested by the difficulty-dependent analysis:

• Revisions excel on easy problems where the model's initial attempts are roughly correct and need local refinement (Figure 7 right, bins 1-2: fully sequential revisions dominate).
• PRM search excels on medium-difficulty problems where the model needs to explore qualitatively different solution strategies, and the PRM's guidance helps navigate toward correct solutions (Figure 3 right, bins 3-4: beam search outperforms best-of-N).

A combined system — using the revision model as the proposal distribution within beam search, or using the PRM to decide which revision chains to pursue versus restart — could outperform either mechanism alone, particularly on medium-difficulty problems where both mechanisms show some benefit. The current results therefore understate what test-time compute can achieve when all tools are deployed jointly.

Furthermore, the paper's FLOPs-matched comparison (Section 7) reports results for revisions and search separately (Figure 9 shows separate panels for revisions and PRM search). The comparison against the 14× larger model never combines the two mechanisms, meaning the reported test-time compute advantage is a conservative estimate.

What evidence exists in the paper. The paper explicitly acknowledges this gap in Section 8. No experiment combines revisions with PRM search at any scale — even a small-scale toy experiment would provide evidence about whether the combination is additive, superadditive, or redundant.

Mitigation status. The paper identifies this as an important direction for future work, but provides no empirical evidence or even conceptual framework for how the combination would work. The complementary difficulty-dependent patterns (revisions excel on easy, search excels on medium) are suggestive but do not guarantee that combining them would yield gains — it is possible that the mechanisms interfere (e.g., the PRM, trained on base model outputs, may not transfer well to revision model outputs, as noted in Appendix J), or that the optimal allocation across mechanisms is simply to use one or the other per difficulty bin rather than both simultaneously. Without experimental evidence, the combined benefit remains speculative.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper introduces a methodological shift in how architectural modifications can be derived: it demonstrates that mechanistic interpretability can be generative—producing actionable architectural improvements—rather than merely descriptive. The standard pipeline in interpretability research has been to identify circuits in trained models, characterize their function, and use that understanding to explain model behavior (e.g., Elhage et al., 2021; Olsson et al., 2022; Bansal et al., 2023; Conmy et al., 2023). This paper closes the loop by asking: given that we understand why a two-layer circuit is necessary for induction heads in standard transformers, can we redesign the attention mechanism to collapse that circuit into a single layer? The answer is KV shifting attention, and its derivation is unusually direct—the shift operation follows logically from the specific mechanical bottleneck (keys and values bound to the same token preventing single-step neighbor-access) rather than from general desiderata like "reduce noise" or "improve sparsity."

The magnitude of this shift is incremental but catalytic. The modification itself is small (4 learnable scalars per head, O(ND) computation), and the paper does not claim to have invented a fundamentally new class of models. But the methodology—using circuit-level understanding to motivate a targeted architectural change, validating it first on toy tasks that isolate the circuit, then scaling to full pretraining—provides a template that other researchers can apply to other identified bottlenecks. The paper explicitly gestures toward this in the Discussion (Section 5):

"If we are not ready to completely overturn the Transformer architecture, it may be interesting to delve into its underlying mechanisms, identify important features, and modify the Transformer to make it easier to implement certain important functions. For example, by modifying transformer to make it easy to learn parity check, which might be very challenging (Wies et al., 2023)."

This reframes the relationship between interpretability and architecture design from parallel tracks to a closed loop, and it suggests that the interpretability community's extensive catalog of identified circuits (induction heads, function vectors, copy suppression, name-mover heads, etc.) is a resource for architecture designers, not just for explainability researchers.

The paper also resolves a subtle contradiction in the literature. Prior work established that single-layer standard transformers cannot implement induction heads (Sanford et al., 2024a), and that two-layer transformers approximate them with error bounded by O(e^(-p^(1))) dependent on positional encoding bias (Wang et al., 2024). But no prior work asked whether this limitation is fundamental to the attention mechanism or merely an artifact of how keys and values are bound. The paper resolves this by showing that a minimal decoupling of keys and values eliminates the depth requirement entirely (Theorem 2: exact representation with one layer, D = d) and halves the width requirement. The contradiction—"transformers need two layers for induction, but one layer could suffice if we change the attention"—is resolved in favor of the latter, with the caveat that the change is architectural rather than discoverable by training a standard transformer differently.

Concretely, this work makes several research directions more attractive:

• Circuit-motivated architecture design becomes a viable methodology. Rather than proposing architectural modifications and hoping they work, researchers can start from a specific circuit-level bottleneck and derive the minimal change that addresses it, validating on toy tasks before scaling. The paper's success with induction heads—arguably the most well-studied circuit in transformers—provides a proof of concept.
• Width efficiency becomes a first-class design criterion. The paper's identification of the width requirement for induction (Property 1, Section 3.1) as a distinct bottleneck from depth means that future architectural modifications can be evaluated not just on whether they reduce the number of layers needed, but on whether they reduce the hidden dimension needed per operation. This insight generalizes beyond induction: any circuit that requires a token's hidden state to simultaneously represent multiple pieces of information (its own identity + a neighbor's identity, or a query + a key, etc.) imposes a width cost that could potentially be eliminated by decoupling the representations.
• The Alibi approximation error in standard induction circuits is a previously unrecognized practical concern. Theorem 1 shows that the standard two-layer induction circuit has approximation error O(e^(-p^(1))) because the copying operation in the first layer distributes softmax attention mass to non-target tokens. This error source is eliminated entirely in KV shifting attention, where the shift is a hard, deterministic copy (Theorem 2 takes an equal sign). This suggests that positional encoding schemes interact with circuit formation in ways that go beyond length extrapolation—they affect the fidelity of learned circuits, which may partially explain why different positional encodings produce different ICL capabilities.

Conversely, this work makes some research directions less attractive:

• Merely making attention more expressive may not help if the bottleneck is circuit-specific. The paper's ablation showing that longer shift windows (Figure 8c: KV shifting 2 and 3) do not improve performance suggests that adding capacity beyond what specific circuits need is wasteful. This argues against architectural modifications that increase attention's general expressiveness without targeting specific computational bottlenecks—the key is identifying what operations the architecture struggles with and addressing them precisely.
• Interpreting induction heads in standard transformers may be partially obsolete as an optimization target. If KV shifting attention or similar modifications become widely adopted, the induction head formation problem that interpretability researchers have studied extensively in standard transformers will manifest differently. Heads that learn α₁ ≈ 0, α₂ ≈ 1 and β₁ ≈ 0, β₂ ≈ 1 are induction heads in a direct, architecturally-supported sense, not emergent circuits that gradient descent must painstakingly assemble from two layers. This doesn't diminish the value of past interpretability work—it validates it by showing that understanding leads to improvement—but it does mean that the specific circuit patterns documented in Elhage et al. (2021) and Olsson et al. (2022) are tied to the standard attention architecture and may not transfer.

---

Follow-Up Research This Work Enables

Mechanistic verification that KV shifting heads actually implement induction in large-scale models. The paper shows that ~40% of heads in the trained 2.9B model converge to induction-like (α, β) configurations (α₁ ≤ α₂ and β₁ ≤ β₂, Table 4), but never demonstrates that these heads actually perform induction on natural text. A follow-up study using activation patching, attention pattern analysis, or knock-out experiments on the released 2.9B checkpoints could confirm that heads with shifted parameters exhibit the defining induction head behavior: attending to previous occurrences of the current token and boosting the logits of their successors. Without this, the connection between the learned parameters and the claimed mechanism remains correlational. A strong experiment would compare induction head behavior in the KV shifting model versus the vanilla model on the same text, measuring whether induction patterns form earlier in training (consistent with Theorem 3's prediction of faster convergence) and whether they remain more robust to distribution shift.

Compute-optimal width comparison between KV shifting and vanilla models at equal parameter counts. Theorem 2 predicts that KV shifting attention can represent induction heads exactly at hidden dimension D = d, while vanilla attention requires D = 2d for approximation. This predicts that a KV shifting model with hidden size D should match or exceed the induction capability of a vanilla model with hidden size 2D, all else equal. No such experiment is performed in the paper—the toy width experiment (Figure 1b) uses hidden size 8 for both, showing KV shifting succeeds where vanilla fails, but doesn't test whether vanilla could succeed at hidden size 16. A controlled experiment would train matched model pairs at multiple widths (e.g., vanilla at D = 512, 1024, 2048 vs. KV shifting at D = 256, 512, 1024) on tasks requiring induction (the synthetic induction data from Appendix D, or multi-hop reasoning from Appendix H), measuring whether the theoretical factor-of-2 width reduction holds empirically and whether it extends to natural language benchmarks. If confirmed, this would have immediate practical implications: for a fixed parameter budget, KV shifting models could allocate fewer dimensions to induction and more to other computations, or equivalently, could achieve the same induction capability with smaller hidden sizes overall.

Difficulty-aware dynamic KV shifting: learning to adjust shift parameters per token. The current implementation learns static α₁, α₂, β₁, β₂ per head—every token processed by a given head uses the same shift ratio. But different tokens likely benefit from different amounts of neighbor-information integration: a token in the middle of a predictable phrase might need little shifting (α₁ ≈ 1), while a token at an induction point (where the current token matches a previous occurrence and the model needs to retrieve the successor) might need maximal shifting (α₁ ≈ 0). A follow-up could make the shift parameters content-dependent by predicting them from the token's hidden state: α₁ = σ(f_α(x)), where f_α is a small learned projection. This would allow heads to dynamically switch between standard-attention and induction-attention modes per token, potentially achieving the best of both: standard attention for bigram prediction, shifted attention for induction. The experiment would compare static KV shifting against dynamic KV shifting on the full pretraining benchmark suite (Table 2), with the hypothesis that dynamic shifting improves performance on tasks requiring both memorization and induction (e.g., MMLU facts requiring recall + reasoning chains).

Combining KV shifting with other circuit-motivated modifications for cumulative improvement. The paper's methodology—identify a circuit-level bottleneck and modify the architecture to eliminate it—is not specific to induction heads. Other circuits that transformers struggle to learn could be targeted similarly. For example, Wies et al. (2023) showed that transformers struggle with parity checking; could a modification analogous to KV shifting (e.g., providing a dedicated "XOR gate" operation in the feedforward layers) make parity learnable in fewer layers? A follow-up could systematically catalog known difficult circuits for transformers (parity, composition, long-range copying, hierarchical structure induction) and propose minimal architectural modifications for each, then test whether stacking multiple modifications (KV shifting + parity gate + composition module) yields cumulative gains or whether the modifications interfere. This would test the broader hypothesis that circuit-motivated architecture design is composable—that fixing bottleneck A doesn't break the fix for bottleneck B. The released KV shifting checkpoints provide a baseline for this composability study.

Negative result: does the KV shifting advantage persist through training to full convergence? The paper's pretraining experiments show a shrinking advantage with more training tokens: at 2.9B, the benchmark gap narrows from +3.98 points (340B tokens) to +2.12 points (500B tokens); at 19B, from +1.31 points (160B tokens) to +0.77 points (200B tokens). The paper speculates that "it may take several TB data for a 2.9B model to converge," but doesn't test this. A critical stress-test experiment would continue training both the 2.9B KV shifting and vanilla models to convergence (as measured by validation loss plateauing), checking whether the benchmark gap closes entirely or stabilizes at some positive residual. If the gap closes to zero, KV shifting's benefit is primarily a convergence speedup—still valuable (reaching equivalent performance in fewer tokens saves compute), but structurally different from a permanent capability improvement. If a non-zero gap persists, it suggests the architectural bias enables learning patterns that vanilla attention cannot acquire regardless of training duration. This experiment is feasible with the released checkpoints: simply continue training from the 500B checkpoint. The outcome has direct implications for whether practitioners should adopt KV shifting only when training budget is limited (convergence speedup) or always (permanent capability gain).

Extending to encoder-decoder and bidirectional architectures. The paper focuses exclusively on decode-only (causal) transformers, motivated by the induction head mechanism's role in autoregressive next-token prediction. But the core operation—decoupling keys and values to allow attending to one token's key while receiving a neighbor's value—has no causal dependency; it could be applied in encoder-decoder models (T5, BART) or bidirectional encoders (BERT). In a bidirectional setting, the shift could operate in both directions (using both the previous and next token), providing a richer neighborhood representation without the causal constraint that restricts KV shifting to the previous token only. A follow-up would implement KV shifting in a BERT-style masked language model, evaluating on GLUE/SuperGLUE benchmarks and probing whether the modification improves performance on tasks requiring local context aggregation (e.g., named entity recognition, where the identity of neighboring tokens is crucial for classification). The theoretical prediction would be that the benefit is smaller than in autoregressive models (because induction heads are primarily an autoregressive prediction mechanism, not a bidirectional understanding mechanism), but that the general principle of decoupled key-value access might aid other operations.

---

Practical Applications and Downstream Use Cases

Production LLM pretraining with faster convergence. The most direct practical application is adopting KV shifting attention in any Llama-style or similar decoder-only transformer pretraining pipeline. The modification adds 4 learnable scalars per KV head per layer—512 total parameters for the 2.9B configuration, a ~0.00002% increase—and requires minimal code changes (the paper provides PyTorch implementation in Appendix F that integrates with flash attention). The practical benefit is that at the 2.9B scale with 500B training tokens, KV shifting achieves +2.12 percentage points higher average benchmark accuracy than vanilla (Table 2: 38.57% vs. 36.45%). For a production team training a model of this scale, that improvement comes at essentially zero additional cost in parameters, computation, or engineering effort. The robustness to learning rate (Figure 5b: KV shifting outperforms vanilla at LRs from 1e-4 to 1e-2, and converges when vanilla diverges at 1e-2) further reduces the cost of hyperparameter tuning. The fact that the advantage shrinks with training duration (+3.98 at 340B → +2.12 at 500B) means the primary value proposition is reaching a given performance level with fewer training tokens, which translates directly to reduced compute cost.

On-device or edge deployment where width is constrained. The width-reduction property of KV shifting attention (Theorem 2: exact induction representation at D = d vs. D = 2d for vanilla) is particularly relevant for deployment scenarios with tight memory constraints—mobile devices, embedded systems, browsers. A KV shifting model with hidden size 1024 should have equivalent induction capability to a vanilla model with hidden size 2048 (all else equal), which means the smaller model could match the larger model's pattern-matching and in-context learning abilities while using ~4× less memory for attention operations. The paper doesn't directly test this equivalence at realistic scales (the width experiment is at hidden size 8 on synthetic data), but the theoretical result is constructive and the synthetic validation at the extreme low end is suggestive. A practitioner deploying a small language model on-device could use KV shifting to either (a) improve induction capability without increasing model size, or (b) reduce the hidden dimension while maintaining induction capability, freeing memory for other components.

Self-improvement and synthetic data generation pipelines that rely on pattern matching. Many self-improvement methods (STaR, ReST^EM, rejection sampling fine-tuning) require the model to generate solutions and then verify or improve them, often using in-context examples to demonstrate the desired output format or reasoning pattern. KV shifting attention's bias toward faster induction head formation means that models trained with this architecture may be more sample-efficient at learning from in-context demonstrations—they acquire the ability to recognize and replicate patterns in the context window more quickly. The convergence speedup on toy induction (Figure 1a: roughly 5-10× faster) and the stronger few-shot MMLU performance (Table 3: +3.62 points for few-shot vs. +1.76 points for cloze) provide evidence for this. In a self-improvement pipeline where the model iteratively generates training data for itself, the quality of the generated data depends on how well the model can follow the pattern demonstrated in the prompt; better induction capability should produce higher-quality synthetic data, potentially accelerating the self-improvement loop.

Mathematical and multi-step reasoning systems where induction compounds. The hop-k experiment (Appendix H, Figure 9) shows that KV shifting attention's advantage grows with the number of sequential induction steps—at hop count 16, the error rate gap is dramatically larger than at hop count 4. The iGSM math experiment (Appendix I, Table 7) shows +7.5 percentage point improvement on problems requiring generalization to longer reasoning chains than seen during training. This suggests that KV shifting attention is particularly well-suited for systems that perform multi-step reasoning by chaining primitive induction operations—the kind of reasoning that chain-of-thought prompting and tree-of-thought search rely on. A deployed reasoning system using KV shifting attention could maintain accuracy on longer reasoning chains where a vanilla model's performance degrades (as the hop-k experiment shows), enabling more reliable multi-hop question answering, program synthesis with longer specifications, or mathematical proofs with more intermediate lemmas.