KV Shifting Attention Enhances Language Modeling
Authors: Mingyu Xu
Bingning Wang
Weipeng Chen
Venue: ICML 2025 poster (2025)
Keywords: Attention mechanism
Large Language model
Induction heads
Primary area: deep_learning->large_language_models
TL;DR: We find that decoupling the key and value in attention helps to alleviate the requirements of the depth and width of induction heads, thereby improving language modeling.

Abstract: Current large language models (LLMs) predominantly rely on decode-only transformer architectures, which exhibit exceptional in-context learning (ICL) capabilities. It is widely acknowledged that the cornerstone of their ICL ability lies in the induction heads mechanism, which necessitates at least two layers of attention. To more effectively harness the model's induction capabilities, we revisit the induction heads mechanism and provide theoretical proof that KV shifting attention reduces the model's dependency on the depth and width of the induction heads mechanism. Our experimental results confirm that KV shifting attention enhances the learning of induction heads and improves language modeling performance. This leads to superior performance or accelerated convergence, spanning from toy models to pre-trained models with over 10 billion parameters.

KV shifting attention enhances language modeling
Mingyu Xu
Bingning Wang' Weipeng Chen
Abstract
analyze the required depth and width for induction heads. Based on the perspective, we propose KV shifting attention 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.
Current large language models (LLMs) predomi- nantly rely on decode-only transformer architec- tures, which exhibit exceptional in-context learn- ing (ICL) capabilities. It is widely acknowledged that the cornerstone of their ICL abiliity lies in the induction heads mechanism, which necessitates at least two layers of attention. To more effec- tively harness the model's s induction capabilities, we revisit the induction heads mechanism and pro- vide theoretical proof that KV shifting attention reduces the model's dependency on the depth and width of the induction heads mechanism. Our experimental results confirm that KV shifting at- tention enhances the learning of induction heads and improves language modeling performance. This leads to superior performance or accelerated convergence, spanning from toy models to pre- trained models with over 10 billion parameters.
Through theoretical analysis and empirical validation, demonstrate that KV Shifting attention achieves comparable or superior performance to Vanilla attention. Moreover, its bias towards learning induction leads to more efficient and effective language modeling across diverse scales, from toy models to pre-trained models with billions of parameters. In summary, our contributions are follows:
To accelerates learning ability for induction heads, we revisit the mechanism of it, and analyze that KV shift- ing attention can effectively represent induction heads with lower depth and width, and learn induction heads from induction data.
We apply KV shifting attention from toy model to large language modeling pre-training and demonstrate its effectiivnnss.
1. Introduction
Transformer-based Large language models (Vaswani, 2017) have demonstrated remarkable capabilities in in-context learning (ICL), largely attributed to the underlying mecha- nisms of induction heads (Elhage et al., 2021; Olsson et al., 2022). These mechanisms enable the models to identify and leverage repeating patterns, which is crucial ICL (Song et al., 2024; Crosbie & Shutova, 2024) and multi-step reasoning (Sanford et al., 2024b).
We arrange the paper in the following order. We introduces our motivation and methods in Section 2, analyzes the KV shifting attention in Section 3 introduces our experiments on large language model in Section 4, discuss, introduce relevant work, and summarize in the last few sections.
Although there many works based transformer for analysis of induction heads of LLMMs, there are few works that utilize analysis of induction heads to modify transform- ers to enhance its ability to learn induction heads. A better learning induction head may potentially contribute to the in-context learning or reasoning ability.
2. Method
2.1. Motivation
Induction heads Our story begins with induction heads. Induction heads mechanism (Elhage et al., 2021; Olsson et al., 2022) is a circuit whose function to find the latest same previous instances of the current token (call it A) and the token that came after it (call it B), then use B as current token's 's next token prediction. (e.g. forming the sequence [A][B]... [A] + [B]). It is generally believed that the implementation of the induction heads mechanism often requires two heads belonging to different layers. For a detailed proof that a layer of transformer cannot
For the goal, revisit the induction head mechanism and 'Baichuan-ine inc, China. Correspondence to: Bingning Wang <god@ bingning.wang. wang>.
Proceedings of the 4 International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

KV Shifting Attention Enhances Language Modeling
implement induction heads, readers can refer to Sanford et al. (2024a). So, can we make slight adjustments to the attention so that single layer of attention can achieve the mechanism of induction heads?
5dh token's value, j E (i- + 13. From different perspective, current token can obtain the value of the ith token by focusing on the keys of the (
i, i
1} tokens.
However, if we directly use the combination of 5th token's 's value (j E {i - 4 i + 13), it breaks causal mask, because (i [i+1] 1++1thh token's value can't be computed when jth token try to do next token prediction. So we can only let 2th token's key connect to (i 1)th and 2th token's value, And if we do similar operation to value, which means let 7th token's value connect to (i - 1)th and ith token's 's key, we can find that we can pay attention to the key of jth token, and then get the value of its near token without breaking causal mask.
Virtual attention heads Another interesting perspective proposed in the article (Elhage et al., 2021) is the concept of virtual attention heads. It demonstrates the cooperation among attentions in different layers. For simplicity, we consider the attention of only two adjacent layers and ignore the presence of residual connections and MLP. The formal expression is:
X1 = A XO oU X2 = A X1 W
(1)
Therefore, we propose the following KV shifting attention. For simplicity, we only formalize the single head attention as follow:
where Xo E RNXD; is the input, X E RNND is the output of 1th layers, A E RNxN is the attention weights of 1hh layer, on = 8 projection, Whi E ROxD is the o projection of 1hh layer, I E 1,2, n R context length, is the dimension of hidden states. Then virtual attention heads is:
Whihhh E RDxD, Whi E RDxD is the v
Q,K, K = XWQ. XWx, XWv
(4)
K, V a + azShiit(K), R.V + BaShift(V)
(5)
Output = Softmax(QK" M/o)
(6)
X2 A x ov h2 A ) (W W
(2)
where X E R is hidden states, Wo. WK, Wv, Wo E
Through virtual attention heads, models can achieve com- plex functions by combining simple attention heads.
RDxD 01 022, 31,32 E R are learnable parameters, o =
VD, M E RDxD is the causal mask, Shift(-) means dis- carding the last token and padding zero at the beginning. Compared to the original attention, 4 learnable parame- ters have been added. In the case of multi head attention, 01, 022, $1, f2 are learned per head, so the additional pa- rameters is 4h. The additional calculation caused by Eq 5 is O(N D). which is much smaller than O(N + N D) in Eq. 6.' We provide the training and inference code for PyTorch implementation in the Appendix M.
With causal mask, there is an interesting things when virtual attention heads do function like induction heads.
Property 1. With causal mask and j 2 i, we have
N
AA
1,,41
Σ
2
Σ Ami+1
A
A
(3)
According to Property 3.2, from the perspective of virtual attention heads, it is difficult for the model to indirectly utilize j tokens to focus on the (i + 1 token through 2th 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 5th token into its hidden states, even if the information of the 7th token is useless for predicting the (i + 2)th token. That's to say, it imposes certain requirements on the dimensionality of hidden states. We will further explain in the next section that the induction heads mechanism requires a certain width.
In order to ensure that the initialization of KV shifting does not affect the initial optimization state, we select on and 31 from (0, 1) and let ag 1 - a1, 32 1 - B1.
3. Analysis
In this section, we will analyze the KV shifting attention. The first subsection will examine how KV shifting attention has better ability to characterize induction heads compared to vanilla attention. The second subsection is to learn in- duction heads on the toy model and analyze their dynamic process. For more small experiments about the capability boundary of KV shifting attention, readers can refer to Ap- pendix E. For an explanation of another perspective on KV shifting attention, please refer to Appendix 3.3.
2.2. KV shifting attention
From more general perspective, induction heads means obtaining information about some tokens by focusing on it surrounding tokens. In order to make it casier for the transformer to learn the mechanism of inductions heads, we can unbind the key and value of the th token. When current token attention to the ith token's 's key, can get the
'In current popular group queries attention (Ainslic et al., 2023) for large language models, the additional parameters is 4/s, where hy refers to the number of KV pairs, not the number of heads, and the additional calculation caused by Eq. (5) is O(Nhyd1), ). where dj is the head dims.
2

---

KV Shifting Attention Enhances Language Modeling
3.1. Better representation for induction heads
related to Alibi's bias. And the KV shifting attention, due to the absence of this noise, Eq. 9 takes an equal sign.
The KV shifting attention reduced not only the depth but also width requirements of the transformers for forming induction heads mechanism. Improving the requirement for depth is very intuitive, while improving the requirement for width is the question we left in section 2. To strictly illustrate this point, we first modeling KV heads attention do induction heads. Then we use the theorem from previous work.
For the theoretical upper limit of the transformer structure in induction heads or more generalized tasks, readers can refer to Sanford et al. (2024b). We believe that after replacing the standard Transformer with KV shifting attention, although there may not be a margin improvement in the theory bound, it still possible to achieve a constant improvement. Next, we will leave the field of representation and enter the field of practice, which will provide experimental support for some of the assertion in this section.
Definition 3.1. (Induction heads) We define the induction heads machine modeling by IH : ULEVI RLxD + RD
with Alibi RPE (Press et al.) as follows:
3.2. Great bias when learning induction heads
IH(x))
Σ softmax (IL ml|-- (0).-s),+ 8)), Fz. (7)
Which has the ability to implement induction heads for any long context, when T > O and m > O, which infinitely approach In practice, transformers processed in a finite length, only a very small and can to be taken."
C5
(a) Various depth
(b) Various width
Previous work (Wang et al., 2024) has shown that two- layer transformer (Alibi RPE, without FFN) can be used to approximate IH, the detail is as follow:
On the left, as the training step size increases, the accuracy of induction varies among different models. In this setting, the only difference between Vanilla and KV shifting attention is the calculation of key and value. The total parameters of Vanilla and KV shifting attention with one layers is the same. And the parameters of Vanilla with 2 layers is twice. On the right is the induction accuracy with different hidden size. There are two layers in Vanilla model, and one layer in KV shifting attention, which means Vanilla model has two times parameters than KV shifting attention.
Figure 1:
Theorem 3.2. (Modify from Wang et al. (2024)). There exists a constant C >> O and two-layer single-head transformer TF(without FFNs), with D 2d, w0.11) w = 0. p(2)- = m, (p(t) means the Alibi bias in 8th layers), and [ 11. (1 s 1/1,1//0) such that
IH - -TFIL 0
sup LEN+
(8)
Various depth The architecture of the model adopts the same architecture as Llama (Touvron et al., 2023)* with approximately 20M no-embedding parameters. We used huge vocabulary with 8000 tokens to randomly generate sentences and ensure that the sequences in them satisfy the condition that when the jth token is the same as the 7th token, then the (j + 1)th token is the same as the 7t token (i < j). We present the accuracy of jth next token prediction's accuracy in Figure la. The tasks is similar to MQAR (Arora al., 2023), but the difference is that each token is key for the next token and also value for the previous token in our setting, while key value pairs are paired and disjoint in MQAR.
If we use KV shifting attention, we can get better estimation as follow:
Theorem 3.3. There exists one-layer single-head KV shifting atlention KVSA, with D = d, such that
IH = KVSA.
The proof in the case of using KV shifting attention is rel- atively simple if we are familiar with the induction heads. We include the proofs of Theorem 3.2 and Theorem 3.3 in Appendix A and Appendix B. From Theorem 3.2 and Theorem 3.3, we can find that KV shifting attention can represent or approximate induction heads with less depth and less width. In addition, 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 We use infinite precision transformers this article.
'Theorems I 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 smaller than the upper bound Theorem
*These is a slightly different from the experiments in (Elhage et al., 2021), which use the attention-only structure, while we chose the Llama architecture for practice.
3

---

KV Shifting Attention Enhances Language Modeling
From Figure la, we can see that it is indeed difficult for one-layer standard transformer to learn the ability of induc- tion. And the 2-layer transformer and I-layer KV shifting attention can perfectly learn the ability of induction. At the same time, KV shifting attention has a bias towards better learning induction, and its convergence speed is much faster than two-layer transformers. At the same time, we also found that increasing the model depth to layer 4 for Vallina does not make the model learn induction faster.
(a) (T)) =
(b) ((T) 10
(c) ( 100
Figure 2: Contour lines and gradient decent derection of L. We simplified ((T) as a constant, and 0g = 1 01 and P2 = 1 - B. Induction heads means (a1, (1) = (0, 1).
attention is equivalent to:
Various width Hidden size for the experiment in Figure la is set to 1024, which is relatively relaxed setting. As we analyzed in sections 2.1 and 2.3. Existing attention may require additional width to perform well on induction heads, so as the dimensionality of hidden states decreases, it will become increasingly difficult for standard attention learn induction heads. Therefore, we conducted pressure testing with hidden states=8. The results are shown in Figure 1b.
min L
e 4/288
log gag3it82/5 + 2e /8tp///++a + 2201 + OTT)
(10)
where 02 =
0j = = c + 2e + OTT).
all (.) here is greater than or equal to 0.
The detail proof is in Appendix C. The dynamic process of optimization is interesting. We consider O(T) as constant and plot the contour lines and gradient directions of - L for 02 = 1 01 and B3 = 1 - B1 in Figure 2.
From Figure Ib, we find that the learning ability of standard attention is very poor, and even failed to cover up one of the answers mentioned in the previous text.
Although this is a toy task, people may think that the current model has 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. Undoubtedly, this will have an negative impact on language modeling.
From Figure 2, as O(T) increases, the contour lines become sparse, which means the convergence speed slows down, In addition, when O(T) is relatively small, the direction of GD may undergo small non monotonic learning pro- cess. When O(T) is relatively large, the direction of GD is relatively consistent. In practice, we often have many attention heads, (a), 31) of some heads are closer to (0, 1) during initialization, making it much easier for them to learn induction heads.
Learning from induction data Regarding how traditional two-layer transformers learn about induced heads, the an- alyzing it is a very complex. Previous work (Bietti et al., 2024; Wang et al., 2024; Chen et al., 2024) has provided analysis under simplified conditions. We also follow their work and provide an analysis of the dynamic process of KV shifting attention in learning induced heads. We use the following simplified conditions: removed the residual con- nection, MLP, normalization, position embedding and use tie embedding and the each component of the embedding of each token is from independently and identically distributed (O,, 2). Wg. WA. W., W. = T, and we assume the sen- tence's 's length is T + 1 and vocabulary size is T (T' 2 3) and every token only appearance once except the last token appearance twice and calculate cross entropy loss only when predicting the next token on the last token. We analyzed the learning process of the four additional variables a 022 S1, and 32 introduced in KV shifting attention. We have:
When T 2, which is actually no longer applicable to Theorem 3.4, training data become limitation, such as [A] [B] [A] + [B]. The model only needs to predict the previous token to achieve low loss on the training data, but this is not an induction head. Learning such a head cannot achieve [A] [B] [C] [A] + [B].
In Bietti al. (2024), they think global bigrams are learned first, then the induction head is formed by learning appro- priate memories in a top-down fashion. But obviously, in KV shifting attention, induction heads become very easy to learn, and even with good initialization, the model has induction capability. But it is difficult to have appropriate initialization to obtain a certain level of bigrams capability without training.
There are some additional small-scale experiments on multi- hop, math and n-gram, which is in Appendix 1. F.2 and E.3. Experiments have shown that performing shift opera- tions on KV can significantly enhance abilities related to induction heads, such as multi-hop and math, but does not improve n- n-gram tasks that rely more on memory.
Theorem 3.4. Under the simplified condition we describe, and as d approx oc, learning induction heads by KV shifting
4

---

KV Shifting Attention Enhances Language Modeling
3.3. Meta learners and checker for in-context learning
Model Configuration Since the baseline of the model trained from scratch is for production environments, not not just for this paper, it uses Group query attention (GQA) (Ainslie et al., 2023) to reduce memory usage during inference and employs a larger vocabulary (48000) to cope with more multilingual environments. For this reason, our experiment on KV shfting attention, which was trained from scratch, was also based on the baseline. Instead of training the model from scratch, we used Llama's original vocabulary size of 36000 and employed standard Multi head attention (MHA), The detail configuration is as shown in Table 16.
In previous section, we present our primitive motivation, which is to reduce the depth and width requirements of induction heads through KV shifting attention. The forma- tion of traditional induction heads mechanism requires the fusion of adjacent token information in the previous layer and the coupling of information into keys and values in the next layer, but fusion is easy and decoupling information is difficult, we think this may also be one of the reasons why Vanilla attention has slower learning. In this section, we will provide another perspective on KV shifting attention.
Table I: Model Configuration.
Firstly, K shifting attention can be seen as a meta learner of in-context learning. For example, after undergoing K shifting, the KV pairs mentioned in the previous text are (ki, o2), (K2, va),..., (k,-1, Therefore, we can use the similarity between kn and historical ky to infer what Un+1 is, shown below:
PARAMETERS
1.5B
2.9B
6.7B
13B
19B
HIDDEN SIZE
2,048
2,560
4,096
5,120
6, 1.44
LAYERS
28
32
32
40
48
HEAD NUMBER
20
32
40
48
KV NUMBER
16
4
32
40
FFN SIZE
5,504
8,704
11,008
13,824
16,384
MAX LENGTH
2,048
4,096
2,048
2,048
12,288
Un+1=
chicke/T U/.
(I1)
TOTAL TOKENS
500B
10B
200B
VOCAB SIZB
36,000
48,000
36,000
36,000
48,000
If we infer 1, we can infer what the (n + 1)th token is. Then, we can find that V shifting attention can be seen as checker of in-context learning. For example, af- ter V shifting, we obtained such a pair of KV pairs, (k2, v1), (k3, u2),... (k2, 1), then if we guess the (n + 17th token is Tn+1, then we can use the guessed + as query to do attention to infer on (although we actually know what the real U/ is) as follow:
Datasets Due to some commercial reasons and data lim- itations, we used non-public private data for pre training, which means that the models are trained on our own dataset. Our data collection and filtering methods are similar to FineWebbedi (Penedo al., 2024). When computing re- sources are available, we will use open-source data, such as RedPajama- -TWeee et al.) to train two models for com- parison, one is the baseline and the other is the KV shifting attention.
22
Un
EE= elatt 1/-1-
(12)
Hyperparameters We used a constant learning rate with linear warmup of 1000 steps. The learning rate for 1.4B / 3B / 7B / 13B 19B model 2c-4 8e-4 2e-4 / 2e-4 2c-4, the batch size is 1M/16M/MM/2//3X/ For optimization, we apply the AdamW optimizer with 31 0.9 and 3 = 0.95, and weight decay = 0.1.
Then we can compare the guessed vn with the real en- If they are close, keep the guess for Fn+1, otherwise delete the guess for 1n+1 1 (keep or delete operation can be easily done in a Gate MLP layers and residual connection).
From this perspective, introducing KV shifting as an in- ductive bias into attention is much suitable for next token prediction based on the decode-only structure. Next, we will leave the toy models and validate the effectiveness of KV shifting attention on large-scale models.
To validate the performance of different mod- els, we used some benchmarks for the 3B and 19B models that were trained more tokens, including Lambada(Paperno et al., 2016), Sinogrande(Sakaguel) et al., 2021), Hel- laswag(Zellers et al., 2019), ARC(Clark et al., 2018), CMMLU(LI et al., 2023a), MMLU(Hendrycks et al.), Math(Hendrycks et al., 2021). While for the models that were trained less token, we could only look at their loss curve. Due to computation limitations, almost experiments are only run once except there an additional notion.
Evaluation
4. Experiments
4.1. Setting
In this section, evaluate KV shifting attention across two models trained from scratch: 2.9B/19B parameters within an architecture similar to Llama2 (Touvron et al., 2023). The experiments trained from scratch are conducted on Nvidia H800-80G GPUs, while others are conducted on Nvidia A 1100-80G GPUs.
'The Vanilla-2.9B and KV shifting-2.9B has been open- sourced.
6IM here means 1,048,576 1,024 1,024, while elsewhere in the paper it refers to 1,000,000
5

---

KV Shifting Attention Enhances Language Modeling
Table 2: Main results. We trained four models on 2.9B and 19B parameters respectively, with the 2.9B model having a total training token count of 500B and the 19B model having a total training token count of 200B. ARC-E short for ARC-easy, and ARC-C short for ARC-Chalee
MODEL
TOKENS
LAMBADA
WINOGRANDE
HELLASWAG
ARC-E
ARC-C
CMMLU
MMLL
MATH
AVERAGE
3408
52.92
$2.09
42.70
27.45
25.97
28.51
29.43
0.80
32.48
VANILLA 2.9B
52.80
43.68
26.02
34.77
30.34
1.20
34.08
500B
51.66
54.06
44.49
36.20
27.90
38,22
37.26
1.80
36.45
340B
55.44
53.91
42.87
36.74
30.04
34.51
36,20
2.00
36.46
KV SHITING- 2.9B
420B
51.91
54.78
43.83
36.66
31.91
37.24
34,30
1.80
36.55
500B
54,51
55.33
44.52
39.02
30.89
40.78
40.88
2.60
38.57
1608
59.93
48.22
48.25
30.34
24.56
39.12
39.22
1.80
36.43
VANILLA- 19B
180B
58.80
48.07
47.78
31.28
25.99
40.80
39.34
2.60
36.83
200B
60.88
49.01
47.36
33.25
25.78
42.92
42.68
2.60
38.06
1608
61.93
48.46
48.28
31.25
25.06
42.10
42.87
2,00
37.74
KV SHIFTING 198
180B
60.20
47.67
48.16
32.45
26.55
43,38
40.49
3.00
37.74
200B
62.35
48.38
48.42
33.28
29.32
42.40
43.29
3.20
38.83
4.2. Main Result
We pre trained language models at two scales with 2.9B and 19B parameters, and the experimental results are shown in Table 2. The results indicate that KV Shifting attention achieved better performance than baseline at various scales and training token numbers. In addition, we also plotted the training loss of them, as shown in Figure 3.
(a) Various Sceds
(b) Various L.R
(c) IR=1e-2
Figure 4: Training loss of 1.5B parameters model among random seeds and learning rate (LR).
In addition, we also conducted experiments on different learning rates, including le-4, 2c-4, Ie-3, and le-2, as shown Figure 4b. can be observed that KV shifting attention achieves better results than Vanilla at different learning rates. And when the learning rate is set to 1e-2, Vanilla has already diverged, while the loss of KV shift- ing attention has not yet diverged as shown in Figure 4c This suggests that the optimization space for KV shifting attention may be flatter. Shifting KV provides a smooth key and value during model initialization, which may make the optimization process smoother. This also partially explains why in Figure 3, the 2.9B KV shifting attention leads more in terms of loss, because the 2.9B model is trained with large learning rate
1.5
3.0
(a) 2.9B
(b) 19B
Figure 3: Training loss curve. We train 2.9B model with 500B tokens, and 19B models with 200B tokens.
In Table 2, we can find that KV shifting attent KV Shifting introduces bias that is more suitable for language mod- eling, accelerating the convergence of the model. And it seems that both 2.9B models are converge on the bench- mark Lambda, and KV shifting attention can achieve better results. Here we slightly argue that KV shifting attention can achieve better performance than the vanilla model when they both converge, although may take several TB data for 2.9B model converge.
4.4. Scaling Experiment
Firstly, we plotted the training loss curves at the 1.5B, 6.7B, and 13B parameters as shown in Figure 5. We found that un- der different settings, KV shifting attention achieved better results compared to vanilla model.
4.3. Robust Experiment
To verify the robustness of the model, we used different ran- dom seeds for 1.5B model, specifically, we set random num- ber seeds for model initialization and data sampling, and conducted five experiments. The Vanilla and KV shifting attention in each experiment using the same re initialization and data. As shown in Figure 4a, although the training loss is quite shaky, it can be seen that under different random seeds, KV shifting attention is always better than vanilla.
Afterwards, we also followed Kaplan et (2020) and used WebText (Radford al.) as the validation set to draw scaling law for vanilla and KV shifting attention. As shown
"To confirm this, we attempted 5 random seeds under the condi- tion 1eR=le-2, and all results of each experiment are Vanilla divergence and KV shifting convergence.
We adopt large learning in practice, because Lohacheva et al. suggests large learning rates improve generalization
6

---

KV Shifting Attention Enhances Language Modeling
5th token is important in language modeling. If K shifting or V shifting is not used, the model needs to two layers of attention to implement the function. And we provide another way to explain the role of K shifting and V shifting in Appendix 3.3.
(a) 1.5B Parameters
(b) 6.7B Parameters
(c) 13B Parameters
In addition, we also performed longer shifts on K and V. extending the shift between i and i to l 2, i] and [i 3, 1010. The results are shown in Figure 7b. The re- sults have shown that using longer shift window does not improve performance, although increases computational complexity. 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 lan- guage modeling. And we also conducted a shift operation on all QKV in Appendix G and H.
Figure 5: Training loss comparison between different size. All models are trained on 10B tokens. The batch size for 1.5B and 6.7B model is 0.5M, for 13B is 1M, so the total steps of 1 13B model is half of others.
in Figure 6a², KV shifting attention also has excellent seal- ing properties and outperforms baseline at every parameter scale. And we continued to train the 1.5B model to 30B tokens, which we calculate the validation set loss every 1000 steps, but the difference in validation loss still does not decrease, as shown in Figure 6b.
3.0
5. Discussion
Overall, based on previous research on the induction heads mechanism in transformer model, we have designed an attention mechanism that is more suitable for learning in- duction heads. This can not only reduce the demand for width, but also reduce the demand for layers to learn induc- tion heads. At the same time, in large-scale language model pre training, the use of KV shifting attention can achieve better results than baseline in many experimental settings, which implicitly demonstrates the importance of learning induction heads for language modeling.
130
(a) Scaling law
(b) 1.5B Parameters
Figure 6: Validation loss across different size and training tokens. For scaling law, while others this paper means the total parameters. models are trained on 1OB tokens and calculate the final checkpoint's validation loss.
4.5. Further experiments
Another noteworthy fact is that the previous research (El- hage et al., 2021: Akyürek et al., 2024) think transformer to outperform LSTMs at in-context learning on text, with induction heads as a major explanation. Our experiments indicate that existing transformer mechanisms still have not fully unleashed the potential of learning induction heads. By making slight modifications to the original ar- chitecture, the model can learn induction heads much faster than before. When the model can learn induction heads in a simple way, it may becomes easier to achieve length gener- alization (Zhou et al., 2024). And it may be interesting to delve into the underlying mechanisms of transformer, iden- tify important features, and modify the transformer to make it easier to implement important functions. For example, it is very challenging to learn parity check (Wies et al., 2023).
3.30
4m
(a) Ablation
(b) More shifting
Figure 7: Further experiments are conducted on a 1.5B model, where we trained 10B tokens.
To further validate the effectiveness of the proposed KV shifting attention, we conducted the following ablation ex- periments. Firstly, in the experiment of ablating the shifting of K and V. shown in Figure 7a, we find that the shifting of Key and Value both plays an important role. 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
Due to our lightweight modifications, the KV shifting atten- tion can easily be compatible with existing training and inference frameworks. We believe that there is not only
"KV shifting 1 is the KV shifting we used in our paper. For KV shifting 2, we randomly initialize an . az from U 1), and then use 1 on az as ag. It is similar for 3 and KV shifting 3.
In this figure, the parameters is the no-embedding parameters, followed by Kaplan al. (2020), while others in this paper mean total parameters.
7

---

KV Shifting Attention Enhances Language Modeling
one way for models better learn induction heads. But some complex operations may be difficult to adapt to under existing training or inference acceleration frameworks.
interpretabilityy. If we start from making the model more in- terpretable, we can also impose some constraints (Friedman et al., 2024), but this may not improve the performance.
From the perspective of operation, the fusion of adjacent information in neural network has a long history, such as Wu et al. (2018) in vision, Zhang et al. (2021) in video, Li et al. (2023b) in speech, and Cobbe et al. (2021) in text. Many of them fuse or shift all information rather than KV. How- ever, existing LLMs may actually easily integrate adjacent token information without introducing additional inductive biases due to position embedding and causal mask. Some works (Chou et al., 2024) also believe that such operations will enhance local interactions. However, as shown in the Appendix F.3, we conduct a 3-gram experiment and the ad- dition of KV shifting does not improve the learning ability of 3 3-grams, which also implies that the key to the effec- tiveness of KV shifting attention is not to promote the fusion of local information which is usually considered an important role of convolution, but to achieve better induction heads by decoupling key-value pairs in time sequence, thereby enhancing in-contextual learning. The fusion of information is easy, but the decoupling of informa- tion may be difficult. We believe our work can enhance the understanding of the use of shift operations in attention.
6. Related works
Induction heads, introduced by Olsson et al. (2022); Elhage et al. (2021), are specialized attention mechanisms that identify patterns in sequences, enabling large language models to predict subsequent tokens based on previous patterns. There are numerous studies (Bansal et al., 2023; Conmy et al., 2023; Wang et al.; Ren et al., 2024; Todd et al.) starting from induction heads to investigate the interpretability of large models, in order to enhance their transparency. Another part of works (Bietti et al., 2024; Wang et al., 2024; Sanford et al., 2024b; Chen et al., 2024) are to study how the model learns induction heads from theoretical perspective. Previous studies motivated us greatly, and the focus of our work is to modify the attention to make the model better learn induction heads.
Induction Heads
Model Structure and Language Modeling Over the past many years, people have attempted various structures to enhance model's language modeling abilities. From early RNN (Mikolov et al., 2010) and LSTM (Sundermeyer et al., 2012) to the dominant transformer (Vaswani, 2017) today. Transformers have quadratic complexity, and many efforts have been made to improve them, such as RWKV (Peng et al., 2023), Mamba (Gu & Dao, 2023), RetNet (Sun et al., 2023). However, transformers still have many excellent properties that cannot be replaced temporarily, especially their ability to retrieve information (Jelassi et al., 2024; Alman & Yu, 2024). Currently, popular language models (Achiam et al., 2023; Dubey al., 2024; Yang et al., 2024a) still use Transformers as architecture. Recently, there have been also some efforts to modify transformers to enhance modeling capabilities, such as reducing the noise of atten- tion (Ye et al., 2024a) or reducing attention to unneeded elements (Leviathan et al., 2024). Our work is slightly modify attention to enhance its ability to learn induction heads, which potentially improves language modeling.
The shifting operation can be seen as short convolution with kernel size 2. We can see that there are some inter- esting connections with another research line in languages modeling, which the extensive use of short convolution structures in non transformer architecture such as RWKV (Peng et al., 2023), H3 (Fu et al., 2022), Hyena (Poli al., 2023), Mamba series (Gu & Dao, 2023; Dao & Gu), XLSTM (Beck et al., 2024), and DeltaNet (Yang et al., 2024b). The short convolutions can be a component in improving per- formance, especially in enhancing their in-context learning ability (Yang et al., 2024b). We believe that our work can provide some insight, such as performing short convolutions on queries may have small contribution to in-context learn- ing, one possible alternative is to set the kernel size of the query's 's short convolution to key's -1 to simulate the learn- ing process of n-gram induction heads. However, from the perspective of induction heads, it is difficult to explain the role of short convolutions, because generalized state space models (GSSMs) cannot arbitrarily retrieve from context (Jelassi et al., 2024)." 1I We can try using a hybrid model if we really want to perform it, and whether it is necessary to apply short convolutions in the non transformer part of the hybrid structure in the presence of attention layers that are proficient in retrieval may be an interesting question.
7. Conclusion
In this work, we analyzed that induction heads have certain requirements for both the width and depth of the transformer, which motivated us to implement the induction mechanism more efficiently. For stronger expressiveness and faster con- vergence for the learning induction heads, we propose the KV shifting attention, enhancing the language modeling ability of the decode-only structure transformers. We con- ducted extensive experiments to verified the effectiveness of KV shifting attention, including pre training of 2.9B and 19B parameter models. We hope this work can provide some inspiration for achieving more powerful language modeling
Besides, using KV shifting may be much easier to locate the induction heads, which can be helpful for mechanistic
For fixed state space, it can perfectly achieve ...[S][B)... [S] >> [B], where [S] is a special token rather arbitrary token.
8

---

KV Shifting Attention Enhances Language Modeling
or enhancing understanding of transformers.
ing in transformers. arXiv preprini arXiv:2409.10533 2024.
Impact statement
Chou, Y., Yao, M., Wang, K., Pan, Y., Zhu, R., Zhong, Y., Qiao, Y., Wu, J., Xu, B., and Li, G. Metala: Unified optimal linear approximation to softmax attention map. arXiv preprint arXiv:2411. 2024.
This paper presents work whose goal is to advance the field of langugage modeling. There are many potential societal consequences of our work, none which we feel must specifically highlighted here.
Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv: 1803005577 2018
References
Achiam, J., Adler, S., Agarwal. S., Ahmad, L., Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S.,
Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L.,, Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10. 14/68, 2021.
Anadkat, S., al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
Ainslie, J., Lee-Thorp, J., Jong, M., Zemlyanskiy, Y., Lebron, F., and Sanghai, S. Gqa: Training generalized multi-query transformer models from multi-head check- points. In Proceedings of the 2023 Conference Em- pirical Methods in Natural Language Processing, pp. 4895-4901, 2023.
Conmy, A., Mavor-Parker, A., Lynch, A., Heimersheim, S., and Garriga-Alonso, A. Towards automated circuit dis- covery for mechanistic interpretability. Advances in Neu- ral Information Processing Systems, 36:16318-16352, 2023.
Akyürek, E., Wang, B., Kim, Y., and Andreas, J. In-context language learning: Architectures and algorithms. In Forty-first International Conference on Machine Learn- ing. 2024.
Crosbie, J. and Shutova, E. Induction heads as an essential mechanism for pattern matching in in-context learning. arXiv preprint arXiv:2407.07011, 2024.
Alman, and Yu, H. Fundamental limitations on sub- quadratic alternatives transformers. arXiv preprint arXiv:2410.042711 2024.
Dao, T. and Gu, A. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. In Forty-first International Conference on Machine Learning.
Arora, S., Eyuboglu, S., Timalsina, A., Johnson, L., Poli, M., Zou, J., Rudra, A., and Ré, C. Zoology: Measuring and improving recall in efficient language models. arXiv preprint ar arXiv:2312.0499 02.92922. 2023.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The llama 3 herd of models. arXiv preprint arXi::2407.211 2024.
Bansal, H., Gopalakrishnan, K., Dingliwal, S., Bodapati, S., Kirchhoff, K., and Roth, D. Rethinking the role of scale for in-context learning: An interpretability-bee -based case study at 66 billion scale. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics Vollme 1: Long Papers), pp. 11833-11856, 2023.
Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., DasSarma, N., Drain, D., Ganguli, D., Hatfield- Dodds, Z., Hernandez, D., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. A math- ematical framework for transformer circuits. Trans- former Circuits Thread, 2021. htts::/transformer
Beck, M.. Pöppel, K., Spanring, M., Auer, A., Prudnikova, O., Kopp, M., Klambauer, G., Brandstetter, J., and Hochreiter, S. xlstm: Extended long short-term mem- ory. arXiv preprint arXiv:2405.045 2024.
Friedman, D., Wettig, A., and Chen, D. Learning trans- former programs. Advances in Neural Information Pro- cessing Systems, 36, 2024.
Bietti, A., Cabannes, V., Bouchacourt, D., Jegou, H., and Bottou, L. Birth of transformer: memory viewpoint. Advances in Neural Information Processing Systems, 36, 2024.
Fu, D. Y., Dao, T., Saab, K. K., Thomas, A. W., Rudra, A., and Re, C. Hungry hungry hippos: Towards language modeling with state space models. In The Eleventh Inter- national Conference on Learning Representations, 2022.
Chen, S., Sheen, H., Wang. T., and Yang. Z. Unveiling induc- tion heads: Provable training dynamics and feature learn-

---

KV Shifting Attention Enhances Language Modeling
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00772. 2023.
Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma,
N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen,
A., Conerly, T., Drain, D., Ganguli, D., Hatfield-Dodds, Z. Hernandez, D., Johnston, S., Jones, A., Kernion, J., Lovite, I.., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. In-context learn- ing and induction heads. Transformer Circuits Thread, 2022. https:/transfommr--rccca htt.://mrrrrrrrrruuuuuuu//////on pub/2022/nn-comexx suadd-iidductioo-eed///nieee html.
Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Measuring massive multitask language understanding. International Conference on Learning Representations.
Hendrycks, D., Burns, C., Kadavath, S., Arora, A., Basart, S., Tang, E., Song, D., and Steinhardt, J. Measuring math- ematical problem solving with the math dataset. arXiv preprint arXiv:2103.088 0000:200330387, 2021.
Paperno, D., Kruszewski, G., Lazaridou, A., Pham, N.-Q.. Bernardi, R., Pezzelle, S., Baroni, M., Boleda. G., and Fernández, R. The lambada dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th Anmual Meeting of the Association for Compu- tational Linguistics (Volume I: Long Papers), pp. 1525- 1534, 2016.
Henry, A., Dachapally, P. R., Pawar, S. S., and Chen, Y. Query-key normalization for transformers. In Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 4246-4253. 2020.
Jelassi, S., Brandfonbrener, D., Kakade, S. M., et al. Repeat after me: Transformers are better than state space models at copying. In Forty-first International Conference Machine Learning, 2024.
Penedo, G., Kydlícek, H.,, Allal, L. B., and Wolf, T. Fineweb: Decanting the web for the finest text data at scale. Hug- gingFace. Accessed: Jul, 12, 2024.
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint arXiv:2001.0834
Peng, B., Alcaide, E., Anthony, Q.. Albalak, A., Arcadinho, S., Biderman, S., Cao, H., Cheng, X., Chung, M., Grella, M., et al. Rwkv: Reinventing mns for the transformer era. arXiv preprint 2rXiv:2305.1300 2023.
2020.
Leviathan, Y., Kalman, M., and Matias, Y. Selec- tive attention improves transformer. arXiv preprint arXiv:2410.02703, 2024.
Poli, M., Massaroli, S., Nguyen, E., Fu, D. Y., Dao, T., Baccus, S., Bengio, Y., Ermon, S., and Ré, C. Hyena hierarchy: Towards larger convolutional language models. International Conference on Machine Learning, pp. 28043-28078. PMLR, 2023.
Li, H., Zhang, Y., Koto, F,, Yang, Y., Zhao, H., Gong, Y., Duan, N., and Baldwin, T. Cmmlu: Measuring mas- sive multitask language understanding in chinese. arXiv preprint ar Xiv:2306.092. 12, 2023a.
Press, O., Smith, N., and Lewis, M. Train short, test long: Attention with linear biases enables input length extrapo- lation. In International Conference on Learning Repre- sentations.
Li, Y.. Wu, Y., Li, J., and Liu, S. Accelerating trans- ducers through adjacent token merging. arXiv preprint arXiv:2306. 2.0,,,9,06.,000, 2023b.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, Language models are unsupervised multitask learners.
Liu, A., Feng, B., Xue, B., Wang, B., Wu, B., Lu, C., Zhao, C., Deng, C., Zhang, C., Ruan, C., et al. Deepscek-v3 technical report. arXiv preprint arXiv:2412.1944 19437, 2024.
Rajaraman, N., Bondaschi, M., Ramchandran, K., Gast- par, M., and Makkuva, A. V. Transformers on markov data: Constant depth suffices. arXiv preprint 24077 17686, 2024.
Lobacheva, E., Pokonechny, E., Kodryan, M., and Vetrov, D. Large learning rates improve generalization: But how large are we talking about? In NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning.
Ren, J., Guo, Q. Yan, H., Liu, D., Zhang, Q., Qiu, X., and Lin, D. Identifying semantic induction heads to understand in-context learning. arXiv preprint arXiv:2402.1305: 2024.
Lozhkov, A., Ben Allal, L.. von Werra, L., and Wolf, T. Fineweb-edu, May 2024. URL
https: http///ggggggncceeeeee/ dataaattsss///////
Mikolov, T., Karafiát, M., Burget, L., Cernocký, J., and Khudanpur, S. Recurrent neural network based lan- guage model. In Interspeech, volume 2, pp. 1045-1048. Makuhari, 2010.
Sakaguchi, K., Bras, R. L., Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106 2021.
10

---

KV Shifting Attention Enhances Language Modeling
Sanford, C., Hsu, D., and Telgarsky, M. One-layer trans- formers fail to solve the induction heads task. arXiv preprint arXiv:2408. (eto:2408.1333333 14332, 2024a.
Weber, M., Fu, D. Y., Anthony, Q. Oren, Y., Adams, S., Alexandrov, A., Lyu, X., Nguyen, H., Yao, X., Adams, V., et al. Redpajama: open dataset for training large language models.
Sanford, C., Hsu, D., and Telgarsky, M. Transformers, parallel computation, and logarithmic depth. In Forty-first International Conference on Machine Learning, 2024b.
Wies, N., Levine, Y., and Shashua, A. Sub-task decom- position enables learning in sequence to sequence tasks. In The Eleventh International Conference on Learning Representations, 2023.
So. D., Mankee W., Liu, H., Dai, Z., Shazeer, N., and Le, Q. V. Searching for efficient transformers for language modeling. Advances in neural information processing systems, 54:6010-6022, 2021.
Wu, B., Wan, A., Yue, X., Jin, P., Zhao, S., Golmant, N., Gholaminejad, A., Gonzalez, J., and Keutzer, K. Shift: A zero flop, zero parameter alternative to spatial convolu- tions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 9127-9135, 2018.
Song. J., Xu, Z., and Zhong, Y. Out-of-distribution general- ization via composition: a lens through induction heads in transformers. arXiv preprint arXiv:2408. 09503, 2024.
Xu, M., Men, X., Wang, B., Zhang, Q.. Lin, H., Han, X., and weipeng chen. Base of roPE bounds context length. In The Thirty-eighth Annual Conference on Neu- ral Information Processing Systems, 2024. URL https: / / //peneevieew net forum? SerruumiddEEieeleeee
Su, J. Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 127033, 2024.
Sun, Y.. Dong, L., Huang, S., Ma, S., Xia, Y., Xue, J., Wang, J., and Wei, Retentive network: successor to transformer for large language models. arXiv preprint arXiv:2307.08621 2023.
Yang, A., Yang. B., Hui, B., Zheng, B., Yu, B., Zhou, C., Li, C., Li, C., Liu, D., Huang, F,, et al. Qwen2 technical report. arXiv preprint 1xXiv:2407.10611 2024a.
Sundermeyer, M., Schlüter, R., and Ney, H. Lstm neural networks for language modeling. In Interspeech, volume 2012, pp. 194-197, 2012.
Yang. S., Wang, B., Zhang, Y., Shen, Y., and Kim, Y. Par- allelizing linear transformers with the delta rule se- quence length. arXiv preprint arXiv:2406.0644
2024b. Ye, T., Dong, L., Xia, Y., Sun, Y., Zhu, Y., Huang, G., and Wei, Differential transformer. arXiv preprint 2024a.
Todd, E., Li, M., Sharma, A. S., Mueller, A., Wallace, B. C., and Bau, D. Function vectors in large language models. In The Twelfth International Conference on Learning Representations.
arXiv:2410.0522
Ye, T., Xu, Z., Li, Y., and Allen-Zhu Z. Physics of language models: Part 2.1, grade-school math and the hidden rea- soning process. arXiv preprint arXiv:2407.2031/ 2024b.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine- tuned chat models. arXiv preprint 000.02230..09288, 2023.
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can a machine really finish your sentence? Proceedings of the 57th Annual Meeting of the Asso- ciation for Computational Linguistics, pp. 4791-4800, 2019.
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Waleffe, R., Byeon, W., Riach, D., Norick, B., Kor- thikanti, V., Dao, T., Gu, A., Hatamizadeh, A., Singh, S., Narayanan, D., et al. An empirical study of mamba- based language models. arXiv preprint arXiv:2406.07885 2024.
Zhang, H., Hao, Y., and Ngo, C.-W. Token shift transformer for video classification. Proceedings of the 29th ACM International Conference on Multimedia, pp. 917-925, 2021.
Zhou, H., Bradley, A., Littwin, E., Razin, N., Saremi, O., Susskind, J. M., Bengio, S., and Nakkiran, What algorithms can transformers learn? a study in length generalization. In The Twelfth International Conference on Learning Representations, 2024. URL https:// openreview net htt///ooumm/m// d=Ass wwd::=ssssssaaa....... TuHnmHX.
Wang, K. R., Variengien, A., Conmy, A., Shlegeris, B., and Steinhardt, Interpretability in the wild: a circuit for in- direct object identification in gpt-2 small. In The Eleventh International Conference Learning Representations.
Wang, M., Yu, R., Wu, L., et al. How transformers imple- ment induction heads: Approximation and optimization analysis. arXiv preprint arXiv:2410. /474, 2024.

---

KV Shifting Attention Enhances Language Modeling
A. Proof of Theorem 3.2
Proof. We consider two-layer single-head transformer without FFN, where the first layer has the residual block, while the second layer does not have the residual block.
We first embed each token into RD as
0 and take W
0 Idxd
0
0 ) then the s-th output token of the first layer
Fs
Us )
Eg (-p(2) - +)) x+ )
- softmax
Then for the second layer, we choose p!²) = m,
w (7(")=((0 0
-((00 (
laxd 9) w. (2.1)
( w* w
(Toxd
0) E T DxD
0
and the projection w:O) (laxd: Odxa) RRdD Then the last output token of the second layer is
softmax(x[us/o ml I. s)ssss
By Lemma D.1, for any L E N+
IH₂
TF 1.00
sup X1
IH( - TF- TX
mll 1 s)).ra
o0tmax(=[=s-1100 - m/L 81)r,
< 0 - m s|}
2
softmax
x-- (e
m3 s|)1
< 2sup 1 v/Vv// x IIIz-1/01
< 2|1x /0/ sup 11 xz-1 00
2 1// sup
softmax
(s x+
-3-1
x-1 softmax ((s -11)
< <2||1/o 11,1 sup
66-
==1
= -
sup
Σラ=0 --(())
e p(1)
ply S 0 (e-(")).
</1|T/ 1,1
B. Proof of Theorem 3.3
we set a1 = 0, a2 = 1, 3. = 1/a, B2 = 0, = m (the bias of Alibi position embedding), W, = W; = W, = W,
E Rddd we can get the proof.
12

---

KV Shifting Attention Enhances Language Modeling
C. Proof of Theorem 3.4
We consider embeddings un E Rd with i..d. Gaussian (o 4) entries.
We recall a few facts:
(Norm) We have uf 1== =1 O(1VVd) Because ui u, is scaled chi-squared distribution, with mean 1, and variance 2/d.
INear-orthogonaliyy For i + j. we have u uj = O(I/Vd). To see this, denoting u, = dds12(ūxx)) ). where üي are the normalized entries 24, note that we have
Vdu, új
Ush Üzh - (0,
Vd
by the central limit theorem, since for each k, the quantities ūs* are zero-mean, unit-variance i.i.d. random. Although we can describe it more accurately, in (10) we actually only used when d tends towards positive infinity, u, us = 1 and up Uj
... Tii I4+1, my UT, 3T+1, where 27+1 = 24, and 2 22. .... 2T are different." Now let's Assume the sequence 21, .2, start our proof:
Calculate attention score The last token attend to the t token, the attention score å before softmax can be calculate as:
If k = i, ag.- = x2_ (a)0; + 02 = a + ( ( );
If k i 1, àz 27 (a1 + a24;) = 02 + +0artt aa );
If k T + 1, az x.. (a12) + (2 01 + 0
(1)+ ):
v/
Other else, åz = x._ (01 + 02 ( di a1 ). vd
So, the attention score after softmax can be calculate as: If = i, ag = e0e "1 3/5;
If k = i 1, az = E0e 21t02)
)/S:
21t02 1/S;
If = T + 1, = c O(e
vd
01402 vi
1/S, where S = 2 Ole 21t02
+ e 21+02
21t02 vi
Other else, ag= = O(e
+ +OTe
Calculate logits
Now we calculate the logits of predict Ath token, donate as l;;
= -i=
I z. (8 + 82 = (31 + 0 (i ))a: + (82 + 0 8 1
)) + (81 + +
Da++ +
vd
O(T)O( f B
Vd O
1/S = (2 + 2e 0 (But
) B2 + e 81 f1
+ 0 -لم
))0 - 3/5.
vd
(Expand according to k=i, i+1, T+1, other else.)
If W' = i 1, 4 = Σ - 41 (Bizk + B2 = 0 )a + (8) + 0 0 1a0+11 + (32 +
0(5 va 11a;42 + 0 |ar+1 + O Bi (1)t va
)0(e
4te0 3/5 (2e 0 01
82 ) + _2 + -² 20 B1 (1+03) ) + B2+
va
vd
0 ((( 1)0(e =1/22)
1/8 (Expand according to k=i, i+1, i+2, T+1, other else.)
If W' = T. In
ΣT TB1Ik + B3Zk-1)aR 0 3 vd
)a; + O ( )ait1 + (3) + 0 06 )ar +
vd
Va
(B2 + 0 ((
DaTH1 + O(T)O O (1++ ()
Bit Va
va
//S
( ( + va0(Site + Bi + 3ze +
Vd
vd
21t02 va
//S (Expand according to k=i, i+1, T, T+1, other else.)
Vd
If K i
lp
Σt (B100 + 1) (81 + 0 01 )))+)))); + (g2 + 0 0
)); +
vid
Here, we assume >1 < T, not, there will slight difference in proof.
13

---

KV Shifting Attention Enhances Language Modeling
O Qitt Jait1 + O( 0 JaT+1 + O(T)O
G+t Vd
(0(e
aL va
3//S (31 + 8ze + 000 f + 2 0((0t()))))
Vd
o ( 100(e vi "va=))/ 1/S5 (Expand according to k i+1, T+1, other else.)
Other else, lN Eker of 1204-1)as
0 0 )a, + 0 )as+1 (B 0
v/ )ak +
va
vil
Vi (B2 + (()t)))))a + O )0(e va 3/S
(E (
+ c020(
(1++22)))
+ 31 + B2 +
vd O (8t (11 ))0(e "2(2)) // (Expand according k=i, i+1, W k'+1, other else.)
vd
vd
vi
As d tending towards positive infinity, we summay the logits of 7th token in Table 3:
predict tokens
logit (31 + 2/S
i-1
i
(2 + 8/5
i+1
(3 32)/8
T
(31 + a2e 1/8
other else
(31 + B2)/S
Table 3: Logits summary, when tend to sc, where 2e + OTT)
Calculate loss
We denote em /"/S as a1, e²2 /S S as 02, than we can get:
c/4t1
Loss = -log(
(13)
2/m1
0201+00/3
= =llog(
(14)
t 20 :22 + O(T)
D. Limitation
Due to limitations in computing resources and other limitations, many experiments cannot be repeated many times, and it is not convenient for us to use open-source datasets or open source our datasets. But our robustness experiments have convinced us that KV shifting attention can achieve better results in many experimental settings, and we have also open-source our pertained model. 13 In theory, we provide a way for KV shifting attention learn induction heads under relaxed conditions. However, remains challenging on how learn induction heads under multi-layer transformer.
Mcanwhile, induction heads can actually be implemented in another complex way. For example, the model can focus on similar tokens, obtain the hidden location information of the 5th position, convert this information into the information of the (i+1)th position through MLP, and then use it as query information in the next layer, so as to focus on the information of the (i+t 1)th token. is currently unclear how this kind of induction heads affects the training dynamics the model. However, the key attention is allowed to shift in the next layer, there is no need for MLP to convert 5th position information into the 1)th, which maybe also be helpful for the learning induction heads.
E. Addition experiments to exploration the capability boundary
E.1. Hop k
In this section, we evaluate the multi-layered form of induction heads, namely multi-hop. We followed the code of Sanford et al. (2024b) and only changed the attention of vanilla to KV shifting. The experimental results are shown in Figure 8. Obviously, KV shifting has a very good bias for learning multi-hop.
Taking the pink line (L=5) as an example, the performance of the vanilla model will significantly degrade when the hop number exceeds 8. However, the KV shifting attention still has a small error when the hop count reaches 16. This powerful ability to perform implicit reasoning implies that KV shifting attention may achieve better results in mathematical or reasoning abilities. Therefore, we present the experimental results of mathematical ability in the next section.
There is open source implements of KV shifting attention. http::/gihuuhcoo/rrgggl/aa/ KV shifting attention, using their code and spending NVIDIA-A100 plus hours of training time is good choice.
If you want verify the effectiveness
14

---

KV Shifting Attention Enhances Language Modeling
L-ayer transfo mar taken-wis class ficatien emror on hop.
L-byer transformer talgn-wise clasoffication errer on hop.
L-byar frai aoroee token-abse classification erro on hop.
0.-
8.-
0.-
03
1-66
0.31
0.2
62
0.2
02-
(a) Vanilla (Origin)
(b) Vanilla (Reproduce by us)
(c) KV shifting attention
Figure 8: Comparison of Error Rates under Hop k Tasks. The smaller the error, the better the performance.
E.2. Grade-School Math
As more direct manifestation of induction ability, learning math problems is natural experiment. At the same time, in order to eliminate the influence of complex syntax, test data leakage, etc. We follow Ye et al. (2024b) and conduct experiments on the iGSM dataset. Due to the code of Ye et al. (2024b) haven't been released yet, we start from code of We use the similar 1024 for all experiments and the learning rate which we set as 2e-4.
except the context length which we set as
In addition, as we are using an open-source code implementation. For simplicity, there no final summary of the answer. So the evaluation of accuracy is based on the fact that 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.
Table 4: Experiments on iGSM. Tr X Te Y means train with the numbers of operation no greater than X and test with the numbers of operation as Y.
MODEL
TR12-TE15
TR21-TE:2
VANILLA
0.8154
0.8711
KV SHIFTING
0.8909
0.9062
The result in Table shown the enormous potential of KV shifting attention. We expect KV shifting attention enhance the reasoning ability of the model by improving its ability in basic induction heads.
We have included this section in the appendix because our experiment was not as thorough as Ye et al. (2024b). We only tested the accuracy without delving deeper into the analysis, such as the reasons for mistake like Ye et al. (2024b). Besides, conducting more dn-tepth experiments is beyond the scope of this article.
E.3. Can KV shifting attention learn n-gram better?
In addition to induction ability, an important part of language modeling is n-gram In this section, we test the model's ability to learn n-grams. We randomly generated approximately 200 pairs of xı and 22, and then we randomly generated 23 for each pair. The model needs to accurately predict 23 when seeing I and 22. The results are shown in Figure 9, Obviously, the KV shifting attention do not enhance the model's ' learning ability for this 3-gram tasks, but it also does not weaken the learning ability for this 3-gram tasks.
We must emphasize here that, 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 Markova data can actually be simulated with just one layer transformer with vanilla attention (Rajaraman et al., 2024). Therefore, the performance of Figure 9 is completely different from that of Figure la.
II our experiment, if 1r=2e-3 is used, Vanilla's performance will be quite poor (In Tr12-Te15, Vanilla will get approximately 0.5, while KV shifting will get approximately 0.7). For parameter tuning is not the most important thing have at hand, will try to find the most suitable hyper-parameters for each model in the future.
15

---

KV Shifting Attention Enhances Language Modeling
10
10
cliting
koroy 06
06-
03
00
300
17300
7148
7003
7530
Fpochs
Epo2ts
(a) 50M
(b) 0.4M
(c) 0.8K
Figure 9: Accuracy of learning 3-gram text using models of different sizes. In this experiments, there are 50M parameters model with 4 layers, 0.4M parameters model with 2 layers, 0.8K parameters model with layer.
Another noteworthy thing that intuitively, the shift operation on KV is beneficial for the fusion of local information, but the KV shift attention that integrates the previous token and current token does not contribute to the learning of 3-grams. This indicates that from the perspective of fusing local information, KV shift attention does provide additional contributions, and perhaps position encoding is already sufficient. Perhaps it's because the bottleneck in learning n-grams is not in attention, but in other areas such as network width or the parameters numbers of MLP, which can be treated as a corollary of Rajaraman et al. (2024).
F. Learnable parameter analysis
One worth studying is how these learnable parameters will change in a pre trained model. We conducted research on the 2.9B model, and due initialization, a and & are independent random variables, but training progressed, they became coupled with each other. As shown in the Table 5, we have counted the number of whether S and B1 S |3 in each KV pair.
Table 5: We caleulate the number of whether a₁ a2 and 31 s B3 in each KV pair in 2.9B model with 500B token, the total numbers is 128.
01
0r
0.1 S 0r
82
50
17
B2
52
As shown in the Table 5, we find that the numbers on the diagonal dominate, which means that KV pairs with same relative size relationships between an and a2. 31 and 82, has become the majority. And we find that the diagonal was already 47,43 when training 20B tokens. This indicates that for most heads, the model tends to use the key and value of the same token.
The heads in the upper right focus on the key of the (i 1)th token to obtain information about the ith token, while the heads in the lower left focus on the key of the tth 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 7th token under the causal mask, as the (i - 1)th token token may contain some information about the (i 1)th token token to some extent. Therefore, the lower left corner will relatively small.
F.I. Different gate for shift
Besides, we find that for the trained 2.9B model with 500B tokens, E, a and Σ, a; is away from 1, although they are I when initializing. However, a common gating mechanism is to use activation functions to control Σ, a, = and Σ, a, = 1 during the training (e.g. Q1 = Sigmoid(a), $1 Sigmoid(b), a2 parameters, call it KV shifting gate). In addition, during initialization, these parameters are all between 0 to 1, but as training progresses, some parameters become negative and far from zero. 15 Do we need to keep the model between 0 to I
a1, 82
a1, where a, E is the learnable
c.gg a1 0.08, 02 0.43, 81 0.34, 92 0.15 the tth KV pairs of the 17th layers. The previous token's 's key dominate the current token's 's key, while subtracting the previous token' value from the current token's 's value.
16

---

KV Shifting Attention Enhances Language Modeling
during the training process (e.g. 0 min(max(a;, 0), B, min(max((;, 0), 1), we call it KV shifting 0 1.)
For this purpose, conducted experiments the 1.5B parameters model with 10B token, shown in Figure 10a. Using more controls does not make the model learn better. Allowing and 3 to have wider range of degrees of freedom may enable the model to learn richer features. For example, Elhage et al. (2021) discovered the presence of "apting prefix-search" heads in vanilla model. Although we don't know what its function is, we restrict ,32 2 0, it is likely to limit generation of this kind heads.
2.30
Venilla
shitting
2.25
KV shifting
KV shitting 0.
Log8
2.15
2.10
2.05
3000
4000
5000
6000
7000
8000
9000
(a) Variant
Figure 10: Further experiments are conducted on a 1.5B model, where we trained 10B tokens.
G. QKV shifting attention?
If we also shift Q in attention, it doesn't 't look like an ablation experiment, after all, it adds an extra part compared to the KV shifting attention. On the other hand, shifting Q is also far from our motivation. So we won't mention this section in the main text. But in order to better compare with some baseline methods which use the similar operation, such as Peng et al. (2023) which we discussed in the discussion section, they can shift all the information of Adjacent tokens. Another highly correlated works Multi-DConv-Hent Attention (So al., 2021), they use architecture search and apply convolution all Query, Key and Value. So conduct shifting experiments on all Query, Key and Value, and the shifting is also per head.
3.5
1.00
ovahiling DVV saiting
13.0
12.0
37.
0
0.3
360
22.3
252
100
110
170
155
1.23
-90
180
170
120
122
120
120
140
12.00
140
170
20
(a) Loss
(b) MMLU
(c) CMMLU
Figure 11: QKV shifting attention vs KV shifting attention in model with 19B parameters.
We trained a model for QKV shifting with 19B parameters. Due to some machine malfunctions, the last checkpoint of our QKV shifting attention saved with 180B training tokens, and we don't plan to continue training until 200B. The loss curve shown in Figure 11. From the loss curve, it seems that the two are similar. Then we evaluated the benchmarks as shown in Figure 11. It can be seen that the performance of shifting QKV is not good as shifting KV, and even not good the vanilla. From the view of induction heads, the shifting of Q is difficult to contribute to the formation of it.
Finally, we will ask myself and answer question when we recall Zhang et al. (2021). In their Figure 3, they try different position for token shift, such as "after Add", "before Norm", "after Norm", "before Add". From the perspective of our article, placing in the place of K and V is the most suitable for forming induction heads. Shifting elsewhere can not alleviate the width requirement of induction heads.
17

---

KV Shifting Attention Enhances Language Modeling
H. More comparison
There is an open source implements of KV shifting attention http:///iibb.cooo//ooo/////aCaC If you want to verify the effectiveness of KV shifting attention, using their code and spending 8 NVIDIA-AI00-80 plus 2 hours of train- ing time is good choice. They use 0.2B model, and train with steps 5100 and batch size 0.5M. The training and validation set is sample from fineweb10B. And we use their code for the following comparison. Due difficulty in compatibility with flash attention, out of memory (OOM) may occur larger scale for some method. So we hwtph//ww..o/mmicooooo/mmmmmbbba fran
use the code in
flashdiff 1.py for Differen- wwwaaaaggnnnnnnnnanaaaaaa ksgk/sslecctiv--atteeniii although it remains OOM when context length is larger than 2048.
tial Transformer (Ye et al., 2024a), And selective attention from http://gihhuu..oo//aag//a//
Note that all experiments here were conducted only once, and we can see that some numbers may be exceptions. However, in almost every learning rate and context length experimented with, KV shifting attention achieved the best results.
Table Vanilla model. Validation loss variation with learning rate and context length.
Learning Rate
Context Length
256
512
1024
2048
4096
0.0001
3.7908
3.7078
3.6786
3.6758
3.6845
0.0002
3.6454
3.5839
3.5400
3.5308
3.5347
0.0004
3.5440
3.4558
3.4133
3.3778
3.3880
0.0010
3.4701
3.3816
3.3244
3.3082
3.2877
0.0020
3.4505
3.3605
3.3080
3.2813
3.2627
0.0040
3.4632
3.3701
3.3373
3.5104
3.2983
Table KV shifting attention. Validation loss variation with learning rate and context length.
Learning Rate
Context Length
256
512
1024
2048
4096
0.0001
3.7689
3.6968
3.6491
3.6402
3.6441
0.0002
3.6272
3.5616
3.5153
3.4975
3,4929
0.0004
3.5308
3.4440
3.3903
3.3686
3,3692
0.0010
3.4603
3.3691
3.3166
3.2838
3,2770
0.0020
3.4381
3.3686
3.2983
3.2618
3.2404
0.0040
3,4766
3,4327
3.3097
3.3109
3.3017
Table QKV shifting attention. Validation loss variation with learning rate and context length.
Learning Rate
Context Length
256
512
1024
2048
4096
0.0001
5.8747
3.7215
3.6750
3.6752
3.6798
0.0002
3.6543
3.5793
3.5301
3.5391
3.5213
0.0004
3.5482
3.4584
3.4190
3.3847
3.3836
0.0010
3.4678
3.3783
3.3307
3.3013
3.2835
0.0020
3.4518
3.3769
3.3081
3.2725
3.2620
0.0040
3.4994
3.3710
3.3063
3.3486
3,2662
18

---

KV Shifting Attention Enhances Language Modeling
Table 9: K shifting attention. Validation loss variation with learning rate and context length.
Learning Rate
Context Length
256
512
1024
2048
4096
0.0001
3.7787
3.7109
3.6567
3.6430
3.6667
0,0002
3,6519
3.5758
3.5356
3.5110
3,5085
0.0004
3.5471
3.4541
3,4037
3.3820
3.3718
0.0010
3.4753
3.3861
3.3240
3.2976
3.2791
0.0020
3.4531
3.3570
3.3038
3.2733
3,2564
0.0040
3,4993
3.3699
3.3080
3.3768
3,2795
Table 10: V shifting attention. Validation loss variation with learning rate and context length.
Learning Rate
Context Length
256
512
1024
2048
4096
0.0001
3.7742
3.7061
3.6698
3.6512
3.6610
0.0002
3.6366
3.5574
3.5253
3,5044
3,5038
0.0004
3.5300
3.4494
3.3980
3.3775
3.3639
0.0010
3.4599
3.3711
3.3234
3.2895
3.2745
0.0020
3.4425
3.3489
3.2915
3.2622
3,2658
0.0040
3.4805
3.3556
3.3798
3.3130
3,4101
Table 11: Differential Transformer (Ye et al., 2024a).
Validation loss variation with learning rate and context length.
Learning Rate
Context Length
256
512
1024
2048
4096
0.0001
3.9601
3.9023
3.8693
3.8336
3.8279
0.0002
3.7360
3.6692
3.6280
3.6149
3.6123
0.0004
3.5909
3.5191
3.4705
3.4462
3,4362
0.0010
5.8422
3.4158
3.3718
3.3356
3.3648
0.0020
3.7052
3.7530
3.7660
4.1150
3.3187
0.0040
7.8280
6.6233
6.4263
3.3948
3,4779
Table 12: Differential Transformer (Ye et al., 2024a).
Validation loss variation with learning rate and context length. We enhance Differential Transformer with query and key normalization(Henry et al., 2020), which is the base setting of http://gittubbcom//nno////Baaa
Learning Rate
Context Length
256
512
1024
2048
4096
0.0001
3.9622
3.8999
3.9122
3.8484
3.8559
0,0002
3.7218
3.6624
3.6245
3.6154
3,6222
0.0004
3.5800
3.5080
3.4703
3.4688
3,4476
0.0010
3.5305
3.4076
3.3840
3.3512
3.3482
0.0020
3.4998
3,3930
3.3478
3.3635
3,3543
0.0040
3.5607
3,4142
3.3566
3.3573
3.3489
19

---

KV Shifting Attention Enhances Language Modeling
Table 13: Selective Attention (Leviathan et al., 2024) Validation loss variation with learning rate and context length. Vanilla model's peak memory consumption is 42584 MiB when context length 4096, KV shifting attention is 43360 MiB. The peak memory consumption of selective attention 69500 MiB when context length is 1024. We will supplement this part of the experiment when selective attention is implemented more efficiently.
Learning Rate
Context Length
256
512
1024
2048
4096
0.0001
3.7922
3.7109
3.6817
OOM
OOM
0.0002
3.6580
3.5842
3.5466
OOM
OOM
0.0004
3,5538
3,4744
3.4240
0OM
OOM
0.0010
3.4768
3.3985
3.3594
0OM
OOM
0.0020
3,4575
3,4100
3.3613
0OM
OOM
0.0040
3,4850
3,4618
3,4128
OOM
OOM
I. Metrics for evaluation
In this section, we conducted MMLU evaluation under three condition, few shot (5 shot), zero shot and cloze (zero shot). The setting of cloze is followed by Waleffe et al. (2024), which intends to break away from the format of standard multiple-choice and directly measure the knowledge. The test results are shown in Table 14. It can be seen that compared to vanilla, KV shifting attention not only enhances the ability of in-context learning, but also accelerates the model's learning of world knowledge. Moreover, the format of multiple-choice questions may be advantageous for KV shifting attention. The model can easily use context to compare the possibilities of various options and select the option with the highest probability.
Table 14: We compare vanilla and KV shifting at 2.9B model with 500B training tokens by using different evaluation metric.
Vanilla
KV shifting
Benchmark
Cloze
Zero
Few
Claze
Zero
Few
MMLU
30.41
33.14
37.26
32.17
37.13
40.88
In addition, under different evaluation metric, KV shifting attention achieved better results, which reflects the robust of KV shifting attention.
J. Compatibility with multi latent attention
Recently, a highly efficient attention method for reasoning, namely multi latent attention (MLA), has achieved good results (Liu et al., 2024). However, due to the fact that the attention calculation of MI. in the inference stage is at the latent level, it is difficult to be compatible with some improvements to MLA, such as QK normalization (Henry et al., 2020), Rotary position embedding (RoPE) Su al. (2024). So much so that their roped vectors were shared across all the heads. But KV shift attention be compatible in this way, that is, we ensure that the key of some the heads in each layer are generated entirely by the previous token, and the value of some the heads in each layer generated entirely by the previous token.
On the other hand, what is the necessity of performing KV Shift on MLA? According to our analysis of attention execution induced heads in Section 3, the width of the model constrains its ability to learn and represent induced heads, while the projection operation of MLA brings the bottleneck of width to the size of the latent space.
We conducted the following experiment with the same basic setup as H. for example, the size of hidden states 768, with a total of 12 heads. We did not perform low rank projection on query. For MLA, we perform low rank projection on key and value to 256 dimensions, and then the dimension with RoPE in key and value is 32, the dimension without RoPE is 32, and the dimension of query is 64. Due to the total number of parameters being about 5% more than in Appendix H, we do not further align the total parameter numbers. For MIA with KV shifting, we shift the key of one head, and the value of another head. We find that Ir=2e-3 is the best for both of them. The results are shown below:
We found that as the sequence length increases, the overall performance of shift KV is also better than MLA. This also indicates the enhancement of the model's 's ability to utilize context.
20

---

KV Shifting Attention Enhances Language Modeling
Table 15: Enhance MLA with KV shifting.
Validation loss
Method
256
512
1024
2048
ML.A
3.5098
3.415
3.3608
3,3188
3.2976
MLA KV shifting
3,4987
3.4069
3.3414
3.3002
3,2753
K. Python script for generating induction data
This is Python script that generates induction data. The model will randomly select number from mid val to max val as token. this token has already appeared, the next token will be the same as the previous one. When the sequence length is less than 512, 0 will be added. there is no induction data the sequence, it will be regenerated. The predict positions returned by this function are the positions where the accuracy of the induction is calculated.
def generate_arreeeeee 512, min_val-1, mid_val=10, max max_val-8000)
lis list (range (mid_val+1, max_val)) random random.shuffles shuffle(lis)
1is [is[:length]
array = []
predict_positions []
di ()
while len (arrayy) <lenghh:
x candem.choics (1is)
if (len -- 0) (x!=array[-1])=
if 2array)))3:
in di: if len [alarray|-1]] =
di [x] -1 array appay.append (x)
else:
predict_positions.appee (1en (len (array))
array (x)
apray. (d)(x)
array [0] (arraay))) predict_positions [0:1]
len cereddcc__psittions 0:
return generate_arra, (length, min_val, mid_val, max_val)
21

---

KV Shifting Attention Enhances Language Modeling
L. Python code for drawing streamline diagram
provide Python code for drawing streamline diagrams in Figure 2.
import torch
import numpy as np
import matplotlib.pyple. as plt
def (alppa, beta, = 0):
alphal alpha
alpha2 1 alpha
betal beta
beta2 = 1- beta
al txrch.exp (2-orrhh. (alphal) +torch. (alpha2) +ot)
a2 corch.exp(alppe (ataaaapaaaa//ertreaaaa
wxt.orrcc....pppppaaaaaa.aa +torch. exp (alpha2) +ot)
(2*torch,
target a2+betal (beta2/S)
nol (betal/S) + beta2+a2
no2 2+al+ber +betal + a2+beta2
exp_ta targttrggttt torch, exp (target)
exp nol torch. (nol)
exp no2 = torch.expi (no2)
exp at torch. (ot)
denominator exp_target
exp_nol + exp_no2 exp_ot
prob_ratic exp_target / denominator
loss = -torch. (proo_ratto))
return less
x,
(0, 1, 1000), np. (0, 1, 1000))
X torch.tensor (x, requires_grad-Iruee)
Y = torch.tenser
Sequires_grad-TTiee
EX,,Y,0.00000 0.00001,01
8. sum () ())ckward
grad_X x.grad
grad_Y Y Y.grad fig, plt..ubplots()
strm ax.streemplot
d.detach (umpy())nm)))).. (), Y.detach (umpp(()
grea__..etacch() numpy (), -grad_ detach numpy ()
levels np.arange 10, 0.05) contours - cot.oontour(
x.detach () numpy (), y.detach () ()rupyy()
2.detach numpy (), levels=levels, colors=/ black'
set xlabel (r' $\alpha 1$()
ax. ax.set_ylabed (x' $\beta_1s()
plt. stows)
22

---

KV Shifting Attention Enhances Language Modeling
M. Pytorch code for KV shifting attention
provide following Python code that can easily implement KV shifting attention with rotary embedding. this example, we used convolution operation perform shifting operations.
from torch import
import torch.nn functional as from flash_attn import flash_atth_fum
def custcon_convelltionn (U, K):
bs, seq, h d U. shape
h, s.shape
padding tw 1, 01
U_padded = F (Ud (0, D, 0, 0, *padding))
(bs, seqtu-1, h, d)
U_unfolded U_padded unfold (1, 1) I (bs, seqtN-1, h, d w)
K_expanded unsqqueeze (0) unsqueeze (0) unsqueeze (bs, seq, d,
(1, 1, h, 1, w)
V_unfolded U_unfolded expanded V V_ (dim=-1) # (bs, seq, h, d) return V
def
torch wtw.:anwwsssll..nuu.......eeee num_kv_hheadss
tarch. wwwpo//see....m.l_____a www_kyu 1)
self.k nn. nn.Paremete (torch.cat cat ([K, 1-K] dimm=1)m
self.v nn. Parameter (torch.cat ([V, 1 1-v], dim=]]
def
foward (self,
custom convolution (k, self.K)
v custem cwstcm__cnnse wow.:mooooooooooooooooooo (v, self.V)
q, k= relarrotary_emit emb (q, k, seqlen
attn_outputs = funs
4 v, causal-True
The following code can be used for inference:
None:
yelf.last kak k,k k[:,-1:) eustem_ coniooutcnvooiiionnnoooo (k, self.k) self.last_vi v, v= -[:],
(v, self.V)
else:
self last_k,kk k, self.K(;, :1] sell..ast_k self.K[: 1:]*k
self.last_v,V v, self.V(;, :1] +slff last_v self.V(;,11 1:]
23

---

KV Shifting Attention Enhances Language Modeling
N. Experimental setup details
conducted toy models for induction heads on Nvidia A100-80G GPUs, with 512 samples per GPU. The learning rate 2e 4 with 1000 steps warm-up. For large language model, due privacy reasons, we are unable to provide detailed description of the training data here. Our training data contains large amount data from Fineweb-edu (Lozhkov al., 2024), well some other filtered data. We have listed the parameters of our models of various sizes below, some of which have also been mentioned in the main text. We used larger RoPE's base here because previous study (Xu et al., 2024) has shown that the longer the context length, the larger the base required, while the default basc=10,000 is relatively small, even for 2048 windows. Therefore, a relatively large value has been uniformly set here.
Table 16: Configuration.
PARAMETERS
1.5B
2.9B
6.7B
13B
19B
HIDDEN
2,048
2,560
4,096
5,120
6,144
LAYERS
28
32
32
48
HEAD NUMBER
16
20
32
40
48
KV NUMBER
16
4
32
40
4
FFN SIZE
5,504
8,704
11,008
13,824
16,384
MAX LENGTH
2,048
4,096
2048
2,048
12,288
TOTAL TOKENS
10B
500B
10B
200B
VOCAB SIZE
36,000
48,000
36,000
36,000
48,000
GPU
A A100-80G
H800-80G
A100-80G
A A100-80G
A800-80G
GPU NUMBERS
64
512
64
128
240
CONTEXT
2,048
4,096
2,048
2,048
12,288
BATCH SIZE PER GPU
8
LEARNING RATE
2E-4
8E-4
2E-4
2E-4
2E-4
LEARNING RATE SHEDULE
CONSTANT
WARM-UP STEPS
1,000
600
1,000
1,000
3,000
OPTIMIZER
ADAM WITH 0n
0.9.0z = 0.95, WEIGHT DECAY=0.1 100,000
ROPE'S BASE
24


Reviews:
[{"review_id": "q7k4zVEwar", "reviewer": "Reviewer_fHVj", "summary": "Based on the analysis of 2-layer attention, the authors found that 2-layer attention cannot effectively represent information flow from i+1 -> k -> j (for j>=i, k<i). Therefore, the authors proposed KV shifting attention to address this issue. The authors also conducted large-scale pre-training experiments to validate its effect.\n\n## update after rebuttal\nRegarding my main question, \"whether the limitation proposed in this paper still exists when attention layers >=3\", the authors mainly use Figure 8 to respond. Figure 8 does provide some explanation, but I believe it's insufficient and indirect. Therefore, I maintain my score.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your valuable review comments. I hope the following response can answer your concerns.\n\n**Regarding expressive power** \n\nFirstly, if I understand correctly, the \"three\" in the three-layer MLP you mentioned refers to the layer meaning of neuron nodes, which actually corresponds to MLP with two layers of parameter. I think using more layers can also achieve the function of induction heads, as shown in Figure 1 (a), four layers of attention can also achieve the function of induction heads. But for the simple task of induction heads, 4 layers cannot improve compared to 2 layers, so more difficult tasks can be considered. For example, you can refer to the hop k task in Appendix E1, where KV shifting attention can achieve better performance in hop k tasks. Here, we compared L=2,3,4,5,6 layers.  In terms of motivation, the two layers situation has inspired the work, and rigorous theoretical analysis of multiple layers is difficult, which we leave for the future.\n\n**The impact of model depth**  \n\nYou can refer to the hop k task in Appendix E1, where KV shifting attention can achieve better performance in hop k tasks. Here, we compared L=2,3,4,5,6 layers. From the experimental results, we speculate that multi-layer did not alleviate this situation. Because if there is relief, multi-layer attention can better achieve multi hop through some shortcut, thereby filling the gap between Figure 8 (b) and Figure 8 (c). In fact, this gap did not decrease with the increase of layers. I think the causal mask has caused such restrictions on information flow.\n\n**Loss and Benchmark**\n\n(a) Firstly, the learning rate of the model trained on the 19B is 2e-4, which is relatively small. As shown in our Figure 4 at a scale of 1.5B, the gap between the two will narrow in the case of primary school learning rate. If a larger learning rate is adopted, it may be speculated that the gap in loss between the two will increase. (b) Secondly, I would like to point out that loss can be confusing during the pre training process of large language models, because there are a large number of tokens that can be easily predicted using n-gram information, without truly utilizing contextual information[1]. Therefore, we compared some benchmarks, among which benchmarks like MMLU are relatively reliable[2], as shown in Figure 11, where KV shifting attention achieved good results on MMLU. (c) The results of some toy tasks in Appendix make us believe that the model can also achieve better in context learning ability at larger scales. \n\n\n[1]This is particularly evident in the training of long texts, where models with similar losses have significantly different abilities for long-distance information retrieval. \"Can Perplexity Reflect Large Language Model's Ability in Long Text Understanding?\"\n\n[2] For example, we can observe that the performance on MMLU can clearly separate transformer  models from mamba model in Table 4 of \"An Empirical Study of Mamba-based Language Models\".", "claims_and_evidence": "1. I'm not entirely sure if my understanding is correct. According to Property 1, for j to indirectly attend to i+1 through token k, k needs to be k>= i+1. This is indeed a limitation of attention. However, this property only considers the 2-layer case. If LLMs have more than 2 layers, would this problem be solved? I think this is similar to MLP's expressiveness: although 2-layer MLPs cannot represent non-linearly separable functions, 3-layer MLPs can represent all functions. Therefore, if the limitations proposed in Property 1 can be solved by 3-layer attentions, is this problem still serious considering that current LLMs generally have dozens of attention layers?\n2. Following point 1, the impact of model depth on Property 1 and experimental results has not been verified.\n3. Following point 1, according to figure 3, the impact of KV shifting on loss is minor, especially in the 19B model. Does this confirm the limitation of this paper's motivation: for larger and deeper LLMs, the limitations of 2-layer attention revealed by property 1 disappear or are no longer a key issue?", "methods_and_evaluation_criteria": "Yes. However, according to Figure 3, its effect is not significant.", "theoretical_claims": "I tried to understand property 1. Please see Claims And Evidence regarding my concerns.", "experimental_designs_or_analyses": "Please see the third point in Claims And Evidence.", "relation_to_broader_scientific_literature": "I think this is similar to MLP's expressiveness: although 2-layer MLPs cannot represent non-linearly separable functions, 3-layer MLPs can represent all functions.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "No", "other_comments_or_suggestions": "Typos: line 84, Property 3.2 -> Property 1\nCitation: Press,O., lack of year information.", "questions_for_authors": "Please refer to Claims And Evidence.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hBvFC59P2i", "reviewer": "Reviewer_fNW3", "summary": "This paper introduces KV shifting attention, a modification to the standard transformer attention mechanism that changes how keys and values are processed. By decoupling the temporal relationship between keys and values, the model can more efficiently learn induction heads - a critical mechanism for in-context learning in language models. Hence, this could suggest a new approach to improve language modeling.\n\nThe authors provide theoretical analysis showing that KV shifting attention can reduce the width and depth requirements for forming induction heads compared to vanilla transformers. While traditional transformers need at least two layers to implement induction heads effectively, the authors prove that KV shifting attention can implement this mechanism with a single layer.\n\nThe authors evaluate their approach on both small-scale experiments designed to test induction capabilities and large-scale language model pre-training (models ranging from 1.5B to 19B parameters). Their experiments demonstrate that models with KV shifting attention consistently outperform vanilla attention baselines across model scales, with improved performance on standard benchmarks and faster convergence during training.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your meticulous review. We hope our response can be helpful to you.\n\n**More analysis** \n\n(a) **What types of tasks or text patterns benefit most from the improved induction capabilities?** To be honest, it is difficult to exhaust what kind of pattern is more suitable, but we can provide some examples about what kind of pattern is suitable and what kind of pattern cannot be gained. As we have stated in our paper, hop k tasks can achieve better results, and math tasks can also achieve better results. However, n-gram tasks do not have very good results. We later conducted experiments on context free gramma(CFG) [1]. We found that KV shifting attention **does not** improve the learning of CFG, the convergence speed and final performace are almost the same as the vanilla model. Therefore, we speculate that the enhancement of language modeling by kv shifting comes from semantics rather than context free syntax. \n\n(b) **A more detailed analysis of reasoning capabilities.** \n\nThe analysis in this section might also serve as part of the answer to the previous question. Based on our test of mathematical ability, we find that KV shifting attention can effectively package the information of adjacent tokens and store them separately in key and value for ease of subsequent operations. For example, \"3+5=\", the information of \"3\" is used as the value, shifted to \"+\", and the value of \"5\" is shifted to \"=\". Then, \"=\" uses two attention heads, the first one obtains the information of \"3+\" by following \"+\", and the other head obtains the information of \"5\" by following \"=\", and then uses MLP to obtain the correct answer \"8\". \n\nOn the other hand, induction heads are also more conducive to helping establish relationships between things. For example, if a sentence contains \"Beijing roast duck\" or \"Beijing, China\", then kv shifting can easily create some key value pairs in one layer of attention, such as<key: Beijing, value: Roast Duck><key: Beijing, value： In the following text, it is easy to associate Beijing with China or roast duck, which is beneficial for the model to answer questions such as which country Beijing is in or what its cuisine is.\n\n(c) **In terms of potential limitations.** We speculate that KV shifting may enhance the pattern of repetition in the model due to the enhancement of induction heads. Although there is no conclusive evidence, I would like to share an interesting discovery. We later trained a model with 14B parameters with 15T tokens. On the  benchmark math500, we find that the optimal generation configuration is to set the repetition penalty to 1.1, but for qwen2.5-14b, only the repetition penalty needs to be set to 1.05.\n\n**More extensive applications**\n\n(a) **For encoder-decoder structures.**  Due to the absence of a causal mask, the direction of shift in the encoder can be either forward or backward. I believe that KV shift can also enhance the encoder decoder model, as it can benefit from easier to implement induction mechanisms.\n\n(b) **Multimodal model.**  In visual models, there is a similar approach to using token shift [2], which they believe is beneficial for \"back and forth for motion captioning\". Therefore, I believe that similar operations can contribute to multimodal models. If it can really work, then the real reason for its effectiveness will be even more interesting. Does it enhance local information exchange, or did it enhance something similar to the induction heads mechanism?\n\n[1] Physics of language models: Part 1, learning hierarchical language structures. \n\n[2] Token Shift\u001cTransformer for Video Classification.", "claims_and_evidence": "There are four main claims in the paper, each of which are well-supported.\n\n- Improved theoretical foundation for induction heads: The authors provide formal proofs (Theorems 3.2 and 3.3) that KV shifting attention requires less model complexity to form induction heads.\n- Better induction capabilities: Experimental results demonstrate that KV shifting attention models can learn induction patterns more efficiently than vanilla attention models of equivalent or even larger size.\n- Enhanced language modeling performance: The comprehensive evaluations on different models sizes show consistent improvements in both convergence speed and final performance across datasets and benchmarks.\n- Robustness to hyperparameters: The experiments with different random seeds, learning rates, and model sizes provide evidence that the improvements are consistent.", "methods_and_evaluation_criteria": "- The authors evaluate on a suite of standard language modeling benchmarks.\n- The scaling experiments compare performance across model sizes and training compute.\n- Experiments with different hyperparameters and random seeds verify the robustness of the improvements.\n- The toy task suggested measure the specific capability the method aims to improve.", "theoretical_claims": "The theoretical claims provided in Section 3 are supported by formal proofs.\n\n- Theorem 3.2 states that a two-layer vanilla transformer with specific parameter settings can approximate the induction heads mechanism with a certain error bound.\n- Theorem 3.3 states that a one-layer KV shifting attention model can exactly implement the induction heads mechanism.\n- Theorem 3.4 provides an analysis of the learning dynamics for KV shifting attention.\n\nThe proofs appear sound, with detailed derivations provided in the appendices.", "experimental_designs_or_analyses": "- The toy model experiments isolate the induction capabilities being studied, allowing clear observation of the learning dynamics.\n- The scaling experiments with models from 1.5B to 19B parameters follow best practices in the field and allow for meaningful comparisons.\n- The ablation studies (Section 4.5) help identify which components of the proposed method contribute to the performance improvements.\n- The robustness experiments with different random seeds and learning rates strengthen confidence in the results.", "relation_to_broader_scientific_literature": "- It builds on the mechanistic interpretability work by Elhage et al. (2021) and Olsson et al. (2022) who identified and characterized induction heads.\n- The theoretical analysis extends recent work on the theoretical capabilities of transformers, such as Wang et al. (2024) and Sanford et al. (2024a, 2024b).\n- The approach relates to other architectural modifications in language modeling, such as RWKV (Peng et al., 2023), Mamba (Gu & Dao, 2023), and RetNet (Sun et al., 2023), but focuses specifically on enhancing the standard transformer architecture.\n- The paper connects to broader work on in-context learning mechanisms in large language models, an active area of research.", "essential_references_not_discussed": "The following paper might be highly relevant: PermuteFormer: Efficient Relative Position Encoding for Long Sequences.", "other_strengths_and_weaknesses": "Strengths:\n- The paper presents a simple yet effective modification to the attention mechanism with minimal computational overhead.\n- The theoretical analysis provides clear insights into why the method works.\n- The comprehensive evaluation across model scales demonstrates practical value for real-world language modeling.\n- The method is compatible with existing transformer optimization techniques (as shown in Appendix J).\n\nWeaknesses:\n- The paper could provide more analysis of what types of tasks or text patterns benefit most from the improved induction capabilities.\n- While the authors show that KV shifting attention helps with math reasoning (Appendix E.2), a more detailed analysis of reasoning capabilities would strengthen the paper.\n- The paper mentions but does not extensively explore potential limitations of the approach, such as whether there are specific scenarios where vanilla attention might be preferable.", "other_comments_or_suggestions": "It would be very interesting to see how this method would interact with more recently suggested transformer architectures, or with other techniques such as prompting.", "questions_for_authors": "- Would this method also apply for encoder-decoder structures? Would there need to be some modifications in the mechanism?\n- Would this method also apply when training with multimodal models, or have benefit when adapting the two modalities?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pQHJl5lm9I", "reviewer": "Reviewer_SEX2", "summary": "- This paper proposes a new attention mechanism named KV shifting attention to enhance the in context learning ability of transformer.\n- The paper aims at enhance inductive head bias of transformer by shifting the key&value vectors.\n- This paper provides extensive analysis to show the effectiveness of the KV shifting attention in enhancing learning ability and reduce model width requirements,.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your review. I hope the following response can address your concerns.\n\n**Causal Conv** \n\nThank you very much for pointing this out. As demonstrated in our discussion, the shifting operation can be seen as a short convolution with kernel size 2. \n\n(a) From an operational perspective, this is not a very new operation. But where to perform this operation is crucial. As shown in Figure 11, using short convolutions on key and value to decouple them in time is beneficial for the emergence of benchmarks such as MMLU. However, if a similar operation is also used on query, interestingly, the effect will actually deteriorate.\n\n(b) On the other hand, I think we have clarified to some extent the true role of short convolutions in attention. Decoupling the key and value sequences to achieve better context learning. Some previous works believed that this convolution was intended to enhance local information exchange.\n\n\n**Training and inference speed**\n\nWe can note that the computational cost of the kv shift operation is $O(nd)$, Compared to $O (nd^2)$ for matrix projection and $O(n^2d)$ for attention calculation, this is much small.  At the same time, under group query attention, the cost of this part is even smaller. Therefore, there was not much additional computational cost added during the training and inference phases. When considering inference, the additional storage space required is $O (d)$ to store the previous time's kv, which has a relatively small storage overhead compared to the original $O(nd)$. Therefore, kv shift does not have significant memory access overhead during the inference phase. The actual end-to-end training and inference speed is related to the framework used. In our own framework, the additional time cost is less than 2% for a 3B model with kv shift.", "claims_and_evidence": "- The paper clearly written and well motivated.\n- The proposed method is simple and elegant, by only introducing few parameters and computation.\n- Yes, extensive experiments are provided to support the claims.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes.\nThe authors aim at formulating the behavior of induction heads in one or two transformer layers,", "experimental_designs_or_analyses": "Yes, Extensive experiments have been conducted to validate the effectiveness of the method.", "relation_to_broader_scientific_literature": "The proposed method is simple and seems effective, it could be included in future work and more scaled-up models to enhance the language modeling ability.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "weakness:\n- As claimed in line 398, the key to the effectiveness is to achieve better induction heads by decoupling key-value pairs in time sequence. I wonder how 1d causal conv performs compared with kv shifting. In addition, the novelty of the proposed might be weak as it is a simplified version of short conv.\n- Training and inference speed comparison can be included.", "other_comments_or_suggestions": "None", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: I recommend accepting this paper. The KV shifting attention mechanism provides a theoretically inspired approach to improve transformer induction capabilities. Two reviewers gave accept ratings, noting the method's simplicity, theoretical foundation, and evaluation across model scales. One reviewer gave a weak reject, the question is around its relevance in deeper models, but the authors' rebuttal offers more evidence showing benefits across different layer depths. The improvements in convergence speed and performance at various scales, combined with minimal computational overhead, make this a solid contribution to the community.
Decision: Accept (poster)