Transformer Training Instability of Softmax and Lipschitz-Kernel Attentions

[W1] Novelty

• Our paper does not explicitly prove that softmax is non-Lipschitz, as this is a well-established fact due to the exponential nature of its formulation. It can be inferred from the kernel in Section 3.1 and Definition 3.1 (to be included in the Appendix). The claim that kernel-based attention outperforms softmax is supported by the fact that Lipschitz kernel self-attention prevents entropy collapse, enabling stable learning. The cause of such collapse is substantiated through controlled experiments examining whether the function is sensitive to variance. Comparisons with existing Transformers were conducted by testing against previous methods that attempted to prevent entropy collapse while still relying on softmax function.

[W2] Lack of theoretical motivation

• While the theoretical basis may be limited, we have provided clear evidence that the attention entropy patterns differ significantly between kernel attention based on Lipschitz functions and attention based on the softmax function. This difference is due to whether the function increases or does not increase the variance, which we have substantiated in our findings.

[W3] Lack of Experiments

• We acknowledge that the description of the architecture may be insufficient, and we plan to add further explanations in the Appendix. The architecture used in our paper follows the structure (Mahankali et al., 2023; Ahn et al., 2023; 2024), utilizing only the self-attention layer. This structure was chosen because it is well-suited for observing how attention entropy collapse in self-attention can lead to unstable training. Additionally, while Transformers or ViTs are known to converge well in some cases, this is only partially true; as model size increases (e.g., LLMs), sensitivity to the learning rate grows, and this phenomenon has been widely reported (Wang et al., 2021; Wortsman et al., 2023; Zhai et al., 2023; Dehghani et al., 2023).

[Q1] The architecture used in the experiments

• The Transformer architecture used in this study consists solely of stacked self-attention layers, with weights applied only to the query, key, and value. This structure was chosen to analyze whether attention entropy collapse occurring exclusively within self-attention impacts training stability.

[Q2] Theoretical insight of how the lack of Lipshitzness leads to entropy collapse

• We are planning to provide theoretical and experimental insights to establish a connection between Lipschitz continuity and entropy.

[Q3] How do you explain the fact that most transformers using softmax attention converge successfully?

• As you mentioned, Transformers, which are based on the self-attention mechanism, have demonstrated strong performance across various domains. However, with current research on Transformers, it has become well-known that as model size increases or with certain hyper-parameters (e.g., learning rate), these models become highly sensitive, making them challenging to optimize (Dehghani et al., 2023; Wortsman et al., 2023). Also, softmax function, which utilizes an exponential function, is also extremely sensitive to input changes, aligning with the fact that it is highly susceptible to variance (Wang et al. 2021; Shen et al., 2023; Hoffmann et al., 2023).

[Q4] 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?

• The main argument presented in the paper is that by re-weighting attention logits with a Lipschitz continuous function, specifically one with a rate of change not exceeding 1, variance can be maintained stably, thus preventing collapse. Therefore, it may be beneficial to use activation functions that are Lipschitz continuous over the entire real range. However, considering that traditional self-attention uses softmax partly due to its non-negative output, it is preferable to use Lipschitz functions like ReLU and ELU+1 as kernels, as they respect this non-negativity characteristic.

Thank you for your response.

I am still not convinced with the authors' work. While interesting I still feel there is still significant areas that need working which I have outlined previously and don't feel the reviewers have adequately addressed. I have gone back and looked at the paper and noticed the authors have not made any real changes that they said they would in response to my initial review. There is still a significant lack of theoretical motivation/contribution as well as a lack of experiments on transformers that are actually used by researchers within the field.

Due to this I will keep my original rating of the paper and thank the authors for the discussion.

[W1] Novelty

• What we are proposing in this paper is not so much a method, but rather an insight into how entropy collapse—a limitation that can arise in attention modules—occurs due to the non-Lipschitz nature of softmax. We provide empirical evidence that replacing softmax with other Lipschitz continuous functions in kernel-based self-attention, aimed at reducing computation, more effectively maintains stable entropy. Furthermore, we theoretically and experimentally investigate how the Lipschitz property impacts entropy and why entropy collapse leads to training instability.

[W2] Experiment Setting

• The purpose of using a high learning rate with a small model scale is to approximate a large language model setting, as done in (Wortsman et al. 2023, Von Oswald et al., 2023; Ahn et al., 2023; 2024). Our experiments demonstrate that even under these conditions, entropy collapse does not occur, and stable training is achieved with Lipschitz continuous self-attention. While previous studies have extensively experimented with standard Transformer scales, our work focuses specifically on attention entropy collapse using only self-attention layers. By employing a small scale model and a high learning rate, we approximate an LLM setting and conduct controlled experiments to analyze the causes of this issue.

[Q1] Line 50 "Previous approaches primarily focused on interpreting the results" is not clear and it would be good to clarify the reference.

• It has been noted that the explanation of related works is insufficient, and we plan to further supplement and expand this section.

[Q2] 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.

• In this paper, the Lipschitz self-attention used approximates the traditional kernel-based self-attention by replacing the softmax function with a kernel function, as shown in eq (4). In conventional self-attention, as seen in eq (1), calculating each row of the $N \times N$ attention probabilities involves first taking the dot product of queries and keys, resulting in a computational complexity of $O(N^2D)$. By restructuring the equation using the associative property, as shown in eq (5), the method computes the product of keys and values first, reducing the complexity to $O(D^2N)$. In our experiments, instead of $\phi$, we employed Lipschitz functions (e.g., ReLU or ELU+1) to observe performance. The assumption that the hidden dimension is smaller than the sequence length $N$ aligns with practices in traditional NLP domains, where the sequence length is generally larger than the hidden dimension, leading to reduced computational complexity (as referenced in Qin et al., 2022).

[W1] The differences between our work and other studies that replace softmax with ReLU or sigmoid functions, as well as those focusing on attention entropy collapse.

• We theoretically and experimentally connect the causes of attention entropy collapse and the resulting training instability by comparing non-Lipschitz and Lipschitz functions of re-weighting attention logits. Therefore, it is not limited to the ReLU function but is used as a representative example of a Lipschitz function that ensures non-negative values. Among the related works, Hron et al. use the ReLU function because it guarantees non-negative values like softmax and behaves similarly to the identity function. Wortsman et al. concentrate on demonstrating that ReLU outperforms other activation functions, while Bai et al., 2023, Fu et al., 2023 use ReLU purely for theoretical convenience without additional reasoning.
• Sigmoid Attention (Ramapuram et al. 2024), similar to our work, leverages the Lipschitz continuity of the sigmoid function for stability and calculates similarity scores independently, resulting in more gradual output changes than softmax. However, these studies lack empirical evidence on how stability differs between the two functions, focusing instead on computational experiments.
• Zhai (et al. 2023) investigate entropy in softmax-based self-attention layers, providing a lower bound for entropy and highlighting the growing weights of query and key as a primary cause. But, our research focuses on how the Lipschitz property of re-weighting functions affects the behavior of attention entropy and systematically conducts experiments to establish its step-by-step connection to training instability.
• In this paper, we show that, unlike softmax, which increases the variance of attention logits, Lipschitz functions maintain stable variance, helping to prevent entropy collapse. We also provide theoretical and empirical evidence linking attention entropy collapse to training instability. This forms the novelty of our research.

[W2-W5] Revised the content of the paper

• That section has been fully revised.

[W6] Plots would be ideally averaged over multiple runs.

• In the experimental results shown in the plot, for non-Lipschitz functions like softmax or QK-LayerNorm, the gradient norm became excessively large, ultimately causing training to halt. Therefore, we believe that averaging over multiple runs is not appropriate in this case.

[W7] Why Lipschitz alternatives are to be preferred to softmax remains speculative.

• The Lipschitz property of a re-weighting function that maps the inner product of queries and keys to values between 0 and 1 can influence attention entropy collapse. Using non-Lipschitz functions with high rates of change significantly increases the variance of attention logits, leading to one-hot-like distributions that focus excessively on a single token. In contrast, when re-weighting with Lipschitz functions, even with large variance, the rate of change remains bounded within a constant, resulting in more balanced distributions. We demonstrate this robustness against entropy collapse both theoretically and experimentally.

[W8] Line 419 being an inequality, an increase of the upper bound does not necessarily indicate an increase in the original quantity.

• The first term in the upper bound, $\||P\||_F$​, represents the norm of the attention probabilities. This norm reaches its maximum value when attention entropy collapses completely (i.e., when all vectors along the sequence axis become one-hot). Therefore, when attention entropy collapses—resulting in a larger $\||P\||_F$—the upper bound increases, leading to training instability. Conversely, as $\||P\||_F$​ decreases, the upper bound lowers, also contributing to stability.

[W9] Considering a temperature within softmax

• Temperature adjustment does not make softmax a Lipschitz function or remove its non-Lipschitz property. It simply adjusts how sensitive softmax is to changes in its logits.

[Q1] 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

• Thank you for encouraging our research. We will review the referenced paper as well as additional related studies.

[W1] I do not acknowledge the theoretical contribution that the authors claim.

• We are supplementing the theoretical contribution.

[W6] True, maybe showing multiple runs in the appendix would do the work?

• We will supplement the result and include it in the appendix.

[W6] The authors claim showing their claim 'theoretically', where can I find the proof? It appears that the high-level arguments presented are mistaken for theoretical proofs.

• It is clear that using the softmax function within self-attention can lead to a reduction in attention entropy, potentially resulting in collapse and contributing to training instability. However, the theoretical foundation for this explanation may be insufficient, and further attempts at a more detailed theoretical proof are underway. Nonetheless, we believe that the experimental results sufficiently support our claims.

[W8] I maintain that having an "upper bound that increases" does not provide any meaningful theoretical insights.

• While the increasing upper bound itself may lack theoretical insights, we have supplemented it with empirical evidence. Specifically, as shown in Figure 5, for softmax-based self-attention, the norm of P increased as the attention entropy decreased. In contrast, for Lipschitz kernel self-attention, stability was observed.

[W9] It is simply not true. The authors might consider looking into the literature, e.g. https://arxiv.org/abs/1704.00805

• The paper you introduced confirms that the softmax function is indeed Lipschitz continuous. However, in self-attention, where $X$ is the input and $P(X)$ represents the attention probabilities obtained after applying the softmax function, the multiplication of $P(X)$ and $X$ in the self-attention mechanism causes the overall mapping to lose Lipschitz continuity. Additionally, since the softmax function is highly sensitive to its input values, when $F(X)=P(X)X$ (self-attention function that dot-product the attention probabilities with the values), the rate of change of this F(X) function is greatly influenced by the saturation of the softmax, which varies depending on the input values (https://arxiv.org/pdf/2006.04710).

Thank you for your response and for sharing the alternative resources.

[W1] Data/Model Setting

• We have specified the experimental settings in Appendix A. Unlike previous studies that analyze attention entropy collapse in larger-scale models with downstream tasks, we focused on this attention entropy collapse with the only attention layer to clearly demonstrate why this issue occurs and how this phenomenon can lead to training instability. Also, we conducted our experiments in a simpler linear regression task, thereby proving that this phenomenon can occur universally.

[W2] Why the concentration of attention on a few tokens should lead to instability?

• Attention entropy collapse occurs when the query-key inner product (also called the attention logits) is converted into an attention probability vector, leading to an excessive focus on a single token (Figure 2). At this point, if all vectors along the sequence dimension take the form of one-hot vectors, the norm of the attention probabilities increases, causing the gradient to explode and resulting in unstable training (line 423-430).

[W3] Under what setting have the experiments shown in Figure 2 been performed?

• Following the similar approach as in (Garg et al., 2022; Zhang et al., 2023; Mahankali et al., 2023; Von Oswald et al., 2023; Ahn et al., 2023; 2024), $x_i$ and $w_i$ were randomly sampled, $x, w \sim N(0, I)$ (It was mentioned on line 209). The $y_i$ was then calculated as the inner product of the two vectors, training the model on the linear regression $y_i=w^{\top}x_i$.

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

• Our paper focuses on empirical evidence showing that training exclusively with a model composed solely of self-attention layers leads to unstable training when using non-Lipschitz functions (such as softmax and QK-LayerNorm). By concentrating on how Lipschitz continuous functions prevent this instability, other issues were unlikely to arise.

[Q1] Data/Model Setting

• In Appendix A, we detail the experimental settings, focusing on attention entropy collapse using only the attention layer to clearly demonstrate its causes and impact on training instability. By conducting experiments on a simpler linear regression task, we show that this phenomenon can occur universally, unlike prior studies that focus on larger-scale models with downstream tasks.

[Q2] Why the concentration of attention on a few tokens should lead to instability and an analysis of how their observations align with attention's purpose in different task.

• When attention is concentrated on a few tokens, meaning that attention entropy decreases and collapses, leading to training instability, we provide evidence that this results in an increase in the norm of the attention probabilities, causing gradient exploding (as mentioned in Section 3.5). Additionally, attention entropy collapse across various tasks has already been explored in previous studies (Zhai et al., 2023; Shen et al., 2023; Jiang et al., 2023). These experiments were conducted on models and tasks of a scale commonly used in practical applications. In contrast, our paper conducted experiments in a more controlled setting to analyze why entropy collapse occurs and how it leads to training instability.

[Q3] Under what setting have the experiments shown in Figure 2 been performed?

• As explained in the response to W2, the data was sampled accordingly, and the experimental settings are provided in the Appendix A.

[Q4] Is there a specific reason that leads to stable training for Lipschitz-Kernel attention as compared to softmax-based attention?

• We focused on the attention entropy collapse occurring within self-attention and accordingly used an architecture composed solely of self-attention layers. Therefore, we cannot determine if other factors outside of attention play a role in this collapse.