Transformer Training Instability of Softmax and Lipschitz-Kernel Attentions

Authors: JongHyun Hong, Sungyoon Lee
Primary area: optimization
Keywords: Transformer, linear attention, self-attention, optimization, training stability, entropy collapse

Abstract:
Transformers have been making significant progress across various domains, and recently, with scaling up of models like LLMs, they have achieved even greater success. Recent findings have shown that the softmax function in the self-attention used to re-weight the attention logits into probability vectors causes \emph{attention entropy collapse}, where the attention is concentrated on a single token, and it leads to unstable training. In this work, we first demonstrate that the (non-Lipschitz) softmax-based attention leads to the attention entropy collapse but the \emph{Lipschitz-kernel}-based attention does not. We show that the Lipschitzness of the attention plays an important role in keeping the attention entropy stable regardless of the variance of the attention logits. Moreover, we argue that the underlying reason why the attention entropy collapse leads to the training instability is that as the attention probabilities become more concentrated, it causes the attention matrix to gradually increase, leading to gradient exploding.

Decision: reject
Reviewer scores: 3, 1, 3, 3

Reviews and discussion:
<thread>
creation date:
1729947679491

thread_type:
official_review

<thread_object>
creator: reviewer_mzod

<content>
summary:
The authors study the issue of attention entropy collapse in self-attention mechanisms that use the softmax activation function. The authors suggest that this issue is due to the fact that softmax attention is not Lipshitz, meaning a small perturbation of the input can lead to a large perturbation of the output in a non-controlled way. They then suggest that kernel based attention, which is Lipshitz, does not suffer from this problem. Thus they argue that Lipshitzness of the attention is the key role in keeping entropy from collapsing. Furthermore, they discuss how entropy collapse can lead to training instability and suggest that this is primarily due to attention probabilities concentrating into one hot vectors.

strengths:
**Originality:** The authors take up the study of an important problem namely the entropy collapse of attention with a softmax activation. Their approach is original rooted in the idea that entropy collapse is arising due to the lack of Lipshitz property of attention with a softmax activation.

**Quality:** The authors perform a variety of controlled experiments showing that kernel based attention does not suffer from entropy collapse.

**Clarity:** The authors take the time to explain the problem of attention entropy collapse and clarify why this is an important problem.

 **Significance:** The authors provide arguments suggesting that entropy collapse is completely due to the lack of Lipshitzness of softmax attention which is significant as it suggest researchers should be designing transformers in a different manner.

weaknesses:
**Novelty:** Unfortunately I feel the paper does not have any real novelty. The authors show that softmax attention suffers from entropy collapse and suggest that this is due to such attention mechanisms failing to be Lipshitz. Yet they only really provide empirical evidence of this and do not perform any sort of theoretical analysis which motivates their empirical work. They then suggest that kernel based attention is better because it is Lipshitz yet again do not provide any theoretical understanding that motivated their experiments. Furthermore, the authors don't seem to provide any analysis on current transformers used within the literature making it very difficult to see any real novelty in their work.

**Lack of theoretical motivation:** The paper has no real theoretical motivation. The authors carry out small scale experiments with ablations that suggest kernel based attention is better than the usual softmax based attention yet they do not motivate their arguments with theory in any way.

**Lack of experiments:** The authors seem to carry out experiments on a small scale transformer yet which transformer and what exact architecture seem to be never been discussed. They give some details in section 5.1 but these are not enough to give the readers a good understanding of what architecture they are using. I also looked at appendix A where the authors show their implementation details but I cannot see any details on what transformer they are working on. They say they are using a simply transformer in appendix A but no details are given. Furthermore, the authors  suggest that softmax attention suffers from entropy collapse and this leads to gradient norm exploding preventing model convergence (p. 2 second paragraph). However, many transformer models with softmax attention, such as Vision transformers [1], are known to converge so this seems to contradict what the authors are saying. They do not back this up with experiments that clearly show modern day transformers used in the community suffer from convergence issues.

[1] Steiner et al: How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers, TMLR 2022.

questions:
1). Could you please explain what type of transformer your experiments are carried out on? It is possible that I have missed the details but I couldn't find any details about this even in the appendix. You give some details in section 3.1 but these are not enough to give the reader insight into the exact transformer architectures you are using within your experiments.

2). Your main argument is based on the Lipshitzness of attention which is a quantitative measure of how the output of a function changes when its inputs are perturbed. Yet you don't give any theoretical argument as to how the lack of Lipshitzness leads to entropy collapse. Would you be able to give some theoretical insight as to how this connection comes about?

3). You mention on p. 2 paragraph 2 (line 062 to 063) that entropy collapse leads to gradient norm exploding which in turn prevents a transformer model using softmax attention to converge. Yet most transformers in the literature use softmax attention and are able to converge. If what you are saying is in fact true we should see that many transformers do not converge, however this is not what we see. Could you please explain this apparent contradiction with what you mention on p. 2 paragraph 2 (line 062 to 063) and the fact that in the literature many transformers use softmax activation? In fact, softmax is still one of the most used activations for attention.

4). On p.2, paragraph 2 you mention that because the softmax function is not Lipshitz continuous it leads to entropy collapse of the softmax attention. Your paper then suggests that kernel based attention is better as it is Lipshitz. Would I be right in then suggesting that as an alternative to softmax one should take an activation that is Lipshitz on the whole real line? Could softmax be replaced with any Lipshitz activation that is non-constant?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730014710204

thread_type:
official_review

<thread_object>
creator: reviewer_1wxl

<content>
summary:
This paper argues that the underlying reason why the attention entropy collapse leads to the training instability and then proposes Lipschitz-kernel-based attention which is able to prevent attention entropy collapse.

strengths:
**Originality:** Introducing Lipschitz Kernel Attention is interesting.

**Quality:** This paper provides mathematical explanations and experiments showing the proposed method having more stable attention entropy.

**Clarity:** Section 3 is well organized and it is clear to follow the logic in general.

**Significance:** Reducing attention complexity is promising for large scale models. Attention mechanism is not understood well and understanding attention through entropy is appreciated and interesting.

weaknesses:
**Novelty:** The method authors proposed in Eq. 5 is not new [1] [2]. The mathematical part in Section 3.5, which is used to support the argument, has also been introduced in other papers(Line 417).

**Experiment Setting:** I am confused with the settings authors choose. The experiment authors conduct is on a Simple Transformer with SGD, very high learning rate(5e-1), low hidden dimension(3), low attention heads(1), and many others. It would be good to have experiments on standard or larger transformer architecture, such as [3].

[1] Efficientvit: Multi-scale linear attention for high-resolution dense prediction. Han Cai, Junyan Li, Muyan Hu, Chuang Gan, and Song Han.  ICCV 2023

[2] cosformer: Rethinking softmax in attention. Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. ICLR 2022

[3] Attention is all you need. A Vaswani, N Shazeer, N Parmar, J Uszkoreit, L Jones, AN Gomez, Ł Kaiser, I Polosukhin. NIPS 2017

questions:
1. Line 50 "Previous approaches primarily focused on interpreting the results" is not clear and it would be good to clarify the reference nad which approaches you are referring to

2. Line 148 why is that "we can reduce quadratic time complexity into $O(ND^2) = O(N)$ with respect to the input sequence length $N$ as the hidden dimension is smaller than sequence length
</content>
</thread_object>

</thread>
<thread>
creation date:
1730303850758

thread_type:
official_review

<thread_object>
creator: reviewer_amhu

<content>
summary:
This paper describes and reaffirms the connection between entropy collapse, a phenomenon recognized in the literature see (Zhai et al., 2023)—and the activation function used in attention layers, comparing row-wise softmax with Lipschitz functions like ReLU or ELU. 
It essentially rephrases already known results in the literature, providing a light high level understanding of the problem. The provided experiments are limited in scale and validate existing results, without presenting any new theoretical contributions.

strengths:
The research question the authors tried to address is of high importance and relevance to the Machine Learning community, with a lot of significant direct implications at stake.

