Value Residual Learning For Alleviating  Attention Concentration In Transformers

Authors: Zhanchao Zhou, Tianyi Wu, Zhiyun Jiang, Zhenzhong Lan
Primary area: foundation or frontier models, including LLMs
Keywords: Transformer, Self-Attention, Cross-Layer Attention, Residual Learning, Language model, KV cache

Abstract:
Transformers can capture long-range dependencies using self-attention, allowing tokens to attend to all others directly. However, stacking multiple attention layers leads to attention concentration. One natural way to address this issue is to use cross-layer attention, allowing information from earlier layers to be directly accessible to later layers. However, this approach is computationally expensive. To address this problem, we propose Transformer with residual value (ResFormer) which approximates cross-layer attention through adding a residual connection from the values of the the first layer to all subsequent layers. Based on this method, one variant is the Transformer with single layer value (SVFormer), where all layers share the same value embedding from first layer, reducing the $KV$ cache by nearly 50\%.  Comprehensive empirical evidence demonstrates that ResFormer mitigates attention concentration problem in deeper layers and enhances representation across most layers, outperforming the vanilla Transformer, DenseFormer, and NeuTRENO in training error as well as downstream tasks. SVFormer trains significantly faster than the vanilla Transformer and performs better than other methods like GQA and CLA, with performance influenced by sequence length and cumulative learning rate.

Decision: reject
Reviewer scores: 5, 3, 3, 3

Reviews and discussion:
<thread>
creation date:
1730479517230

thread_type:
official_review

<thread_object>
creator: reviewer_dqdc

<content>
summary:
This paper presents ResFormer and SVFormer, two Transformer model variants that address challenges with deep Transformers, particularly attention concentration where deeper layers focus too much on fewer tokens. ResFormer incorporates a residual connection from the initial layer's values to subsequent layers, thus approximating cross-layer attention without heavy computational costs. SVFormer simplifies further by sharing the value embedding from the first layer across all layers, reducing memory usage by nearly half and accelerating training.

strengths:
1. The paper offers an interesting twist on standard residual connections by applying them specifically to V instead of the usual hidden state H. This approach targets the common issues of over-smoothing and information loss in deep Transformers.

2. SVFormer aims to make Transformer models more efficient by sharing the same value embeddings across layers, reducing memory and computation needs. This design could help make large models faster and more practical for applications with long sequences.

weaknesses:
1. **Problem Definition and Motivation**: The problem of "attention concentration" is not clearly defined or sufficiently justified. It is essential for the authors to establish a precise understanding of this issue and clarify why it is a significant challenge within Transformer architectures. Without a thorough introduction and motivation for addressing "attention concentration," it remains unclear what gap this work aims to fill, and the importance of resolving it is left ambiguous.

2. **Novelty and Theoretical Basis**: The proposed approach largely resembles existing residual connections in Transformers, as seen in architectures like ViT and LLaMA. The primary difference with ResFormer appears to be the application of residuals to the value V alone, rather than to the hidden state H as in traditional models. However, this adjustment lacks theoretical grounding and rigorous analysis, especially with regard to the SVFormer, which further simplifies by removing layer-specific values. This simplification seems ad-hoc and trivial, as no theoretical guarantees or insights are offered to support the effectiveness or necessity of such changes.

3. **Experimental Setup and Comparisons**: The experiments are limited and do not provide a thorough benchmark. Although the models are trained on a LLaMA-like architecture, there is no comparative performance evaluation against other prominent Transformer-based or SSM-based models. Furthermore, there are no tests involving visual downstream tasks, which would have strengthened the claims of improvement in Transformers and provided a more comprehensive evaluation across different modalities, especially for encoder-only tasks.

4. **Evaluation of Attention Mechanisms**: An essential part of evaluating any modification to Transformer architectures is understanding how the attention patterns differ from those in the vanilla Transformer. Although the paper discusses attention concentration, it does not provide visualizations or statistical analysis of the multi-head attention weights to demonstrate the proposed method's effect on attention distribution. Such an investigation is critical for validating the claims and understanding how the modifications impact attention dynamics.

questions:
same as weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730492657091

thread_type:
official_review

<thread_object>
creator: reviewer_ncnk

<content>
summary:
This paper studies the problem of attention concentration in Transformers and proposes solutions that try to approximate cross-layer attention by incorporating the "value" from first layer into subsequent layers. There are two solutions: ResFormer that uses residual mapping and SVFormer that uses the same V across all layers. Experiments show that the proposed solutions perform better than baselines on language modeling tasks.

strengths:
- The paper introduces a relatively new and important problem that affects existing Transformer architecture. This is useful towards understanding the dynamics and behavior of Transformers.
- The proposed solutions only require small changes to existing Transformer architecture. They can be immediately useful for many existing Transformer-based models.
- The paper provides a good analysis and ablation study on ResFormer and SVFormer that demonstrate their benefits over existing Transformer. Particularly, ResFormer is shown to be achieving higher token importance entropy (i.e., less attention concentration) than traditional Transformer.

weaknesses:
- The authors claim that cross-layer attention is useful at reducing the effect of attention concentration but it is unclear why this would be the case. This work is built on the premise that ResFormer approximates cross-layer attention and thus it is effective against attention concentration. But we do not really know that cross-layer attention provides such a benefit. The author should perform some analysis and/or small-scale experiment on a baseline that actually uses cross-layer attention to check its behavior against that of ResFormer.

- It is hard to disentangle the effects from: (1) reducing attention concentration; (2) ease of optimization in the proposed solutions. Using V in the form of residual mapping (ResFormer) or layer sharing (SVFormer) should make it easier to optimize network parameters during training. It may be possible that the accuracy improvements are largely attributed to the ease of optimization rather than attention concentration reduction. The authors should explain this.

- It would also be interesting to see how well the proposed methods work for non-language tasks and architectures like ViT (image recognition).

questions:
- What are the reasons of using LLama-like architecture and SlimPajama dataset?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730517443174

thread_type:
official_review

<thread_object>
creator: reviewer_tuj8

<content>
summary:
This manuscript presents a novel framework for approximating cross-layer attention. Within this framework, the authors introduce ResFormer as a practical implementation, demonstrating its effectiveness in mitigating attention concentration challenges. In addition, they propose SVFormer within the same framework, which further enhances efficiency by reducing the memory requirements for KV caching, thus lowering overall computational costs.

strengths:
The manuscript proposes a framework for reducing the computational cost of cross-layer attention, offering a unified approach that integrates and extends existing methods, including NeuTRENO and DenseFormer.

weaknesses:
1. The paper lacks discussion of prior work on the attention concentration problem and the connection to the over-smoothing issue addressed by NeuTRENO is unclear. A more detailed review of relevant literature would enhance clarity and better contextualize the impact of this work.
2. Using training loss as a criterion for comparing model performance is unconvincing (e.g. in Section 4.2, 4.3, 4.6), as it may not accurately reflect generalization. A more reliable evaluation metric, such as accuracy or perplexity on a separate validation set, would provide a clearer assessment of the model's effectiveness.
3. Minor comments:
- The term “gold attention matrix” in Section 4.3 should be clearly defined for better understanding.
- Right margin violated at line 659.
- Some references list only the first author; please ensure consistency in citation formatting.

questions:
Could you provide a detailed comparison of training time and memory requirements for SVFormer and ResFormer relative to other baseline models? Such a comparison is crucial for understanding the extent to which these models benefit from the proposed framework.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730749261730

thread_type:
official_review

<thread_object>
creator: reviewer_2rae

<content>
summary:
Paper proposes SVFormer, a way to reduce the size of the KV cache in Transformers by almost 50%. The authors propose sharing the values from the first self-attention layer across all layers. They find that this outperforms other approaches that reduce the KV cache size and perform extensive ablations to find when SVFormer works.

strengths:
- Paper is straightforward and easy to read. 
- It's interesting that values from the first layer can be used throughout the network for a small loss penalty. 
- Authors thoroughly discusses prior work and explains the contributions of this work. 
- Lots of ablations and experiments.

weaknesses:
- The paper leaves out many important details. See the "Questions" section for specifics.
- Results are not well organized, and appear to have contradictory findings. Fig. 13 (c) in particular shows that SVFormer only outperforms a vanilla transformer when they have 2M parameters, which is very small.  At 82M parameters, SVFormer already is worse than the baseline. Fig. 13 (d), 14, and 15 also indicate that SVFormer hurts loss. However, Fig. 6 shows that SVFormer does better at larger scales
- I don't like the practice of subtracting the transformer performance and showing the difference. It potentially (a) hides bad baseline performance, and (b) potentially hides the fact that the difference between methods is tiny compared to the overall training loss curve.

questions:
- Fig 4: 
  - What model is this? 
  - It seems very shallow -- only 6 layers?
  - These seem like such shallow models. Is “current mapping” Eq. 4 or Eq. 5?
- Eq. 8: why is the identity matrix  $J$ and not $I$?
- Effect of scale unclear:
  - Which figures correspond to the 700M parameter model described in 4.1.1?
  - How are the hyperparameters tuned for the baselines (especially the vanilla Transformer)?
  - Why is Fig. 6 (right) inconsistent with Fig. 13 (c) on the effect of model size?
  - Authors should show scaling laws to show much better or worse their method is they scale up their model. 
- L460: "SVFormer will always be gradually surpassed by vanilla attention during training while its training speed is much faster than vanilla attention." How much faster can it be during training?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732482344086

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Updated Manuscript and General Response

comment:
We sincerely appreciate all the reviewers for their time and constructive comments. 

In general, we proposed **ResFormer** and **SVFormer**. Both of them mitigate attention concentration through **elegant and concise approaches**, whereas previous methods attempted to divert attention concentration by introducing additional tokens or similar techniques. Specifically, they alleviate value state drains by incorporating the first layer's value into the current layer's value, thereby indirectly suppressing attention concentration. 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**. Besides, SVFormer, while reducing the KV -cache by nearly half, requires a 12.2% increase in parameters to achieve the same validation loss as Transformer. 

In light of the reviewers' feedback, we have made several enhancements to our manuscript. **All changes are highlighted in blue in the updated PDF**. A summary of the **key updates** is provided below:
1. **Figure 1**: Add illustration of the over-smoothing and attention concentration in randomly initialized models; add illustration of the attention concentration issue in Llama and Mistral models; and substitute the entropy analysis of 82M models with that of 468M models. 
2. **Figure2**: Add the analysis of token importance, value-state norms, and hidden-state norms of the first token across layers of 468M models (Transformer, NeuTRENO, ResFormer, SVFormer).
3. **Section 1**: The connections between over-smoothing, attention sink, and attention concentration are elaborated in detail. Recent research progress related to attention concentration is introduced. The performance and novelty of the model are emphasized.
4. **Figure 6 & Figure 15**: Validation loss scaling curves for ResFormer (SVFormer) and Transformer are shown, comparing their performance across model and dataset sizes.
5. **Figure 8 & Figure 22**: Robustness analyses of hyperparameters are conducted for both ResFormer and NeuTRENO.
6. **Section 4**: Add analysis of scaling results and visualization outcomes.
7. **Appendix A.1**: Visualization of the over-smoothing phenomenon is added.
8. **Appendix A.2**: More visualization of the attention concentration phenomenon is added.
9. **Appendix A.5**: Detailed hyperparameters for models of various sizes and different training sequence lengths are added.
10. **Appendix A.6**: Validation loss results for various models is added.
</content>
</thread_object>

</thread>


Metareview:
metareview:
All reviewers converged on rejecting the paper post rebuttal. The AC checks all the materials, and while appreciating the additional efforts including results and analyses and making major modifications to the draft, the AC resonates with the reviewer consensus that the paper currently has issues to address and would benefit from another cycle.