KV Shifting Attention Enhances Language Modeling

Thank you very much for your review. I hope the following response can address your concerns.

Causal Conv

Thank 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.

(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.

(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.

Training and inference speed

We 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.

Thank you very much for your meticulous review. We hope our response can be helpful to you.

More analysis

(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.

(b) A more detailed analysis of reasoning capabilities.

The 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".

On 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.

(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.

More extensive applications

(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.

(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?

[1] Physics of language models: Part 1, learning hierarchical language structures.

[2] Token ShiftTransformer for Video Classification.

Thank you very much for your valuable review comments. I hope the following response can answer your concerns.

Regarding expressive power

Firstly, 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.

The impact of model depth

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. 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.

Loss and Benchmark

(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.

[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?"

[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".