weaknesses:
1. This submission overall reads more like a report than a research paper due to a significant lack of novelty. Previous works have extensively explored the replacement of softmax with the ReLU activation in both practical (Hron et al., 2020) and theoretical contexts (Bai et al., 2023; Fu et al., 2023). Wortsman et al. (2023a) replace softmax with many normalized different variants. Ramapuram et al. 2024 investigate sigmoid activation function. Additionally, Zhai et al., 2023 directly connect softmax layers to entropy collapse, leaving this submission with little to no new insight to contribute.

2. The manuscript appears unpolished and difficult to read. Below are a few specific issues among the many present: 

- line 40: a verb is missing
- line 51: word "bridge" shouldn’t be there
- lines 128 and 133: the quadratic bottleneck issue is repeated twice
- lines 55, 138, 142, etc.: articles ("the" or "a") are often missing
- line 112 is not a complete sentence and should be separated from the rest by a comma
- lines 114 to 117, lines 151-154 are unreadable
- over the whole manuscript: past and present tenses are inconsistently used to describe the literature
- line 205: "that it can learn" -> "it can learn"

3. Figure 1.a should be presented on another scale for better visibility.

4. Definition 3.1 is redundant given the formal definition provided right before in the paragraph.

5. Definition 3.2 does not make mathematical sense as the left hand side, which is not a function of j, is indexed by j.

6. Plots would be ideally averaged over multiple runs.

7. The primary focus of the paper which is to explain why Lipschitz alternatives are to be preferred to softmax remains speculative: line 227: "To compare with the softmax-based attention, we experiment with a kernelized self-attention with a lipschitz kernel function […], which we expect to mitigate entropy collase".

8. The only mathematical assertion in this paper is questionable and not convincing. Line 419 being an inequality, an increase of the upper bound does not necessarily indicate an increase in the original quantity. The bound can simply become vacuous and uninformative as opposed to the authors’claim (line 430): "the norm of the attention probabilities increases", which is therefore unsupported.

9. Minor comment: Introducing the variance as in definition 3.2 is a good idea but it essentially reduces to considering a temperature within softmax, the latter being analysed in Zhai et al., 2023.

questions:
In its current form, this submission is not suitable as a research paper. However, I encourage the authors to consult the following papers for insights into how the Lipschitz property of the attention layer influences training dynamics (Ramapuram et al. 2024, Jha et al. 2024). This may offer useful directions if they wish to further pursue this topic.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730659936837

thread_type:
official_review

<thread_object>
creator: reviewer_juem

<content>
summary:
This paper addresses the problem of entropy collapse in self-attention-based transformer models, which leads to training instability in such networks. The authors show that softmax-based self-attention suffers from attention entropy collapse due to non-Lipschitz property and can thus be mitigated using Lipschitz kernel-based attention.

strengths:
1. The paper presentation is clear and easy to follow.

2. The issue of training instability in self-attention-based transformers is an interesting problem for understanding the inner workings of transformer models.

weaknesses:
1. The paper does not mention the data/model setting under which experiments have been performed. Details given in section 3.1 are not enough to reproduce the results, the setup section needs more details to better understand the observations in the paper.

2. It is not clear why the concentration of attention on a few tokens should lead to instability, as the premise of attention is to give higher attention to the relevant tokens and ignore irrelevant tokens for specific downstream tasks.

3. Under what setting have the experiments shown in Figure 2 been performed? e.g. share the details regarding the training data and how it is sampled, these details are missing from the paper

4. Apart from stable entropy, is there a specific reason that leads to stable training for Lipschitz-Kernel attention as compared to softmax-based attention?

questions:
Questions to authors: 

1. The paper does not mention the data/model setting under which experiments have been performed.

2. It is not clear why the concentration of attention on a few tokens should lead to instability, as the premise of attention is to give higher attention to the relevant tokens and ignore irrelevant tokens for specific downstream tasks. Authors should include a detailed analysis of  how their observation relates to the intended function of attention in different types of tasks.

3. Under what setting have the experiments shown in Figure 2 been performed? Please clarify. e.g. share the details regarding the training data and how it is sampled, these details are missing from the paper

4. Apart from stable entropy, is there a specific reason that leads to stable training for Lipschitz-Kernel attention as compared to softmax-based attention?
</content>
</thread_object>

</thread>


Metareview: