Value Residual Learning For Alleviating Attention Concentration In Transformers

Q4: Evaluation of Attention Mechanisms

We add several experiments to visualize the attention patterns in different models. We visualize attention weights, value-state norms, and hidden-state norms of tokens at different layers and positions, with detailed results in Appendix A.2 Figure 19. Given that attention clustering often occurs on the first token, we separately show its results in Figure 2 and Figure 20. Besides, we illustrate the distribution of token importance for different models at different layers in Appendix A.2 Figure 21. The results indicate that Resformer and SVformer largely alleviates the problem of attention sinks, value-state drains and hidden-state peaks. And they also exhibit a more uniform distribution of token importance wchich represents the distribution of attention weights.

Q2: Novelty and Theoretical Basis

The effectiveness of Resformer and SVformer can be attributed to three key factors:

1. The application of shortcuts to the value component
2. The utilization of information from the first layer
3. The implementation of shortcuts before each layer's attention operation

Our findings support these claims:

• Figures 9 and 16 demonstrate the necessity of applying shortcuts to the value-states.
• Figure 11 illustrates the importance of using the value from the first layer.
• Figure 10 shows that when the identity map (equivalent to applying the shortcut after the attention operation) is used for residual value, the performance decreases significantly.

Theoretical grounding and analysis:

Attention concentration typically emerges after the second or third network layer [1] [2] [3] and is associated with value-state drains (decreased magnitude of value states) [1], and hidden-state peaks (increased magnitude of hidden states) [2]. [3] shows a mutual reinforcement mechanism exists between value-state drains and attention concentration.

The first layer always shows no attention concentration, and our methods leverage this characteristic. ResFormer indirectly mitigates attention concentration. It leverages the absence of value-state drains in the first layer by introducing a value residual connection. This alleviates value-state drains in deeper layers, thereby disrupting the mutual reinforcement between attention concentration and value-state drains, as shown in Figure 1 (Right) and Figure 2. Specifically, for tokens lacking semantic information, such as the first token, a large value state magnitude can adversely affect the prediction of future tokens if it is overly attended to, resulting in higher loss. Besides, SVformer can also alleviate attention concentration for the same reasons.

Novelty:

Recent studies have linked attention sink to implicit biases during pretraining, with most existing solutions focusing on additional tokens (registers) [4] or additional keys and values (explicit attention bias) [2]. While these methods redirect attention concentration from tokens within the sequence to additionally introduced bias terms, Resformer and SVformer effectively utilize the inherent absence of attention concentration in the network's first layer. This strategy provides a simple yet efficient means of mitigating attention concentration.

In language modeling tasks, Resformer leverages the first-layer values already stored in the KV cache or intermediate activations to achieve significant performance improvements across multiple settings. This enhancement is accomplished almost without additional memory overhead and computational cost (involving only several matrix addition operations).

Building on CLA [5], we found that sharing values across layers is more useful than sharing both keys and values. SVFormer reduces KV cache size by nearly half through value sharing across all layers, with only a small performance penalty. SVFormer's effectiveness improves with increased learning rates and training sequence lengths, outperforming other KV-efficient methods like GQA and CLA in these scenarios.

[1] Zhiyu Guo, et al. "Attention score is not all you need for token importance indicator in kv cache reduction: Value also matters." Arxiv(2024).

[2] Mingjie Sun, et al. "Massive activations in large language models." ArXiv (2024).

[3] Tianyu Guo, et al. "Active-dormant attention heads: Mechanistically demystifying extreme-token phenomena in llms." ArXiv (2024).

[4] Timothee Darcet, et al. "Vision transformers need registers." ICLR (2024).

[5] William Brandon, et al. "Reducing transformer key-value cache size with cross-layer attention." ArXiv (2024).

Thank you for your thoughtful review and suggestions. We have revised the manuscript's presentation based on your comprehensive feedback. We address the concerns below:

Q1: Problem Definition and Motivation

We clarify three important concepts in language models and their relationships:

Key concepts:

1. Over-smoothing: As the model's depth increases, token representations become increasingly similar.
2. Attention sink [1]: Language models allocate substantial attention to low-semantic tokens, particularly the start token.
3. Attention concentration: The model's attention focuses on fewer tokens as network depth increases.

Relationships:

