Generalized Probabilistic Attention Mechanism in Transformers

Authors: DongNyeong Heo, Heeyoul Choi
Primary area: foundation or frontier models, including LLMs
Keywords: Attention Mechanism, Transformer, Rank-Collapse, Gradient Vanishing

Abstract:
The Transformer architecture has become widely adopted due to its demonstrated success, attributed to the attention mechanism at its core. Despite these successes, the attention mechanism of Transformers is associated with two well-known issues: rank-collapse and gradient vanishing. In this paper, we present a theoretical analysis that it is inherently difficult to address both issues  simultaneously in the conventional attention mechanism. To handle these issues, we introduce a novel class of attention mechanism, referred to as generalized probabilistic attention mechanism (GPAM), and its dual-attention implementation within the Transformer architecture. Unlike conventional attention mechanisms, GPAM allows for negative attention scores while preserving a fixed total sum. We provide theoretical evidence that the proposed dual-attention GPAM (daGPAM) effectively mitigates both the rank-collapse and gradient vanishing issues which are difficult to resolve simultaneously with the conventional attention mechanisms. Furthermore, we empirically validate this theoretical evidence, demonstrating the superiority of daGPAM compared to other alternative attention mechanisms that were proposed to address the same issues. Additionally, we demonstrate the practical benefits of GPAM in natural language processing tasks, such as language modeling and neural machine translation.

Decision: reject
Reviewer scores: 6, 5, 5, 6

Reviews and discussion:
<thread>
creation date:
1730693064995

thread_type:
official_review

<thread_object>
creator: reviewer_jq8m

<content>
summary:
Despite the success of Transformers, the standard attention mechanism is associated with two well-known issues: rank-collapse and gradient vanishing. This paper proposed generalized probabilistic attention mechanism (GPAM) to mitigate both the rank-collapse and gradient vanishing issues simultaneously. Evaluation results on real-world datasets demonstrate the effectiveness of GPAM.

strengths:
This paper studies how to alleviate the rank-collapse and gradient vanishing problems in standard attention mechanism, which is an important problem of profound implication in practice. Authors provided a concrete analysis on the intrinsic trade-offs between avoiding rank-collapse and gradient vanishing in standard attention mechanism and showed that GPAM is able to mitigate both problems simultaneously. The theoretical analysis in this paper seems solid and the presentation is clear.

weaknesses:
Evaluation in this work can be significantly improved. Given the prevalent success of Large Language Models (LLMs) on various advanced tasks (e.g., solving math and coding problems), evaluation on only small language models (20M-200M) and classic NLP tasks (e.g., neural machine translation) seems insufficient. More evaluation results on open-sourced LLMs (e.g., Llama-3 or Phi-3 family) and widely used benchmarks (e.g., MMLU or GSM8K) will add great values to this paper.

questions:
Q1. Empirical improvements seem quite marginal according to Table 1. Specifically, a roughly 1% improvement on PPL is not significant enough given that the proposed approach has approximately 1% more parameters (Line 450-451).

Q2. In Line 250, the range for PG should be [1 - \lambda^-, 1 + \lambda^+].
</content>
</thread_object>

</thread>
<thread>
creation date:
1730700299611

thread_type:
official_review

<thread_object>
creator: reviewer_jf2u

<content>
summary:
The paper introduces GPAM, an attention mechanism that allows negative attention scores while maintaining a fixed sum, addressing both rank-collapse and gradient vanishing problems in Transformers. Their proposed dual-attention GPAM (daGPAM), adds a negative attention matrix computation alongside the original one, requiring small additional parameters. The authors show the theoretical advantage of their model and examine their models on language tasks.

strengths:
1. The writing is easy to follow
2. The paper shows quite interesting results gradient vanishing and rank-collapse is hard to solve simultaneously total norm of gradients is maximized can lead to rank-collapse.
3. The authors theoretically show that their method can mitigate both the rank-collapse issue and reduce the the gradient vanishing problem.

weaknesses:
My main concerns are:
1. The improvement of daGPAM is very marginal, and not significant in both Language modeling (on Wikitext-103 and Enwiki8, the improvement is around 0,5 PPL) and machine translation tasks (on IWSLT14 and WMT14, maximum improvement is around 0.7% BLEU score.) 
2. Computational cost: the dual-attention structure requires computing two attention matrices instead of one, which is very inefficient. The marginal improvement does not justify this performance-efficiency trade-off.

questions:
1. The optimal $\lambda^{+}$ and $\lambda^{-}$ values seem to vary across different tasks. Can the authors provide an ablation study on choosing different settings of the hyper-params. 
2. Can the authors provide the inference time and memory cost of the model?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730866460793

thread_type:
official_review

<thread_object>
creator: reviewer_jbvx

<content>
summary:
The authors propose daGPAM -- an attention module that computes an additional negative coefficient attention matrix aiming to solve rank collapse and vanishing gradients in transformers.

strengths:
- The paper includes theory on how daGPAM can reduce rank collapse and vanishing gradients. It does so by showing the residual norm of daGPAM and also its gradient is greater than the respective values for vanilla attention (the original values for attention having been computed by [[Dong et al](https://proceedings.mlr.press/v139/dong21a.html)]).
- The paper is well written and easy to follow, and contextualizes relevant work well.

weaknesses:
- I think the downstream experiments for this paper could be fleshed out more. I appreciate that some sort of comparison with alternative attention mechanisms like CoDA is present, but I think only having these comparisons on PTB is insufficient. So many attention variants have been proposed over the years, and none have really taken hold, so I think the empirical bar should be high for current and future variants. Even though this paper shows evidence of reducing rank collapse, rank collapse is ultimately something that matters as it reflects in downstream performance, so I think rank collapse on its own is not enough in the absence of strong empirical results. I think this could be addressed by the authors including more attention variants (CoDA etc) in the LM and NMT experiments.

questions:
- Do the authors believe other attention variants solve the rank collapse / vanishing gradients problem, and can this be experimentally validated?
- How does rank collapse (Figure 4) reflect in downstream performance? It would be interesting to see that a higher degree of rank collapse results in lower downstream results, but without this correlation it is hard to reason about the value of reducing rank collapse.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730963328822

thread_type:
official_review

<thread_object>
creator: reviewer_i3m3

<content>
summary:
The paper introduces a new attention mechanism that replaces the attention scores. The attention scores in softmax for each query are a convex sum, so scores must be positive. They argue for using an affine sum which also sums to 1 but allows for negative scores.

They are motivated by rank collapse and gradient vanishing. Rank collapse is mitigated since the negative sums introduce more diversity. They show that gradient vanishing is a related problem.

To validate their method they perform several small scale benchmarks, namely perplexity and translation and find their methods improve against the baselines.

strengths:
The method and presentation are good. They give intuitive, theoretical and experiential justification of the ideas.

The experiment results improve over the baseline.

weaknesses:
The range of evaluations is pretty limited and extra computation is needed.

A benchmark against just increasing the number of attention heads would be useful.

questions:
How does the method work in long context scenarios? Even a simple benchmark like Long Range Arena would be helpful. My intuition feels like there should be a larger improvement.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This work introduces daGPAM, an architecture update addressing two problems in transformers : rank collapse and vanishing gradients. Reviewers enjoyed the theoretical contributions, but the empirical validation of the proposed approach is limited. The improvements observed on standard benchmarks are marginal, and the paper lacks a comprehensive comparison with other alternative architectures. Overall, the paper falls short of the acceptance threshold.