Generalized Probabilistic Attention Mechanism in Transformers
Attention Mechanism;Transformer;Rank-Collapse;Gradient Vanishing
Yes
No ethics review needed.
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.
The method and presentation are good. They give intuitive, theoretical and experiential justification of the ideas.

The experiment results improve over the baseline.
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.
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.
Yes
No ethics review needed.
- 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.
- 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.
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.
- 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.
Yes
No ethics review needed.
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?
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.
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.
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.
Yes
No ethics review needed.
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^+].
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.
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.
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.
We introduce a novel attention mechanism to mitigate rank-collapse and gradient vanishing in Transformer.
@inproceedings{
anonymous2024generalized,
title={Generalized Probabilistic Attention Mechanism in Transformers},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=dIoLjHet58},
note={under review}
}
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.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Attention Mechanism Transformer Rank-Collapse Gradient Vanishing
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/7c916b51c7578291499a8b111758f735f981fd0c.pdf
foundation or frontier models, including LLMs
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
/attachment/d3921ed8254b5418ca3766ca316d37e4f71c56af.zip
Generalized Probabilistic Attention Mechanism in Transformers
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission