Softmax is not Enough (for Sharp Size Generalisation)

Dear Reviewer n5x1,

We are highly pleased to read your review, and really appreciate your positive view on our results and their significance!

In what follows, we reply to all of the points you raised:

On the other hand, the adaptive temperature sampling scheme is less convincing, more ad hoc, and leads to seemingly small gains on the problems tested. Therefore, I overall find highlighting the dispersion phenomenon to be the more valuable contribution in this paper, since it may well motivate future work on this topic.

We fully agree with you that the key outcome of our work should be highlighting the dispersion effect, improving understanding of it, and stimulating future work towards addressing it. Adaptive temperature was designed as a mostly ad-hoc method to illustrate that even simple interventions can counter dispersion in a way that leads to measurable improvements, but it does not escape the confines of our theoretical results – as we clearly highlight in our Conclusions.

I had some trouble following the explanation of the adaptive temperature sampling scheme.

We appreciate this remark and commit to adding a new section (potentially in the Appendix) which will provide a step-by-step overview of how the sampling scheme was arrived at.

However, in particular this sentence could be improved: "We prove this important result now" before Theorem 2 could be changed to simply read "We prove this result now".

This is a great suggestion, and we will tone down by removing the word ‘important’ here.

Dear Reviewer jdu5,

Thank you so much for the very careful review and the vast amount of useful suggestions for improving our work! We are very happy you appreciated our efforts to raise awareness of the dispersion issues of softmax in the ICML community!

We provide detailed answers to your comments below, and we hope they will be to your liking.

Bound looseness

We highly appreciate your remark about the looseness of our bound, and wholeheartedly agree it is worth discussing in more depth. We stratify our answer by your specific comments, and commit to updating the paper to incorporate all aspects of our discussion!

The experiments in the paper go up to n=16,384=2^14, while machine precision limits would only provide a lower bound of m=−10^38=−2^126, so this sequence length and lower bound is going to induce minimal dispersion on the attainable sharpness.

Just to assert our understanding: due to the exponents you presented, we find it likely that your bounds were derived under the assumption of the bfloat16 type being used. This is a fair assumption, though we remark that large-scale frontier models are frequently served at even lower-precision types – sometimes with aggressive quantisation. This might well affect the practical lower bound obtained for $m$ due to machine precision.

The proofs provided by the authors demonstrate that softmax must disperse in the limiting case, but not that it does in practice for the scales actually used with transformer models.

In this regard, it is worth taking Theorem 2.2 together with Proposition 3.1, which shows models aiming to improve sharpness have no choice but to amplify their weights at some point in the architecture, and that the weight magnitudes directly control the empirically observed values of $m$ and $M$. But amplifying weights is generally risky due to the danger of overfitting—as such, the model has no choice but to keep the empirical spread somewhat contained. This immediately brings us to your next point:

For example, what is the distribution of logit values seen in LLM attention heads at inference time across a range of tasks? The minimum and maximum values from this would be informative to establish a practical range for m and M.

This is a fantastic suggestion and we really appreciate it!

To measure this, we fed an entire code sample of Gemma’s Modules file (available at https://github.com/google-deepmind/gemma/blob/main/gemma/modules.py, ~4,000 tokens) to Gemma 2B and 7B models. The empirically observed values of the logit spread, $\delta = M - m$, across all attention heads, were as follows:

This shows that empirical logit spreads in practical queries are, even in their maximal occurrence, rather small compared to the machine epsilon-induced bounds mentioned before, and should give further credibility to the practicality of our results. We of course will include this table in the paper!

Additionally, I think there could be more discussion on the factors which prevent the model from achieving an arbitrarily sharp distribution

Another excellent proposal – we fully agree with all of your specific factors, and will make sure to thoroughly discuss those in our revision.

Specifics of Adaptive-$\theta$

As suggested, we have performed preliminary experiments with a custom oracle multiplier for Adaptive-$\theta$. While this led to (2–3%) percentage point improvements on lesser OOD sizes, there was no consistent improvement on longer sequence lengths.

We are also happy to enrich the discussion of the comparisons in Figure 8. While we cannot make very strong claims, a unifying property of Heapsort, MST Kruskal and Bubble Sort in CLRS-Text is that they all occupy relatively large chunks of Gemma 2B’s context window, which stretches further beyond the largest contexts over which the polynomial fit of Adaptive-$\theta$ was fit—and this might cause an unintended shift, which is in line with your suggestion.

On Figure 5

• We acknowledge your point regarding jet. We agree that perceptually uniform colormaps like viridis or plasma offer better visual representation, and will switch!
• We used power series (controlled by $\lambda$) as a simple model to represent sequences with varying degrees of "sharpness" in their logit distributions.
• While negative $\theta$ does not have a direct thermodynamics meaning, it allows us to explore the behavior of the softmax function beyond the typical regime.

Miscellaneous issues

We are very grateful for your thorough comments about several miscellaneous minor issues in the paper. We fully agree with your remarks, and commit to correcting them in our revision.

Dear Reviewer 5ZBY,

Thank you for your careful review and the positive assessment of our contribution! We are very grateful for your comments, and provide our responses below – hoping that they are to your liking!

The authors clearly discuss the disperse problem of softmax function. I wonder how big the problem is when we zoom out to the entire transformer, given all the residual connections and normalization etc?

This is an excellent question!

We study exactly this, to varying extents, in Appendix B and Appendix C of the original submission. We provide a brief summary of these results for your convenience, and are very happy to discuss them further:

• [Corollary B.1] We prove that the dispersion effect necessarily leads to classification failures in a single-layer Transformer architecture on simple reasoning tasks (such as predicting the maximum).
• [Remark B.2] We make an informal sketch of how Corollary B.1 can be extended to deep Transformer architectures (both BERT- and GPT-style) to show the same kind of classification failures must occur past a certain number of input tokens. This argument explicitly takes into account residual connections.
• [Remark C.1] We prove that, in BERT-style Transformers without residual connections, the situation is particularly dire: when a particular layer disperses, all layers after it will be immediately dispersed as well.

The residual connections play an important role with depth, as evidenced by Remark C.1: they allow a model to “shortcut” a dispersed layer and retain its original embeddings for longer. However, Theorem 2.2 shows that, no matter how many residual connections are used, each individual layer still must disperse past a certain size.

The normalisation layers often play a counter-productive role, which we discussed in Section 3: they clamp the input to a certain expected norm, meaning that there is higher pressure on key/query matrices ($\mathbf{K}$, $\mathbf{Q}$) in order to achieve a sufficient logit spread (cf. Proposition 3.1).

However, the evaluation is rather "toy". I understand that max retrieval provides a clean setting. However, only CLRS-Text makes it hard to understand the usefulness of adaptive temperature or the problem from the dispersed nature of softmax in the real world.

Thank you for your remarks! Since our study concerns out-of-distribution generalisation specifically, we have focused our analysis on tasks requiring out-of-distribution generalisation (such as CLRS-Text, a collection of thirty challenging algorithmic execution tasks across many problem sizes). In most other static benchmarks, it might be very difficult to measure the distribution shift in the test set.

We also remark that focusing on synthetic execution tasks is the standard approach in papers studying length generalisation in LLMs. As a standard representative we refer to “Exploring Length Generalization in Large Language Models” (Anil et al., NeurIPS’22), which studies only two synthetic problems: parity and variable assignment. In contrast, CLRS-Text studies thirty such problems, with a significant increase in their complexity.

Dear Reviewer WiLd,

We would like to thank you for carefully considering our paper. While we regret that your initial rating of our paper was negative, we believe you raised important points and that there is a clear discussion to be had, and that we may be able to provide relevant arguments for you to reconsider the relevance of our work.

To that end, we address all of your points in order:

On significance of our results

We do not dispute that this result was not very difficult to prove, but we argue that it is definitely not evident to a significant part of the ICML community. And we hope you agree we should preferably judge proofs’ significance using the latter criterion rather than the former -- simplicity is not in and of itself bad.

Therefore, we will focus this part of the response on discussing whether our results are evident.

Due to the anonymous nature of the reviewing process, the only concrete evidence we can provide towards this are the reactions of the other three reviewers:

• [Reviewer n5x1] The paper's observation is simple, but thought-provoking.
• [Reviewer jdu5] I appreciate that this work presents the problem well and raises awareness of this issue.
• [Reviewer 5ZBY] The limitation on softmax and adaptive temperature is relevant for the boarder audience, given the prevalence of softmax in modern ML system

We wish that we could provide more concrete evidence, but we cannot for obvious reasons. Suffice it to say, we have had numerous discussions with our previous collaborators (who are at many varying levels of seniority and expertise about self-attention) about this result, and the overwhelming majority of them initially reacted to our result with surprise; i.e., not expecting that the dispersion effect is guaranteed in softmax. In fact, it was exactly these interactions that compelled us to write this paper in the first place!

It is interesting to ponder why our result is surprising to such a broad audience of AI researchers. Our hypothesis is that this is due to several current trends:

• The overall prevalence of the Transformer architecture and the many expressivity results for it;
• The elevated importance of the dataset for training AI models, and the diminished importance of architecture choice;
• Mechanistic interpretability research, which reverse-engineers sharp behaviours in trained LLMs.

Such trends may easily lead to a naïve intuition that softmax should always be able to pick out the key elements / circuits to apply to the data, so long as we choose the “right data” to train on.

However, the length generalisation setting challenges this preconception, because it by default focuses on evaluating beyond the largest training data point. Further, mechanistic interpretability research typically does not operate in such regimes, and the circuits discovered therein do not generalise to ever-increasing inputs.

We believe our paper plays an important part in grounding the limitations of the softmax function, especially when considering how it is leveraged in modern LLMs. We hope you will agree with this motivation!

On thresholding contributions

We appreciate your comment and believe addressing it will improve the rigour of our argument! In our revision, we will explicitly mention thresholding contributions when defining sharpness, and how our Theorem 2.2. proves that no fixed threshold ($\epsilon > 0$) is sufficient to maintain sharpness on larger inputs.

On the infinity dimensions and the ‘scaling up sharpness’ framing

To be clear, our work does not assume infinitely deep or wide architectures. We start with the practical assumption of a model of fixed depth & width, and then quantify how its coefficients’ sharpness decays with respect to input size – exactly as you suggested. We commit to adding further sentences around the problem description to make this crystal clear.

Improvements to softmax

Our argument is slightly more nuanced than what you suggested. We suggest that it is the combination of the softmax function and how it is used within Transformers (e.g. tokenisation, global attention, etc.) that causes limitations over arbitrarily large inputs. That is, certainly the Transformer itself could be improved, but there are also possibilities of improving the softmax function itself.

We cited several examples of possible alternative aggregation functions in the Conclusions: linear attention, sigmoidal attention, and stick-breaking attention. There are also proposals such as selective attention (Leviathan et al.), which retain softmax but modify the algorithm which allocates logits in a more size-robust manner. For several of these proposals, Theorem 2.2 would not apply.

We are happy to add this discussion to the revised paper!

Miscellaneous issues

We are happy to correct all minor nits you pointed out, as well as avoid usage of terms like ‘reasoning device’ in a revised paper.