softmax is not enough (for sharp out-of-distribution)

Dear Reviewer 4NvJ,

Thank you so much for acknowledging our response and amplifying your score towards acceptance! We highly appreciate that our new theoretical content is seen as useful.

Your remaining presentational suggestions are indeed minor and we are happy to address them in the final revision.

Regarding your remarks on the theory positioning:

The contents in Appendices B and C add substantial value to the paper (much more so than the proofs of Lemma 2.1 and Theorem 2.2). I would suggest moving these proofs to the appendix, and instead highlighting some of the key results in the main text.

We appreciate your suggestion (and are very happy that you feel that the Appendices add valuable content!), though we would like to offer a few points of discussion here:

• We would prefer to keep the proof of Lemma 2.1 in the main body of the paper. This is because it introduces the all-important logit spread quantity, $\delta$, that is used throughout the paper afterwards. It's not easy to motivate why this quantity is important without the proof being there, in our opinion.

• The case for the proof of Theorem 2.2 is more interesting:

  • It is our opinion that this an important result -- all of the subsequent results we added on the reasoning failures of softmax depend on it being true -- and therefore it is very important that the reader believes us that there are no hidden caveats in the theory.
  • Also, this is not a result that feels intuitive at first sight (at least, it didn't feel intuitive to us at first).

In comparison, the results in Appendices B and C are indeed important, but we would argue that, once we already know that the coefficients must collapse in Transformers, there is an intuitive sense in which this is going to be bad for the network at some point. At least to us, the results in these Appendices are hence less surprising ("lower entropy") in light of Theorem 2.2 being true, and this inspired the current way we layout the results in the revised paper.

Because of the reasons above, and the constrained space in the main body of the paper, it is our current belief it would do more harm than good to relegate the proof of Theorem 2.2 into appendices. Of course, this is not a belief that cannot be swayed, especially if you feel very strongly about this -- but we believe our case is sensible.

Secondly, on your remark about checking the theoretical work:

What you wrote in the response makes intuitive sense to me. However, I must admit that I might not have the bandwidth to check the proof of the new results in detail.

We are very happy that the response makes sense, and highly value your time. If you do get some bandwidth in the remaining time to look at the results more, we would suggest focussing on checking the proof of Corollary B.1. Once this result is established, all of the subsequent Remarks either directly follow (as relatively natural specific cases), or require minor analysis on the dataflow of the various Transformer architectures.

Finally, concerning your remark about the changes we made on the whole:

I believe (and I expect the authors would agree) that the changes made are substantial. It may be best for the AC/SAC to decide whether another full round of review is warranted.

We do agree that the changes are significant in their volume, though we also believe that there is a case to be made that they are not as substantial in their substance. This is due to the fact that a lot of the modifications we made concern clarifying definitions that already existed in the paper at first (e.g. dispersion, sharpness, OOD), making the paper more self-contained by covering well-known material (e.g. introduction to Transformers), or numerically validating facts we've already proved (e.g. Proposition 3.1.).

It is our desire to be fully transparent about this and invite any debate while we still have a chance to respond to it. As such, later today, we will post a global response (addressed to all Reviewers and the AC) explicitly enumerating the changes made to the paper, along with our opinion on the significance of their substance.

That being said, we are hopeful for your support towards acceptance, given that you found the additional theoretical results to be intuitive and following from Theorem 2.2 (which was already in the paper at the start). We are happy to provide any additional or pointed insight that can resolve any outstanding doubts you have -- we still have plenty of time left.

We sincerely thank you once again for kindly considering our rebuttal and amplifying your score!

Dear Reviewer 4NvJ,

For starters, we would like to give you particular thanks – while your review was negative in its rating, we’ve been able to clearly detect positive sentiment about our work, and highly pointed suggestions for how to improve our work. Many of these suggestions were indeed very much called for.

We truly believe that, as a result, we now have a much stronger contribution on our hands – and we hope you will agree! The paper has already been revised on OpenReview. We summarise specific things we have done in response:

Improving the presentation

We believe that we have now addressed all of your remarks on presentation:

• A full primer on Transformers (which also defines attentional heads and coefficients) has been added within our Introduction section. We also more precisely defined dispersion, sharpness and the specific out-of-distribution generalisation setting we study in response to the other reviewers – all of these are in the Introduction.
• The captions of Figures 2, 3, and 7 have all been amended to properly define the various terms you called out: batch / rows (Fig. 2), curves and shaded blue region (Fig. 3), x-axis and individual curves (Fig. 7).

In addition, we now have a more proper description of how adaptive temperature is applied, both in the main body of text:

That being said, there is an alternate route to make the Gemma model still benefit from our adaptive temperature module exactly as-is (i.e., with exactly the same polynomial fit as in Figure 5); it just has to directly learn how to leverage it. As such, in our CLRS-Text ablation we apply adaptive temperature both during fine-tuning and at inference time. What this means is, we replace all instances of jax.nn.softmax within all the attentional heads of Gemma 2B with our adaptive\_temperature\_softmax function, both during fine-tuning of the model and during inference. This allows the model to learn how to compute key/query embeddings that can maximally exploit the temperature adaptation.

Further, in Appendix F, where we now detailedly describe a – to the best of our knowledge – novel algorithm for computing the entropy values in a streaming manner (akin to Flash Attention), allowing us to compute adaptive temperature in $O(n)$ space complexity and deploy it on large context lengths – up to 131,072 tokens on a single A100 node.

Improving the theoretical conclusions

As already mentioned, we really appreciate your pointed feedback. As such, we will dedicate the majority of our rebuttal response to you on how we improved the theoretical conclusions in the previous couple of weeks—we believe we have managed to meaningfully address all of your questions. (For what it’s worth, it is also our belief that we’ve done useful things (e.g. the aforementioned streaming algorithm for scaling adaptive temperature) to strengthen the algorithm/experiment points as well.)

it is unclear what the consequence of such "dispersion" of attention coefficients is: does it imply the failure of the underlying reasoning task?

Yes it does! We prove (Corollary B.1 and Remark B.2 in Appendix B) that attentional dispersion will be a provable culprit of misclassifications for even rudimentary reasoning tasks such as finding maximal tokens. We elaborate a lot more in our proof, but the core idea is to leverage pairs of input sets of the form {A B B … B} (max is A) and {B … B} (max is B). Under dispersed softmax, for sufficiently many Bs, at least one of these sets must be misclassified, as soon as the dispersion decay surpasses machine epsilon. Note that modern LLMs are often heavily quantised when deployed for fast inference, so effective machine epsilon may be significantly higher than what one might expect e.g. for single-precision floating point types.

Will it still be problematic if the ground truth token has a coefficient of $O\left(\frac{\log n}{n}\right)$ while the other tokens have coefficients of $O\left(\frac{1}{n}\right)$, where their ratio still goes to infinity?

Yes (and we elaborate on this in Remark B.3). There is nothing in the proof of Corollary B.1 that depends on the rate of dispersion—so long as there is decay to zero, the attentional mass on the sole ‘A’ element in {A B B ... B} will eventually disappear below machine epsilon.

We also remark that Theorem 2.2 shows we cannot get this kind of a situation (a single token decaying with rate $O\left(\frac{\log n}{n}\right)$) in the current way we build frontier Transformer language models, though it might be very interesting to study what kinds of architectures might induce such behaviours!

We continue our response in a new message, due to character limitations.

Thanks again for the response and the summary of the changes.

• I think it is fine to leave the proof of Lemma 2.1 as it is, but I genuinely feel that Theorem 2.2 is a straightforward corollary of Lemma 2.1. Perhaps you could elaborate on why it doesn’t feel intuitive to you? I might be biased (having seen the proof before thinking about the problem from scratch), but to me, the proof is all about "the continuous mapping of a compact set remains compact," which is essentially college-level math.

• On the other hand, the connection between the dispersion of attention coefficients and the failure of reasoning tasks was not intuitive to me—and still isn’t. Here’s why: although all attention coefficients are on the order of $\frac{1}{n}$, the constants could differ significantly. This is exactly why I brought up the $O(\frac{\log n}{n})$ example. My intuition is that as long as the ground truth token corresponds to a large constant, there is a high chance that the model will function correctly.

• Finally, I briefly checked the proof of Corollary B.1.

  • The statement: "At this point, the value of $\alpha_1^{(n)}$ will be indistinguishable from zero, and the weighted sum will reduce to $g(f(v_b))$" lacks mathematical rigor. The term "indistinguishable" is vague, and it’s unclear what assumptions are being made. To conclude properly, you need to assume that $g$, which maps a vector to a class, is continuous at $f(v_b)$. While this is generally true, care must be taken: (1) points on the decision boundary of two classes require special consideration, and (2) this continuity property/assumption should be explicitly stated in the proof.
  • The example is not very general and even in some sense degenerate. In particular it still does not explain why my previous intuition regarding the large constant might be wrong. I was wondering if this is the best result you can achieve, or do you think it’s sufficient for now and stop exploring?

Thank you for validating our suggestion, the paper is now revised with the logit inequality assumption explicitly called out.

We highly appreciate the time you have invested in helping us strengthen our theoretical contribution, which we do consider to be among the key takeaways of our work.

And of course we are very happy about your support for acceptance based on our interactions! We will make sure to triple-check all of our arguments.

Dear Reviewer CFHF,

We would like to thank you for your positive feedback as well as insightful pointers on our theoretical findings. We have now revised our paper in a manner that should address your concerns better, and have follow-up thoughts on some of your points:

Defining dispersion

Thank you so much for raising this point – we fully agree, and have now concretely specified what we mean by dispersion in the paper:

Through one of our key theoretical results (Theorem 2.2), we demonstrate that modern deep learning architectures, operating over a fixed vocabulary of input tokens and leveraging the softmax function, are fundamentally incapable of learning functions that remain sharp under such out-of-distribution instances. This is due to the fact that the coefficients emitted by the softmax function must disperse as we increase the number of input items. Here by dispersing we mean that, as the number of input items grows, the coefficient attached to each individual item must decay towards zero. This makes it impossible to robustly compute functions that depend on any particular finite amount of input values, such as the aforementioned max, as we show in Appendix B (Corollary B.1 and Remark B.2).

On the use of $\Theta\left(\frac{1}{n}\right)$

We follow the usual definition of asymptotic growth: a function $f(n)$ is in $\Theta\left(\frac{1}{n}\right)$ if there exist constants $A$ and $B$, such that, as $n\to\infty$, $\frac{A}{n}\leq |f(n)|\leq\frac{B}{n}$.

The dependence on $\delta$ can be abstracted away in the big-Theta notation because $\delta$ is assumed to be a constant – and Theorem 2.2.’s main contribution is, among other things, proving that $\delta$ must be a constant in modern Transformer architectures over tokenised vocabularies.

On the “surprise” of the distribution shift result

We fully agree with you that it is unsurprising that models tend to behave worse under certain distribution shifts. What is surprising about our result is that our results hold under any statistical distribution of added items, so long as sufficiently many are added. For example, we could add many identical irrelevant items, or a mix of randomly sampled items, and dispersion would still occur in both cases! This, in our opinion, is not a necessarily expected result.

In terms of quantification of how softmax responds to these distribution shifts, we believe that Lemma 2.1. provides a useful quantification of how the dispersive decay rate is controlled, in terms of the spread, $\delta$. We now explicitly call this out in the paper:

The difference of these bounds (the spread, $\delta = \max_{i} e^{(n)}_i - \min_j e^{(n)}_j$) directly controls the rate of dispersion.

On whether Theorem 2.2 considers the worst-case scenario

From our understanding, Theorem 2.2. does not consider only worst-case scenarios, but actually the best-case scenario.

It proves dispersion will happen out-of-distribution in every self-attention architecture using softmax to attend over all items, no matter how well it has been trained or which dataset has been used to train it.

Please let us know if we misunderstood something!

Numerically verifying Proposition 3.1

Indeed we can do this, and we have! You may find a plot of our upper bound and the empirically observed spread values, over the evolution of a trained model, in Appendix D of the revised paper (Figure 8).

Adaptive temperature only slightly improves the performance

In our opinion, Figure 7 demonstrates a more significant outperformance of adaptive temperature; namely, it is apparent that the Gemma 2B model gets better at executing several algorithms in CLRS-Text, with gains of 40+% accuracy on several input sizes, when it gains access to adaptive temperature. Since this is the only modification to the architecture, we can attribute these gains to learning to make good use of the adaptation scheme.

Please let us know if there are any other discussion points you have for us during the remainder of the rebuttal period – we are always happy to engage further.

Dear Reviewer kvZv,

We would like to thank you for your careful review of our paper, and the ample suggestions for improvement! We hope that our recently-uploaded revision satisfies your concerns.

We also provide more detailed responses to accompany those revisions:

Improving the definition of sharpness

We fully agree that our previous definition of sharpness did not clearly imply why max is sharp. We have now amended the definition in the paper, with a clear rationale:

Here we call a function taking a variable number of inputs sharp if its output value can be expressed using only a constant number of these inputs. For example, max is sharp, as its output value is equal to exactly one of its inputs'. The average function is not sharp, as its output value depends on all of its input values (with factor $1/n$ for each of the $n$ items).

Clear definition of our out-of-distribution setup

Thank you very much for this suggestion! We wholeheartedly agree, and have now added this paragraph into the paper, drawing directly on our sharpness redefinition:

We define sharp functions by their behaviour as their number of inputs varies. This directly motivates the out-of-distribution setting we study: generalising to different amounts of inputs. Specifically, when we analyse neural networks that learn sharp functions, we assume that they are trained on problem instances containing no more than $n$ input items, and we take a particular interest in their sharpness on instances with $n' > n$ items; these are considered out-of-distribution instances because they go beyond the maximal number of inputs the model had been prepared for. In language modelling, this setting is also known as length generalisation \citep{anil2022exploring}; in graph machine learning, it is known as size generalisation \citep{yehudai2021local}.

On the scalability of adaptive temperature

Thank you very much for raising this issue! Your remark inspired us to investigate how to most efficiently implement adaptive temperature, and we have successfully developed an iterative streaming algorithm for computing the entropy values on-the-fly, in a manner akin to Flash Attention.

This algorithm allows us to compute adaptive temperature in linear space complexity. Using it, we have successfully scaled adaptive temperature on large context windows – up to 131,072 tokens on a single A100 node. We now fully describe this algorithm in Appendix F. To the best of our knowledge, this algorithm is novel.

On the generalisability of conclusions

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 tasks 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.

Please let us know if there are any other discussion points you have for us during the remainder of the rebuttal period, or anything we can do to improve your opinion of the paper in the time remaining – we are always happy to engage further.

How is adaptive temperature implemented?

We now elaborate in more detail on how this is done in the paper:

That being said, there is an alternate route to make the Gemma model still benefit from our adaptive temperature module exactly as-is (i.e., with exactly the same polynomial fit as in Figure 5); it just has to directly learn how to leverage it. As such, in our CLRS-Text ablation we apply adaptive temperature both during fine-tuning and at inference time. What this means is, we replace all instances of jax.nn.softmax within all the attentional heads of Gemma 2B with our adaptive\_temperature\_softmax function, both during fine-tuning of the model and during inference. This allows the model to learn how to compute key/query embeddings that can maximally exploit the temperature adaptation.

Additionally, we have been able to deploy the entropy computations in a streaming manner (akin to Flash Attention), allowing us to compute adaptive temperature in linear space complexity, scaling it up to context windows up to 131,072 tokens on a single A100 GPU. The detailed elaboration of this algorithm is provided in Appendix F and, to the best of our knowledge, it is novel.

Please let us know if there are any other discussion points you have for us during the remainder of the rebuttal period – we are always happy to engage further.

Dear Reviewer HP5U,

We would like to thank you for generously supporting our submission! Additionally, your review raises important points of discussion, which we now address in a revised paper.

On the significance of dispersion

We would like to reiterate that the key aspect of our study is out-of-distribution behaviour of softmax coefficients. We now also explicitly call this out in the paper, in the following passage:

We define sharp functions by their behaviour as their number of inputs varies. This directly motivates the out-of-distribution setting we study: generalising to different amounts of inputs. Specifically, when we analyse neural networks that learn sharp functions, we assume that they are trained on problem instances containing no more than $n$ input items, and we take a particular interest in their sharpness on instances with $n' > n$ items; these are considered out-of-distribution instances because they go beyond the maximal number of inputs the model had been prepared for. In language modelling, this setting is also known as length generalisation \citep{anil2022exploring}; in graph machine learning, it is known as size generalisation \citep{yehudai2021local}.

With this in mind, we briefly comment on a few of your observations:

In various prior work studying transformer mechanisms, the heads appear to be sharp.

In all of these prior works, the heads were studied in-distribution examples (oftentimes, these may have even been in the training data!). We now explicitly call this out in the paper:

...and the discovered heads always appear sharp when inspected on in-distribution samples

In conclusion:

Whether softmax is "sufficiently" sharp is, after all, dataset and problem dependent.

We wholeheartedly agree, though in light of our study, we make explicit that it provably cannot be sufficiently sharp for problems requiring robust behaviour under size distribution shifts.

This is exactly the reason why we chose the CLRS-Text dataset to study the effects of adaptive temperature---it is a collection of thirty challenging algorithmic execution tasks where out-of-distribution generalisation is explicitly required. In most other static benchmarks, it might be very difficult to measure the distribution shift in the test set.

Important real-world scenarios where we often run into distribution shifts include tasks such as scientific discovery and agentic behaviours---both known to be rather challenging for LLMs today, and also really challenging to rigorously evaluate.

On the utility of adaptive temperature

Your review rightfully realises that adaptive temperature is not meant to be a solution to the problem in Section 2 (which has now been strengthened with additional Corollaries in the Appendices), but rather a diagnostic approach to validate some of the implications of our theory and hopefully stimulate some future research directions.

As such, we do not intend to dwell significantly on these points, though we would like to briefly respond to two of your remarks:

As noted by the authors themselves, the correct adaptive function for $\theta$ is dataset-dependent and determining what this function is can be highly non-trivial in attention models.

This is true, though our CLRS-Text results leverages the same adaptive temperature scheme we deployed specially for max-retrieval (though with additional fine-tuning), and we believe our results indicate that the model learnt to use such an adaptive temperature just fine given additional training signal. With that in mind:

It is not clear from the empirical results that merely having an adaptive temperature parameter is a meaningful fix to counter the dispersive tendencies of softmax. E.g. the results in Table 1 are mostly incremental and the visual differences in Figure 6 are minor.

In our opinion, Figure 7 is also an important aspect of these results: it is apparent that the Gemma 2B model gets better at executing several algorithms in CLRS-Text, with gains of 40+% accuracy on several input sizes, when it gains access to adaptive temperature. Since this is the only modification to the architecture, we can attribute these gains to learning to make good use of the adaptation scheme.

We continue our response in a second message due to character limitations.

I thank the authors for their detailed response to my review, as well as to the other reviewers. Overall I accept the comments the authors have made and maintain my score.

As to AC questions, OOD literature is indeed very broad, and my understanding is that this paper focuses on a specific kind of OOD (generalization to increasing problem lengths) and specific model architecture (although as a cursory remark, it does not seem to be too difficult to generalize to other architectures if compactness is the only requirement).

This does limit the scope of the proof and so the impact of the paper. Nevertheless, I think the authors have done a good job tightening up the proofs (including the post-rebuttal edition) and offering a potential solution. More so, I think the paper is an important step in the direction of probing the inductive biases that individual (and often overlooked) modeling choices might make, and how these biases can affect specific kinds of distribution shifts. I think it would be a helpful step towards future research in this direction.

Dear Reviewer H5PU,

Thank you so much for engaging with our response and accepting our comments. We are delighted that you are maintaining your acceptance score.

Just to note we are unable to see the "AC questions" being referenced here, but the points you raised are worth briefly replying to, in case they are helpful!

We indeed focus on size / length generalisation in this work and focus on the most currently used attentional architectures in our proofs (BERT / GPT / Set Transformers).

While this is of course only one kind of distribution shift, it is a shift that has been very extensively studied in recent years (as evidenced by the vast body of literature on length generalisation in LLMs [1, 2, 3, 4...] or size generalisation in GNNs [5, 6, 7...]).

We agree with the reviewer that, as our proof relies on compactness and continuity, it should easily work with many other attentional architectures relying on softmax. To be fair, already the compositional formulation of the architecture in Theorem 2.2 may be sufficient to describe most such architectures (though it of course most naturally aligns with the Transformer)!

Nevertheless, I think the authors have done a good job tightening up the proofs (including the post-rebuttal edition) and offering a potential solution. More so, I think the paper is an important step in the direction of probing the inductive biases that individual (and often overlooked) modeling choices might make, and how these biases can affect specific kinds of distribution shifts. I think it would be a helpful step towards future research in this direction.

We are extremely grateful for your assessment, and we fully agree about the future potentials of the work.

Best, Authors

[1] Anil et al., "Exploring Length Generalization in Large Language Models", NeurIPS'22

[2] Ruoss et al., "Randomized Positional Encodings Boost Length Generalization of Transformers", ACL'23

[3] Chowdhury et al., "Monotonic Location Attention for Length Generalization", ICML'23

[4] Zhou et al., "What Algorithms can Transformers Learn? A Study in Length Generalization", ICLR'24

[5] Bevilacqua et al., "Size-Invariant Graph Representations for Graph Classification Extrapolations", ICML'21

[6] Yehudai et al., "From Local Structures to Size Generalization in Graph Neural Networks", ICML'22

[7] Buffelli et al., "SizeShiftReg: a Regularization Method for Improving Size-Generalization in Graph Neural Networks", NeurIPS'22

Dear Reviewer HYsK,

We would like to thank you for the support you have expressed towards our paper as well as calling out important points where additional mathematical clarity would be required.

We wholeheartedly agree and have now addressed both of the issues you raised in our revised paper:

On sharpness

We redefine sharpness in the following way, along with a more clear elaboration on why max is sharp:

Here we call a function taking a variable number of inputs sharp if its output value can be expressed using only a constant number of these inputs. For example, max is sharp, as its output value is equal to exactly one of its inputs'. The average function is not sharp, as its output value depends on all of its input values (with factor $1/n$ for each of the $n$ items).

On feature spaces

Thank you for calling out this issue; we believe that the misunderstanding stemmed from the fact that $\mathcal{X}$ is a feature space rather than a vector space. We have now corrected this to read:

Let $\mathcal{X}\subset\mathbb{R}^m$ be a set of possible $m$-dimensional input features…

Please let us know if there are any other discussion points you have for us during the remainder of the rebuttal period – we are always happy to engage further.

Dear Reviewer HYsK,

Thank you for elaborating on your position, and we appreciate your suggestions. We would also like to do our best to improve the chances of the paper.

The definition above is not mathematically rigorous to me, and it will potentially hurt your theoretical results (Lemma 2.1. for example, which only holds asymptotically and you cannot really tell what would happen if $\infty > n' > n$).

Could you please elaborate on why Lemma 2.1 cannot tell us what happens for $n' > n$?

The asymptotic behaviour discovered in Lemma 2.1 is obtained through first deriving an established upper and lower bound on the coefficients (Equation 8), which depends on the number of logits, $n$.

Therefore we can use the derivation of Lemma 2.1 to exactly give upper and lower bounds for coefficients for any $n' < \infty$ we want.

In fact, we use this exact bound in our proof of Theorem 2.2, to not just say that all heads must disperse, but give an exact lower bound on $n' < \infty$ after which dispersion below any given $\epsilon > 0$ must happen.

With this in mind, do you have any suggestion for us on how to improve the rigour of the definition?

However, other reviewers are complaining about the triviality of your theoretical results and their lack of insights (not me, I simply see this paper as an empirical paper and the added results are just the cherry on the cake) so you might want to make it more concrete in that part to avoid being rejected with a good score.

We presume you are referring to Reviewer 4NvJ.

Just to remark, we had a very useful discussion with the Reviewer during the rebuttal period and have managed to provide additional theoretical results, after which the Reviewer improved their score. We are unaware of other Reviewers calling out triviality or lack of insight of the theoretical results (though we haven't yet had a chance to engage with the other Reviewers).

Best, Authors

Clarified definitions of key terms (Section 1)

• Sharpness
• Dispersion
• Out-of-distribution

These are all terms we used and leveraged throughout our original paper. The updated definitions serve to make it very clear (e.g. using mathematical concepts) what these terms mean for the scope of the paper---they do not add any information to the paper that was not at least implied originally.

Primer on Transformers (Section 1)

We now provide a full introductory material on Transformers, attentional heads and coefficients, prior to using them in our analysis. This constitutes standard knowledge in LLM architecture design, and serves to better introduce the reader to the specific notations we use in the paper. Hence, we do not consider this to be a semantic change to the paper, merely one that improves the ease-of-reading.

Clarifications in various descriptions:

• Clarifying that the logit spread, $\delta$, controls the constant of the dispersion rate (Section 2)
• A minor modification to the definition of a feature space using sets rather than spaces (setup of Theorem 2.2)
• A more precise description of how Gemma 2B experiments leverage adaptive temperature (Section 4.2)
• Improving descriptiveness of figure captions (Figure 2, 3 and 7)
• Adding a few recent relevant related work citations across the text

These correspond to minor (often single-sentence level) modifications to the text, and they mainly serve to clarify the description that was already there during the initial submission. We do not consider this to be a semantic change to the paper, mainly a change that improves convenience.

Further theoretical results (Appendices B and C)

We have derived a useful corollary of Theorem 2.2 (Corollary B.1, Remark B.2, Remark B.3) showing that attention head dispersion in Transformers provably leads to reasoning failures on simple tasks such as max-retrieval. This directly ties our observation to OOD failure modes of language models, and is hence valuable and important.

That said, the line of proof used is relatively straightforward given Theorem 2.2:

• We know (Thm 2.2) all attention heads' coefficients will eventually disperse towards zero
• Take a set with one maximal item and $n$ smaller items
• At some point, the attention coefficient on the max item will be negligible against the total coefficient on other items.
• Once this coefficient falls below macheps, the set is indistinguishable from a set containing only smaller items, hence those two sets must be classified the same way by a Transformer.
• This implies at least one of them must be misclassified.

And it should be reasonably intuitive that attention head dispersion can lead to undesirable outcomes. Hence, these corollaries' aim is to validate already established intuition in the original paper, rather than introduce brand-new theoretical advances. As such, we hope they will be seen as appropriate for the rebuttal period.

We also provided an informal analysis of how vulnerable different attention heads are to dispersion as a function of their depth. This mainly relies on an intuitive observation -- dispersion of a given layer will bring embeddings closer to an averaged embedding, limiting embedding (and hence logit) spread, and accelerating dispersion elsewhere. In addition, we informally proved (Remark C.1) that for a specific BERT-style architecture, full dispersion in layer L immediately implies full dispersion in all future layers.

As above, we find these observations to be reasonably intuitive given the results in Theorem 2.2 -- further, the idea of amplifying or attenuating logit spread was already explored in the original paper, by Proposition 3.1 -- and hence we also consider them validation of these intuitions, and hence appropriate for the rebuttal period.

Numerical validation of Prop 3.1 (Appendix D, Figure 8)

We show that Prop 3.1 holds in practice during training of self-attentional models. As we already proved this bound in the original paper, this is just a sanity check and we do not consider it to constitute a substantial change.

Streaming algorithm for attentional entropy (Appendix F)

This is an algorithm allowing for attentional coefficient entropy to be computed in an online manner, allowing us to serve adaptive temperature in a manner akin to Flash Attention, enabling highly scalable $O(n)$ memory usage for a context window of $n$ tokens.

While this significantly improves the applicability of our adaptive temperature proposal to large language models, the result is derived mainly through careful algebraic manipulation of the individual terms; all steps of the algorithm can be easily checked for correctness. We also numerically verified that the outputs of the online algorithm match (are sufficiently close to) the naïve JAX implementation provided in the original paper.

All things considered, we find this change to be appropriate for the rebuttal.