1. Over-smoothing vs. Attention concentration: Over-smoothing is inherent to model architecture, while attention concentration emerges during training. Figure 1 (Left) demonstrates that randomly initialized models exhibit over-smoothing but not attention concentration. These two phenomena interact, as shown in Figure 1 (Middle) and Appendix A.1 Figure 18.
2. Attention sink vs. Attention concentration: Attention sink is a specific manifestation of attention concentration in certain layers, while attention concentration describes the changing attention patterns as information propagates through layers. Deep models exhibit a "concentration - dispersion - concentration" attention pattern, as illustrated in Figure 1 (Middle).

Reasons for alleviating attention concentration:

On the one hand, the concentrated phases may lead to potential information loss. On the other hand, mitigating attention concentration could result in more interpretable attention maps and potentially improve downstream task performance [2][3].

[1] Guangxuan Xiao, et al. " Efficient streaming language models with attention sinks." ICLR (2024).

[2] Timothee Darcet, et al. "Vision transformers need registers." ICLR (2024).

[3] Mingjie Sun, et al. "Massive activations in large language models." arXiv (2024).

Q3: Experimental Setup and Comparisons

Loss-based experiments:

Perplexity (Loss) is a crucial metric for evaluating language models during pretraining [1]. We conducted extensive experiments using Perplexity as the primary indicator to demonstrate ResFormer's superior performance compared to Transformer and its variants across various model sizes, training token counts, and sequence lengths. Through ablation studies, we explored ResFormer's ability to achieve lower Perplexity.

For comparing KV-efficient methods, we adopted the approach from [2], combining different techniques and evaluating Perplexity under various KV-cache reduction levels. Our experiments consistently show that SVFormer achieves 50% KV cache reduction with minimal loss penalty across different learning rate schedules, model sizes, data sizes, training sequence lengths, and architectures.

Added content:

1. Our experimental setup was informed by several previous studies [3] [4]. We have provided detailed hyperparameters for models of various sizes and different training sequence lengths in Appendix A.5, Tables 3 and 4.
2. To further assess the model's robustness, we tested variants of ResFormer defined as $U_n = A_n(V_n + λV_1)$ and NeuTRENO defined as $U_n = A_nV_n + λ(V_1 - V_n)$ with different $λ$ values. The results are shown in Figure 8 and Appendix A.3 Table 22 respectively. Our results indicate that ResFormer can accommodate a wider range of $λ$ values compared to NeuTRENO. Specifically, ResFormer achieves optimal performance at $λ=2$, whereas NeuTRENO's performance peaks at $λ=0.4$.
3. As mentioned in Section 4.1.2, we trained our models for only one epoch on the training dataset. Consequently, the training loss on each new batch can be considered equivalent to validation loss. In Appendix A.6 Table 5, we provide a detailed report of the loss for models of different sizes on various domain-specific validation sets. These validation sets, provided officially, do not overlap with the training data. Additionally, the scaling curves in Figure 6 are also based on validation losses.

Downstream tasks evaluation:

We evaluated the trained ResFormer and Transformer models on several classical downstream tasks following [5] (see Table 1). Our best-performing models required 2-4 days of training time on 8 A100 GPUs. Due to the scale limitations of our models, we did not evaluate them on more challenging benchmarks such as MMLU [6].

Added content: Following your suggestion, we conducted object classification experiments on ImageNet-1K using the DeiT-base model [7]. We adhered to the hyperparameters specified in the original paper and trained for 300 epochs. Results indicate that ResFormer exhibits faster initial training but encounters overfitting in later stages. This overfitting was not observed in language modeling tasks, where models do not overfit most of the time. These findings suggest that for vision recognition tasks, ResFormer may require additional regularization techniques to optimize performance.

[1] Jared Kaplan, et al. "Scaling Laws for Neural Language Models." Arxiv(2020).

[2] William Brandon, et al. "Reducing transformer key-value cache size with cross-layer attention." ArXiv (2024).

[3] Albert Gu, et al. "Mamba: Linear-time sequence modeling with selective state spaces." COLM(2024).

[4] Stella Biderman, et al. "Pythia: A suite for analyzing large language models across training and scaling." ICML(2023).

[5] Haoyi Wu, et al. "Layer-Condensed KV Cache for Efficient Inference of Large Language Models." ACL(2024).

[6] Dan Hendrycks, et al. "Measuring Massive Multitask Language Understanding." ICLR(2021).

[7] Hugo Touvron, et al. "Training data-efficient image transformers & distillation through attention." ICML(2021).

Thank you for your constructive review and appreciation of our work!

Q1: The effect of cross-layer attention on attention concentration.

The visualization results in Appendix Figure 20 can support that cross-layer attention alleviates attention concentration. Attention concentration occurs when the angles between certain tokens' keys and other tokens' queries are small [1], often accompanied by low value-state norms for these tokens [2]. As shown in Figure 1 and Figure 2, the first layer typically does not exhibit attention concentration, and its key-value pairs do not display the aforementioned characteristics. This suggests that applying cross-layer attention using the first layer's key-value pairs could enable the model to maintain appropriate attention distribution across important tokens throughout the entire sequence, even in deeper layers.

In Figure 10, the ``Gold Attention" actually corresponds to the model using cross-layer attention on the first layer. We have revised this potentially misleading representation in the manuscript. Results indicate that cross-layer attention underperforms compared to ResFormer. This is because cross-layer attention does not directly affect the current layer's tokens, which still experience value-state drains, as shown in Figure 20. In contrast, ResFormer's impact on the attention optimization process is more fundamental through mitigating value-state drains.

Q2: Performance gains: faster training or true improvement?

While this is a good question, completely decoupling these two performance-influencing factors is challenging. We attempt to provide some indirect insights into this issue. There are two perspectives: (1) additional visualization experiments demonstrating significant changes in transformer attention patterns (Appendix A.2 Figure 19), and (2) scaling curves for ResFormer and the vanilla transformer regarding model and data size (Figure 6). This common method [3] [4] verifies consistent and stable performance improvements in ResFormer. Notably, SVFormer's final performance is a little lower than the original Transformer due to its reduced parameter count.

Q3: Experiments on image recognition based on ViT.

Please refer to our response to Reviewer dQDC ``Experimental Setup and Comparisons”.

Q4: Reasons for dataset and architecture.

Llama is one of the best open-source models available and is widely used for study [5]. We have extended our SVFormer-Transformer comparison to other architecture like GPT-2 architecture, see Figure 15 (d). Besides, SlimPajama is the open-source corpus that most closely resembles Llama's pretraining data distribution and is really comprehensive. Notably, pretraining corpora from different domains do not significantly affect language model scaling trends [3].

[1] Xiangming Gu, et al. "When Attention Sink Emerges in Language Models: An Empirical View." Arxiv(2024).

[2] Mingjie Sun, et al. "Massive activations in large language models." arXiv (2024).

[3] Jared Kaplan, et al. "Scaling Laws for Neural Language Models." Arxiv(2020).

[4] Albert Gu, et al. "Mamba: Linear-time sequence modeling with selective state spaces." COLM(2024).

[5] Haoyi Wu, et al. "Layer-Condensed KV Cache for Efficient Inference of Large Language Models." ACL(2024).

Thank you for your thoughtful review and suggestions. We address the concerns below:

Q1: Problem definition of attention concentration and relevant work.

Please refer to our response to Reviewer dQDC Q1 Problem Definition and Motivation"** and **Q2 Novelty and Theoretical Basis". Generally speaking, over-smoothing is an inherent model characteristic, while attention concentration emerges during optimization. These two phenomena are interrelated. Previous studies have well studied attention sink, a specific manifestation of attention concentration. Attention sink describes static features at a particular layer, whereas attention concentration represents dynamic changes in attention patterns as the model deepens. In practice, attention concentration and dispersion interweave as model depth increases. Existing approaches mitigate attention concentration by redirecting it from sequence tokens to additionally introduced bias terms, merely shifting its focus. In contrast, ResFormer and SVFormer alleviate attention concentration by preventing value-state drains, based on the principles and characteristics of how attention concentration occurs. NeuTRENO, however, fails to effectively influence attention pattern learning due to its use of first-layer values post-attention. Our new visualizations results support these findings, see Figure 1, Figure 2, Appendix A.1 Figure 18, Appendix A.2 Figure 19, Figure 20 and Figure 21.

Q2: Evaluation criterion is unconvincing.

Please refer to our response to Reviewer dQDC ``Experimental Setup and Comparisons". We evaluate our model using perplexity (loss) and assess its performance on several widely used downstream tasks. This approach aligns with common practices in language model evaluation. In Appendix A.6 Table 5 (Table below), we provide a detailed report of the loss for models of different sizes on various domain-specific validation sets. Additionally, the scaling curves in Figure 6 are also based on validation losses.

Q3: Training time and memory comparison.

ResFormer does not increase memory usage compared to Transformer, either during training or inference. It requires retaining the value states from the first layer throughout forward propagation. During training, Transformer already stores these as intermediate activations, while during inference, these value states can be directly retrieved from the KV-cache. In terms of computational complexity, ResFormer introduces an additional matrix addition operation per layer. However, this extra computation is negligible compared to the matrix multiplications performed in each layer.

We define $n_{layer}$ as the number of layers, $d_{model}$ as the dimension size of hidden state, $d_{ffn}$ as the dimension size of intermediate feed-forward layer and $l$ as the sequence length. Taking the GPT-2 model as an example, where $d_{ffn} = 4d_{model}$, the number of non-embedding parameters is $12 n_{layer}d_{model}^2$ [1]. SVFormer reduces the total parameter count by $(n_{layer}-1)d_{model}^2$, which is $\frac{n_{layer}-1}{12n_{layer}}$ of the original non-embedding parameters.

Q4: Other questions.

We have revised the references as your suggestions. Besides, the term 'gold attention matrix' actually means that we adopt cross layer attention and use $Attn(Q_i, K_1)$ rather than $Attn(Q_i, K_i)$ as mapping matrix for $V_1$. We recognize that this may have been misleading and have revised the relevant text and figure captions.

[1] Jared Kaplan, et al. "Scaling Laws for Neural Language Models." Arxiv(2020).

Thank you for your detailed review! Thank you for your interest in SVFormer. However, I believe you may have overlooked ResFormer. I look forward to hearing your suggestions of it.

Q1: Effect of scale is unclear.

We reorganized our results and presented the scaling trends of valid loss for both ResFormer and Transformer as model size and data size increase in Figure 6. Additionally, we illustrated the scaling behavior of SVFormer across different model sizes in Figure 15(d). The results demonstrate that ResFormer achieves equivalent validation loss with 10.4% fewer model parameters and 13.6% less training data compared to Transformer, while maintaining similar memory usage and computational cost. SVFormer, while reducing the KV-cache by nearly half, requires a 12.2% increase in parameters to achieve the same validation loss as Transformer.

We have provided detailed hyperparameters for models of various sizes and different training sequence lengths in Appendix A.5, Tables 3 and 4. All our trained models have fewer than 1 billion parameters (approximately 700 million non-embedding). Results in Appendix A.4 show that with 20 billion training tokens, none of the models overfit. Note that we trained all models (NeuTRENO, DenseFormer, vanilla Transformer, and Resformer) under the same experimental conditions and then calculated the relative training loss. The parameters of SVFormer are less than Transformer. In Appendix A.3 Figure 22, we present the results of NeuTRENO under various values of λ. The optimal performance was achieved with λ = 0.4, which is just the model we used for comparison.

Q2: Relative training loss is not straightforward.

We utilized relative training loss to better illustrate the consistent benefits of new models throughout the training process. While absolute improvements are influenced by various factors such as learning rate, it is more important to demonstrate that the new models consistently provide stable benefits throughout the training process. We believe that the use of relative training loss more clearly highlights this trend of improvement. Additionally, we wish that the scaling experiments mentioned above can address your concern about the perceived small differences between methods.

Q3: Other questions.

1. The results in Figure 6 and Figure 15(d) are not contradictory; they demonstrate the relative performance of ResFormer and SVFormer, respectively, against the standard Transformer.
2. We updated the Figure 4 to clarify: The lines illustrate the similarity of layer outputs using current attention (Eq. 4) or identity attention (Eq. 6) compared to those using cross-layer attention (Eq. 2). Using a trained 82M Transformer model, we investigate the output space differences when mapping values from the previous layer to the current layer's attention output using different matrices.
3. In Equation 8, I should be used instead of J to represent the identity matrix. Thank you for pointing this out.
4. SVFormer initially trains faster but achieves slightly lower final performance than the standard Transformer due to its reduced parameter count. We apologize for any previous imprecise statements.