Softmax is not Enough (for Sharp Size Generalisation)
Authors: Petar Veličković
Christos Perivolaropoulos
Federico Barbero
Razvan Pascanu
Venue: ICML 2025 poster (2025)
Keywords: softmax
size generalisation
attention
transformers
length generalisation
sharpness
entropy
Primary area: deep_learning->theory
TL;DR: We prove that the softmax function cannot robustly model sharp functions with increasing size, based on several controlled experimental observations over simple attention heads, as well as in language models.

Abstract: A key property of reasoning systems is the ability to make sharp decisions on their input data. For contemporary AI systems, a key carrier of sharp behaviour is the softmax function, with its capability to perform differentiable query-key lookups. It is a common belief that the predictive power of networks leveraging softmax arises from "circuits" which sharply perform certain kinds of computations consistently across many diverse inputs. However, for these circuits to be robust, they would need to generalise well to arbitrary valid inputs. In this paper, we dispel this myth: even for tasks as simple as finding the maximum key, any learned circuitry must disperse as the number of items grows at test time. We attribute this to a fundamental limitation of the softmax function to robustly approximate sharp functions with increasing problem size, prove this phenomenon theoretically, and propose adaptive temperature as an ad-hoc technique for improving the sharpness of softmax at inference time.

Softmax is not Enough (for Sharp Size Generalisation)
Petar Veličkovié' Christos Perivolaropoulos'
Federico Barbero 2
Razvan Pascanu'
Abstract
(ia-distribution)
(outtdf-listtibutiia))
A key property of reasoning systems is the abil- ity to make sharp decisions on their input data. For contemporary AI systems, a key carrier of sharp behaviour is the sof tmax function, with its capability to perform differentiable query-key lookups. It is a common belief that the predic- tive power of networks leveraging tmax arises from "circuits" which sharply perform certain kinds of computations consistently across many diverse inputs. However, for these circuits to be robust, they would need to generalise well to ar- bitrary valid inputs. In this paper, we dispel this myth: even for tasks as simple as finding the max- imum key, any learned circuitry must disperse as the number of items grows at test time. We attribute this to a fundamental limitation of the softmax function to robustly approximate sharp functions with increasing problem size, prove this phenomenon theoretically, and propose adaptive temperature as an ad-hoc technique for improving the sharpness of softmax at inference time.
Figure 1. Illustration of Theorem 2.2, one of our key results. As- suming a tokenised input from a fixed vocabulary and non-zero temperature, for every sof tmax attention head inside an architec- ture comprising only MLPs and tmax self-attention layers, must hold that, given sufficiently many tokens, its attention coefficients will disperse, even if they were sharp in-distribution.
of logits, e E R", into a probability distribution, in form that is part of the exponential family. Further, sof tmax allows for application of a temperature parameter, A E R, to adjust the amount of probability mass attached to the highest logit-a concept borrowed from the Boltzmann distribution in statistical mechanics.
Initially, the primary utilisation of softmax in deep learning was within the final layer of classifiers. Its influence in this domain vastly expanded after it saw use in the internal layers-as differentiable key-value store (Graves et al., 2014) or mechanism for attending over the most relevant parts of the input (Bahdanau et al., 2015). This attentional framing of sof tmax was critical in defining important mod- els for sequences (Vaswani et al., 2017, Transformers), im- ages (Dosovitskiy et al., 2021, ViTs) and graphs (Veličković et al., 2018, GATs).
1. Motivation
no understatement to say that the sof tmax, : R"
[0, 1|" function':
exp(e1/0)
exp(en/0) Σn exp(ex/0)
sof
Σn exp(ex/0)
one of the most fundamental functions in contemporary artificial intelligence systems.
Several efforts attribute the success of tmax to its capa- bility of modelling computations relevant to reasoning. This can be related to the concept of circuits in theoretical com- puter science (Arora & Barak, 2009). Several interpretable pieces of "circuitry" (Olah et al., 2020) have already been discovered in large Transformers, primarily under the um- brella of mechanistic interpretability (Elhage et al., 2021; Olsson et al., 2022; Wang et al., 2022).
The role of sof smax in deep learning is convert any vector
'Work performed while the author was at Google DeepMind. 'Google DeepMind 'University of Oxford. Correspondence to: Petar Veličković <petarv @ goog...com>>.
Proceedings of the nddd International Conference on Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Here we study the robustness of such circuitry, especially when going beyond the distribution the models are trained on-a a critical regime for reasoning engines. We find that, in spite of its many successes, sof soft is fundamentally
'Strictly speaking, the proper name for this function should be softargmax. We choose to retain the terminology introduced by Bridle (1989), primarily for reasons of alignment with modern deep learning frameworks.

---

Softmax is not Enough (for Sharp Size Generalisation)
unable to robustly generalise such circuits out of distribution, especially as it provably cannot approximate sharpness with increasing problem size (Figure 1).
The analysis of attentional coefficients and attempting attribute interpretable operations to them dates back to the earliest deployments of internal sof tmax layers at scale; examples include Graves et al. (2014, Figure 6), Bahdanau et al. (2015, Figure 3), Vaswani et al. (2017, Figures 3-5) and Qiu et al. (2018, Figure 5). A strong current topic in this space analyses the self-attentional heads of Transformers (Voita et al., 2019; Jain & Wallace, 2019).
Background
Here we call 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' values?. 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 items).
With the rise of large language models, mechanistic inter- pretability has taken charge in detecting and elucidating various circuits in Transformers (Elhage et al., 2021). Some prominent discoveries include induction heads (Olsson et al., 2022), indirect object identification (Wang et al., 2022), multiple-choice heads (Lieberum et al., 2023), successor heads (Gould al., 2023), attentional sinks (Darcet et al., 2023), comparator heads (Hanna et al., 2024) and retrieval heads (Wu et al., 2024). Most recently, these efforts have relied on sparse autoencoders (Kissane et al., 2024).
Key Theoretical Result We define sharp functions by their behaviour as their number of inputs varies. This di- rectly motivates the size generalisation 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 con- taining no more than input items, and we take partieular interest in their sharpness instances with n' > items; these are considered -ittof-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 (Anil et al., 2022); in graph machine learning, it known as size generalisation (Yehudai et al., 2021).
The skills above are quite impressive and span many rules one might hope a robust reasoning system would have, and the discovered heads always appear sharp when inspected on in-distribution samples. However, it is also known that many easy tasks requiring sharp attention -schhsuchhhh as find- ing minima-and are hard to do reliably with LLMs out-of- distribution, particularly on larger inputs (Markeeva et al., 2024, Figure 6). More challenging sharp order statistic tasks, such as finding the second minimum (Ong & Veličković, 2022) may even be hard to learn in inddistribution. The discrep- ancy of such results with the previous paragraph motivate our study, and formalisation of sof dispersion.
Through one of our key theoretical results (Theorem 2.2), we demonstrate that modern deep learning architectures, operating over fixed vocabulary of input tokens and lever- aging the softmax function, are fundamentally incapable of learning functions that remain sharp under such larger instances. This is due the fact that the coefficients emitted by the sof function must disperse 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.I and Remark B.2).
Certain dispersion effects in sof toftmaxx-..2 as effect of increasing temperature- are already well-understood in thermodynamics. A core contribution of our work is under- standing dispersion a setting where the amount of logits can vary, which is relevant for generalisation in Trans- formers. We are not the first to observe dispersion in this setting empirically; prior works studying the capability of Transformers to execute algorithms (Yan et al., 2020) and perform random- lookups (Ebrahimi et al., 2024) also note dispersion patterns. Our work is the first to rigorously prove these effects, directly attribute them to the sof tmax operator, as well as propose ways to improve sharpness empirically within softmax. The proof technique we will use to demonstrate this is inspired by Barbero et al. (2024), though unlike their work, our key results apply regardless of whether the computational graph is bottlenecked or not.
We hope that our results will encourage future study of alternative attentional functions, in light of the problems we identify, especially for building reasoning engines of the future. That being said, we also believe our findings indicate ways to modify the smax function to support sharpness for longer as one simple example, we propose an adaptive temperature mechanism for smft
"In practice, as softmax coefficients that are exactly zero or one are often not casily obtainable, we may consider an item to be expressed in the final answer if its coefficient is greater than some e 0 (which may depend on the problem size). That said, the exact choice of does not affect the essence of our results, as we are able to hound the coefficients both above and below.
Primer on Attentional Heads and Transformers Within this paper we will primarily study the use of sof within self-attentioaal neural network architectures, such as Trans- formers (Vaswani et al., 2017). The core building block of such models is the (dot-product) attentional head, which op-
2

---

RKKK V xk are learnable parameter matrices. Then, dot-prue lot-products between the query and all of the key vectors are taken to compute unnormalised attentional coef- ficients of each item, also known as logits, c(")) R. These coefficients are normalised using the sof smft function to Softmax is not Enough (for Sharp Size Generalisation)
erates collection of nodes (or tokens), with features (n) R for node f i N n, for a given query vector
(n)
R
First, the attentional head computes key, (updated) query and value vectors via matrix multiplication:
Figure 2. Visualising the attentional head for the max retrieval task for batch of 32 randomly-sample input sets (each represented by one of the rows), over the 16 items with largest key (columns). If the head operates correctly, must allocate sharp attention the rightmost item. From left to right, in each frame we double the number of items the head has to process (starting from 16 items).
[!m) = Kx (1))
(n) = Qq (n)
(m)
Vx; (1)
(2)
where K, Q V
Entropy in Gemma 26 attention heads
5
obtain attentional coefficients, a; (n) E R. Finally, the at- tentional coefficients are used for a weighted sum of value
vectors, which represents the output of the attentional head, ((m) E R
Shannon Entropy
(n) = (q (1))
k!n) (n)
(n) = sof tmax (e( (3)
yo(m)
Σ
oi (n) (n) ((
1
ISiSn
With regard how attentional heads are used within Trans- formers, will mainly analyse two of the most popular strategies: BERT-style (Devlin et al., 2019) and GPT-style (Radford et al., 2018). In both cases, each of the input nodes computes its own attentional output, i.c. there is one query vector per node, computed as q Qx("), leading to per-node attention coefficients oij and outputs y (m) by dis-
20
21
22
2
24
25
2
27
29
2
Number of items
Figure 3. Entropy of attention heads in the first block of Gemma 2B with prompt "What is the maximum in the following sequence: {seq}? The maximum is:" " and varying the number of elements in seq. Fach curve is one attentional head; the blue shaded curve the mean and standard deviation across of them.
tributing Equation across queries. The main difference is in the choice of keys.
Our key dispersion results hold for both styles of attention- this mainly due to the fact that all predictions made by GPT-style architectures are dependent on the final token embedding, (x), which will attend over all items, much like any BERT head. The main difference hetween the two will be in qualitative effects on certain corollaries of the theory (Appendices B-C).
In BERT-style self-attention, each node's query vector at- tends over all of the key vectors, i.e. it is obtained by directly distributing Equation 3 across all queries:
(n) = (aa
T kmm
(n) Oij
sof tmax (e, (4)
Σ
(1)) (n)
2. Dispersion in soft and Transformers
15/50
In comparison, GPT-style attention (also known "causal masking" or the decoder-only Transformer) only allows information to flow forwards; each node's query vector may only attend to the key vectors from nodes that precede it. This yields the following modification:
To motivate our theory, we train simple architecture in- cluding a single attentional head to predict a feature of the maximum item in a set. Each item's features are processed with a deep MLP before attending, and the output vector of the attention is passed to deep MLP predictor (see Ap- pendix A for experimental details). We train this model using sets of < 16 items, and in Figure 2 we visualise the head's ''s attentional coefficients, computed over sets of vary- ing size at inference time.
(n)
q:
(n)
k(n1
i
(n)
= softmax;(e)) (1))
aij
00
j i>i
(5)
y(m))
Σ
v
(n) (n)
While the model indeed attributes focus sharply and cleanly on the maximum item, this only holds true on the problem
ij
1515i
3

---

Sof tmax is not Enough (for Sharp Size Generalisation)
sizes that the model was trained on. As simulate a gener- alisation setting where the problem size increases (without changing the value distribution), the attentional coefficients eventually disperse towards the uniform distribution.
Let e())) = (q k 5n) and K(a) = x(x) (n) ....,,nn) x,))), where : x" - R and K X" + Ruxk are continuous functions, each ex- pressible as a composition of L layers 9L fL ........ o 9i S1 where each layer contains a feedforward component fz kn)s fi(zx) or self-atientional component I (z),,...... Zn)k = -isg (11,;(2)) where Ouk E [0, 1] are sof tmax-normalised attention coefficients and or is a feedforward network. Then, for any 0 > O and e > 0, there must exist an n E N such that sof tmaxo(c < c for all 1 S n.
(n) where q (nn) = o(x
(n)
Xn (*)
This effects manifests in the attention heads of Transformers as well-we visualise the entropy (a proxy for sharpness) of Gemma 2B (Gemma Team et al., 2024)'s 's heads when answering a similar maximisation task in Figure
In fact, we can show that this effect is inevitable in softmax using the following Lemma:
Lemma 2.1 (moftmax must disperse). Let e!n) E Re be a collection of n logits going into the sof function with temperature 0 > 0, bounded above and below s.l. (x) S M for some choice of constants m, M E I. Then, as more items are added (n + +00), it must hold that. for each item 1 S k S n, softmax, (e())) = (.))
That is, attention coefficients must disperse in all global Transformer heads if the input vocabulary is finite.
m S 6k
Proof. Firstly, note that since X' is finite set of m- dimensional vectors, then it is also part of a compact space spanning all convex combinations of those vectors. Then, all feedforward layers, fü and 2i being continuous functions, move inputs from a compact set to another compact set. Similarly, every self-attentional layer, Gi. computes con- vex combination of the outputs of e; and as such, if outputs of
The computed attention coefficients hence must disperse for all items with increasing problem size.
Proof. Let us denote the attentional coefficient assigned to [0, 1]. Then can bound
(n) sof (e E
A by d&
compact space, the outputs of gi remain the same compact space. Therefore, if the input space of & and x is compact, then the output space of & and (each row of) x on R must be compact as well, regardless of the choice of n. Further, the dot product of two vectors (q(")) Tk(")) com- ing from compact spaces must be compact as well. Hence, the logits must be bounded by m S 6
(n) Okk above as:
exp(el") /0) /0) Σ, exp(((()) 10)
1 exp(M/0) 1 n emp(mm n
M m
exp
0
Similarly, we can bound of below as:
(n) < M for con-
stant m and M. Then, letting ò = M m. we know (Lemma 2.1) that sof boftmaxx( (c(")) ) S 1 exp (6/0), so for all n exp (5/0)
exple("//0 10)
2 exp(m//) 1
m M
exp
n
E, exple(("))11 18) n ex)(MM/0) n
0
this value will be below r.
Hence, if we let = (M- - m)
It might seem intuitive that attention head dispersion is a potentially destructive event, which forces the Transformer into misclassifying certain inputs. We prove this intuition in Appendix B. We also discuss the rate at which dispersion occurs at various model depths in Appendix C.
1 8 S a
(n)
1
8
exp-
exp
(8)
n 0
n
Which implies a
(n) ( as ō and 0 are both constants.
Practical Values of 5 It evident that the logit spread, 5, is the key controller of dispersion in our results. If we were to naively apply our theory (to the maximal and minimal values representable by the types used to represent floating point numbers such as bfloat16) 16,, it would result in fairly loose bounds which are not particularly informative.
Lemma 2.1 relies bounding logits. The difference of these bounds (the spread, 8 = max, e,") min, (n)) con- trols the rate of dispersion. In modern Transformer LLM architectures operating over a vocabulary of possible token values, we can actually bound the logits in every single attentional layer-implying that dispersion must happen everywhere in a Transformer for sufficient problem sizes:
However, it can easily be observed that empirical values of 8 do not reach those theoretical limits. For example, when feeding Gemma 2B and 7B models with an entire code sample of Gemma's Modules file* (~ 4, 000 tokens), the
Theorem 2.2 (softmax in Transformers over vocabular- ies must disperse). Let X C. IMM be a set of possible m- dimensional input features, and let (()) E X² be a matrix of input features for n items. Further, assume that input fea- tures come from a finite set of possible values, ie. |2| < [N].
httpps///iihhub com/ www/google-eeoooooooooooooooooee groogle-deepmind/gemeeee main/ heaan//g//mmmduuess.p..
4

---

tmax is not Enough (for Sharp Size Generalisation)
empirically observed values of ó across all of their attention
This motivates our use of adaptive temperature, where we vary l depending on the entropy in the input coefficients. Adaptive temperature can be elegantly motivated by the fact that decreasing the temperature must monotonically decrease the entropy, which is well-known in thermodynam- ics:
heads are substantially more contained:
Gemma 2B: 8 E [2.28, 14.78] (average 5.69 + 2.05);
Gemma 7B: 8 E [0.09, 32.74] (average 5.82 + 2.61).
Proposition 3.2 (Decreasing temperature decreases en- tropy). Lei e E Rn be a collection of n. logits. Con- sider the Boltzmann distribution over these n. items, ps x exp(-3e.)) ()) for B E R. and let H = -; Pi log be its Shannon entropy. Then, as 3's magnitude increases, H must monotonically decrease.
There are several possible reasons why practical pre-trained models keep logit spreads contained; for example, any label noise prevents the model from learning arbitrarily large parameter values, or the derivative needed to make off-target logits arbitrarily negative may vanish when the softmax output is already sufficiently sharp.
Thus, since 33 x z where @ is the temperature in sof tmaxp, decreasing the temperature must monotonically decrease the entropy. We provide a full proof in Appendix E. To supplement Proposition 3.2 empirically, we also provide- in Figure 5-a visualisation of how the Shannon entropy varies with temperature, for 10-logit input with varying spread between the logits.
3. A daaptive Temperature
Since we now know dispersion is inevitable, are there any ways we can leverage our theory's 's findings to make sof tmax sharper? One obvious constraint our theory rests on is the assumption that 0 > 0, i.e. that our temperature is nonzero. While zero temperature-albe known as hard attention (Denil et al., 2012; Ranzato, 2014; Mnih et al., 2014; Xu et al., guarantees sharpness, training large- scale Transformers with tends to not work well in practice (Bica et al., 2024).
To compute the approximate temperature value as a function of entropy, we generate a dataset of inputs to our model where the maximal items do not obtain the highest logit (see Appendix for detailed overview). For each such input, we find the "optimal" value of 0 that would maximise its probability. Then we fit an inverse degree-4 polynomial to this data Figure 6-and use it to predict temperatures to use at inference time. Note we do not wish to increase entropy; as such, we do not correct O to values greater than 1.
What about applying 0 = 0 to an already-trained Trans- former? We can show this is also problematic since, for any attention head where the Transformer has learnt to induce sharpness, it necessarily did so by increasing magnitude of its weights (see Appendix D for a proof and numerical validation):
The JAX (Bradbury et al., 2018) implementation of our adaptive-0 omax is provided in Figure 4, and we use it as drop-in replacement for jax, nn. smft in all of our experiments.
Proposition 3.1 (Sharpness in Transformers necessitates (m)
large weights). Let ( E R" be collection of n log-
its, computed using dot product attention; i.e. 6g
QQ, Kxt), where y E R is a query vector and Q K E Rmm xm are parameters. Let 8 max
(n) min C (n)
While this approach requires two calls to jax. nn. tmax in place of one, as well as computing several additional intermediate tensors, we are able to implement it in a way that allows the entropy correction computation to be fully streamed, and hence compatible with efficient, scalable ap- proaches like Flash Attention (Dao al., 2022) that uses ( rather than O11²) memory to compute attention. We provide the derivation of our streamed algorithm in Ap- pendix G.
Isisn
1sisn"
be their maximum difference. Then ô is upper bounded as (K) 8 < 20 2o max |y max ||x:]], where o mmax, dmax E
(Q) (K)
IR are the largest singular values of Q and K.
That is, the sharpness of the softmax in Transformers de- pends on the norm of its parameters. Note that there is a common practice of leveraging operators such as layer nor- malisation (Ba et al., 2016) extensively within Transformer architectures, which clamps ||x:|| || and ||y|| if applied right before the query-key mechanism, accentuating the impact of Q and K's singular values.
Note we are not the first to propose dynamically adapting temperature- Neumaan et al. (2018); Radford al. (2021) do this in the classification layer (and hence do not have to handle ever-increasing amount of items), whereas Chiang & Cholak (2022); Cao et al. (2024) perform it over interme- diate attentional heads, but in a way that only depends on problem size (e.g. multiplying logits by log n), hence not taking into account initial logit sharpness. It is important to also call out AERO (Jha & Reagen, 2024), a method
However, forcing large parameters promotes overfitting, and the likelihood that the incorrect item gets the largest logit- see Figure 2. Setting temperature to zero will then degrade accuracy- we might prefer to make the coefficients sharper while making sure that the chosen item not left behind.
5

---

Softmax is not Enough (for Sharp Size Generalisation)
def hennrdiveteemprrture___ssss (logitss: original_probs jax. sof taa.fttaaxllliooaaaa
poly_fit = jnp. array([-O. 037, 0.481, -2.3, 4.917, -1. -917) # see Figure 6 1e-9),
entropy jnp. wunp.summ-orriii caum-original_praa * jnp. log.orrggiall compute the Shannon entropy
axis=-1, keepdims-True)
beta jnp. where( # beta theta
entropy 0.5, don't overcorrect low-entropy heads jnp. www.mmxximmipppppp polyval entropy), 1.0), # never increase entropy 1.0)
return jax. nn. sof tmax (logits * beta)
Figure 4. Our implementation adaptive temperature JAX.
which introduces learnable temperature, and Entropix (xjdr & doomslide, 2024), a notable library for (var)entropy- LLM sampling.
4. Experimental Results
To validate the utility of our proposed adaptive temperature scheme, we evaluate it on both our previously-mentioned max retrieval task-which allows us pristine environment for evaluating whether adaptive temperature leads to more useful attention heads well the CLRS-Ter algorith- mic reasoning benchmark (Markeeva et al., 2024), which represents challenging reasoning task for decoder-only Transformers, and is hence likely to require low-entropy behaviour.
Figure 5. Entropy of softmax, (a) for 10 elements of a power series a; Ai, split into four regions depending on range of (A, 0). Degenerate cases: near À and À 1 (all logits equal).
4.1. Max Retrieval
For this task, we first train our single attention head archi- tecture as described in Appendix A; then, we evaluate it at various numbers of input items, with and without apply- ing adaptive temperature its sole sof tmax function call. Note that this is a "pure" inference time adjustment- no modifications the learned parameters performed!
Polynomial optimal factor entropy
Observed data Polynomial
3.5
Optimal temperature
The results-averagee over ten seeds and with statistical significance tests applied are summarised in Table 1. As is evident, applying adaptive temperature leads to a more performant retrieval head on larger inputs, with statistical significance ascertained via a paired 1-test.
2.0
1.5
1.0
0.5
These results are further supplemented by a qualitative com- parison of the sof tmax coefficients before and after apply- ing the temperature adaptation. As can be seen in Figure 7, our proposed adaptive temperature adaptation indeed leads to sharper coeflicients on larger inputs and higher attention being directed to the desired item, even situations where it did not receive the largest logit.
0.0
Shannon entropy
Figure 6. The polynomial fit used derive our adaptive formula for as a function the Shannon entropy, H. The fit degree-4 function was ze 1/(-1.791 + 4.917H - -.3 0.481 H 0.037 We do not apply the correction to 0 predicted greater than 1.
We have now successfully validated the predictions of
6

---

Sof tmax is not Enough (for Sharp Size Generalisation)
Table 1. Improvements observed when applying adaptive temperature on the max retrieval task (without changing the parameters), averaged over ten seeds, p-vuuess computed using a paired t-test.
ID size
Out-of-distribution sizes
Model
16
32
64
128
256
512
1. 024 2. 048 4, 096 8, 192
16, 384
Baseline
98. 6% 97.1%
94.3%
89.7%
81.3%
70.1%
53.8%
35.7%
22.6% 15.7%
12.4%
Adaptive 0 98. 6% 17 94.5%
89.9%
82.1%
72.5%
57.7% 39.4%
24.9% 17.5%
14.0%
p-value
0.4
0.4
0.002
2- 10 2- 10
43
10-5
10 6- 10 0.02
10
4 10
tional heads" behaviours, and even less so to ascertain that the model's 's robustness depends on those specific heads in those ways. As briefly discussed before, any such individual endeavour typically leads to a brand-new research project in mechanistic interpretability, and we do not find this be in-scope of our 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 polyno- mial fit as in Figure 6); 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, replace all in- stances of jax. nn. softmax within all the attentional heads of Gemma 2B with our adapti ve. wemperature sof tmax 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 tem- perature adaptation.
Figure 7. Visualising the attentional head for the max retrieval task with (below) and without (above) adaptive lemperature, for the same batch and model as in Figure 2. Note the increased sharpness in the coefficients, especially as the amount of items increases.
theory in a controlled environment. What about a more challenging benchmark with a baseline model comprising many attentional heads?
4.2. CLRS-Text
These final comparative results may be found in Figure 8, and they demonstrate a significant advantage of the adaptive temperature-bacee backed model on nearly all of the thirty algo- rithms under study. While we cannot make strong claims about where the underperformance on outlying algorithms comes from, a unifying property of some of them (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 our adaptive temperature was derived; this might cause unintended distribution shifts.
In this benchmark, we follow the protocol established by Markeeva et al. (2024) and fine-tune Gemma 2B models (Gemma Team al., 2024) on the thirty algorithmic execu- tion tasks in CLRS-Text, plotting their performance profiles in- and out -ii-o---dittibttionn. at various problem sizes.
While it may be tempting to directly re-apply our learned adaptive temperature function from Figure 6 solely at in- ference time- the same way we did in the max retrieval experiments-this approach does not empirically work well in the CLRS-Text regime. This is due to the fact that CLRS- Text inputs are often textual representations of floating-poinl numbers and therefore individual numbers often span multi- ple tokens. It is therefore insufficient and inappropriate to aim for entropy levels where all the focus would be on one token only, as was desirable the max retrieval task.
Our result indicates that, even in a complex system with many interactions between attentional heads, it is possible to extract benefits from the simple idea of dynamically adapt- ing the temperature-annt we hope our result paves the way for more involved future investigation of such approaches.
One follow-up on this could be to perform exactly the same polynomial fit exercise leading up to Figure 6, only this time focussing on "optimal" values of temperature for Gemma's ss attentional heads. However, in this regime, we argue this ex- ercise is substantially less trivial to do-as as we are now deal- ing with a system spanning many attentional heads across many layers, it is not easy to even discover relevant atten-
5. Conclusions
"Energy continuously flows from being concentrated To becoming dispersed, spread out, wasted and useless." The 2nd Law: Unsustainable, by Muse
7

---

tmax is not Enough (for Sharp Size Generalisation)
vity selector
articutatioo points
beliman ford
815
Dinary search
bridges
10-
0.6
0.2
15
30
12
14
15
20
25
20
32
18
24
30
16 24
32
26 kodane
Royd warshall
0.-
3.44-
graham SCan
heapsort
Insertion sort
jerats march
kmp matcher
T05 length
1.0-
0.6-
0.0
24
10
mst prfm
str ng
optimal
0-2
32 Quicisort
euickseect
segments intersect
strongly connected components
1O0K Schoduling
topological sort
200-
3.44
0.0
10 15
55 64
15
301 longl
102. kr
1030 kangh
100 longi1
0431 No
Figure 8. Resampling test results on CLRS- Text of variants of Gemma 2B, fine-tuned with and without adaptive temperature applied, various problem sizes. Each point on the axis corresponds a particular problem size in the corresponding algorithmic task. For example, on sorting tasks, this corresponds to the number of items being sorted; for graph tasks, it corresponds to the number of nodes in the graph. The blue curves represent the accuracy the baseline fine-tuned Gemma 2B model, whereas the red curves represent the accuracy of that same model, fine-tuned with adaptive temperature, Both Gemma 2B variants were explicitly trained on CLRS-Te tasks-the training set sizes are denoted by dots-and and evaluated zero-shot. Note that we limit our sample length to 2, 048 tokens. and only show performance metrics for sizes where the answer fits in this constraint.
8

---

Softmax is not Enough (for Sharp Size Generalisation)
In this paper, we have provided extensive theoretical and empirical evidence that the softmax-1 key function in the design of modern frontier architectures--ss leveraged within those architectures in way that is fundamentally unable to sustain robust reasoning behaviours across all possible inputs. This is due to the fact its output coefficients are nec- essarily dispersing provided sufficient input elements. Our evidence relies on fundamental properties of the sof tmax, as well specific popular design decisions as tokenisation, global attention, and similar.
problem size the head saw at training time, not the largest context deployed training time.
Lastly, our key Theorem relies on the model being built out of continuous building blocks. Inserting dis- continuities in the feedforward payhas using approaches like Dudzik et al. (2024) as inspiration would also break the assumptions of our theory, though it comes with obvious challenges to learning at scale.
While such approaches haven't 't seen as much success at scale as the "vanilla" Transformer, we hope our results inspire future work into making them stable, especially for constructing reasoning systems.
Beyond illustrating and proving these dispersion effects, we also attempted to use our theoretical framework to propose an adaptive temperature approach that is able- at least to a certain extent-be to hold the dispersion effect at bay. It is our opinion that the favourable results we observe with adaptive temperature warrant further investigation, and indicate that such adaptive layers are strategy worth dedicating further attention to in future work.
Acknowledgements
We would like to deeply thank Daniel Johnson, the author of Penzai (Johnson, 2024)-without it, our exploratory anal- yses would not be nearly as fruitful. Further, we thank Alex Matthews, Andrew Dudzik and João Araújo for helping us with the proof of one of our propositions, and Arthur Conmy, Neel Nanda and Csaba Szepesvári for reviewing the paper prior to submission and numerous highly useful comments and references. Lastly, we show deep appreciation to Alex Vitvitskyi. Olga Kozlova and Larisa Markeeva, for their endless engineering support with CLRS-Text.
We conclude by remarking, once again, that adaptive temper- ature is merely an ad-hoc method and it does not escape the conclusions of our theory! The key takeaway of our paper not the adaptive temperature proposal; it is the fact that we find it worthwhile to more seriously invest in research of hybrid architectures that will not fully rely on the tmax function, at least within the confines of the assumptions of our theory. To name a few possibilities:
Any kind of unnormalised attention, such as linear (Schmidhuber, 1992), sigmoidal (Ramapuram et al., 2024) or stick-breaking attention (Tan et al., 2025) does not have the dispersion issues presented here. That heing said, becomes substantially harder to meaning- fully rank items using them, see e.g. the GATv2 paper (Brody et al., 2022).
Impact Statement
This paper presents theoretical and empirical study into the numerical properties of the soft function, placing into the spotlight several issues that models leveraging will experience when generalising to longer inputs at test time, This is setting that has clear implications to research in long context generalisation, which an important frontier for reasoning in Al systems. Hence, any societal conse- quences of our work would potentially be contained within the consequences for general-purpose purpose research in either of those established areas.
Explicitly removing excess attention could be another attractive alternative, as robustly done in selective at- tention (Leviathan et al., 2025) in way that is capable of alleviating dispersion. The Differential Transformer (Ye et al., 2025) follows similar idea, though the ex- act way it is realised - subtraction of two distributions after the application of sof toax-ax- - may not effectively evade dispersion, assuming it already happened in the individual sof tmax calls.
References
Anil, C., Wu, Y., Andreassen, A., Lewkowycz, A., Misra, V., Ramasesh, V., Slone, A., Gur-Ari, G., Dyer, E., and Neyshabur, B. Exploring length generalization in large language models. Advances in Neural Information Pro- cessing Systems, 35:38546-38556 2022.
Similarly, forcing the attention to be hard or local (Martins & Astudillo, 2016; Correia et al., 2019; Peters et al., 2019; Vitvitskyi al., 2025) would also escape the confines of our theory. We already briefly discussed the challenges of learning with hard attention- local attention provides a very interesting alternative, but it behaviours
Arora, S. and Barak, B. Computational complexity: a mod- ern approach. Cambridge University Press, 2009.
Axler, S. Linear algebra done right. Springer, 2015.
must be stressed that for certain heads may appear even at highly "local" scales; OOD here refers to going outside the largest
Ba, J. L., Kiros, J. R., and Hinton, G. E. Layer normalization. arXiv preprint /607 06450, 2016.
9

---

tmft is not Enough (for Sharp Size Generalisation)
Bahdanau, D., Cho, K., and Bengio, Y. Neural machine translation by jointly learning to align and translate. In wwwpg////////////// 0473..
Darcet, T., Oquab, M., Mairal, J., and Bojanowski, P. Vision transformers need registers. arXiv preprint arXiv:2309 2023.
ICLR, 2015. URL
Barbero, F., Banino, A., Kapturowski, S., Kumaran, D., Madeira Araújo, J., Vitvitskyi, O., Pascanu, R., and Veličković, P. Transformers need glasses! information over-squashing in language tasks. Advances in Neural Information Processing Systems, 37:9811 111111199142 2024.
Denil, M., Bazzani, L.,, Larochelle, H., and de Freitas, N. Learning where attend with deep architectures for im- age tracking. Neural computation, 24(8):2151-2184, 2012.
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. BERT: Pre-training of deep bidirectional transformers for lan- guage understanding. In Burstein, J., Doran, C., and Solorio. T. (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Com- putational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pp. 4171 17864 Min- neapolis, Minnesota, June 2019. Association for Compu- tational Linguistics. doi: 10. 10.18653/v1/N19-11
Bica, I., Ilie, A., Bauer, M., Erdogan, G., Bošnjak, M., Ka- planis, C., Gritsenko, A. A., Minderer, M., Blundell, C.,, Pascanu, R., and Mitrovic, J. Improving fine-grained understanding image-text pre-training. In Forty-first In- ternational Conference on Machine Learning, 2024. URL https://openreviae htwpiiiww...eeeeeeee mett/forum?id=5nxIRR08.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Millee S., and Zhang. Q. JAX: composable transformations of Nython+NumPy programs, 2018. URL cww/ googl///oo http://gi wtw..ttoooooooooo.oooooo
URL
https:
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x16 words: Transformers for image recognition at scale. In International Confer- ence on Learning Representatiins, 2021. URL https: //operreiiew wwtp:iiwwwweeeeeeeeeeeeeeeeeeeee
Bridle, J. Training stochastic model recognition algorithms
as networks can lead to maximum mutual information estimation of parameters. Advances in neural information processing systems, 2, 1989.
Brody, S., Alon, U., and Yahav, E. How attentive are graph attention networks? In International Confer- ence on Learning Representations, 2022. URL https: wwwpeeeewwwwwneeeeereeeeeeeeeeee het.//orrmiid-f2xxims./..
Dudzik, A. J., von Glehn, T., Pascanu, R., and Veličković, P. Asynchronous algorithmic aligmment with cocycles. In Learning on Graphs Conference, pp. 3-1. PMLR, 2024.
//openreview.r net
Ebrahimi, M., Panchal, S., and Memisevic, R. Your context is not an array: Unveiling random access limitations in transformers. arXiv preprint arXiv:2408.05500 2024.
Cao, C., Ren, X., and Fu, Y. MVSForme+++: Revealing the devil in transformer's s details for multi-view stereo. In The Twelfth International Conference Learning Represen- tations, 2024. URL https: http:///pppevviiieeee wwtp://ee/ooouumme id=wXWfvSpYth.
Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., et al. A mathematical framework for transformer circuits. Transformer Circuity Thread, 1(1):12, 2021.
Chiang, D. and Cholak, P. Overcoming theoretical limita- tion of self- attentonnn In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 7654-7664, 2022.
Gemma Team, Mesnard, T., Hardin, C., Dadashi, R., Bhu- patiraju, S., Pathak, S., Sifre, I.., Rivière, M., Kale, M. S., Love, J., et al. Gemma: Open models based on gemini re- search and technology. arXiv preprint 2Xx::2403300299 2024.
Correia, G. M., Niculae, V., and Martins, A. F. T. Adap- tively sparse transformers. In Inui, K., Jiang, J., Ng. V., and Wan, X. (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural pp. 2174-2184, Hong Kong, China, November 2019. Association for Computational Linguistics, doi: 10. 10.11663/11//11-12223 URL htttp:::/allantoolog..org///////
Gould, R., Ong, E., Ogden, G., and Conmy, A. Successor heads: Recurring, interpretable attention heads in the wild. arXiv preprint arXiv:2312. 2023.
Language Processing (EMNLP-IJCN)
Graves, A., Wayne, G., and Danihelka, I. Neural turing machines. arXiv preprint arXiv:1410.540. 2014.
Dao, T., Fu, D., Ermon, S., Rudra, A., and Ré, C. Flashat- tention: Fast and memory -fffrr-eeffccient exact attention with io-awareness. Advances in Neural Information Process- ing Systems, 35:16344-41 1144-16359, 2022.
Hanna, M., Liu, O., and Variengien, A. How does gpt-2 compute greater-than?: Interpreting mathematical abili- ties in a pre-trained language model. Advances in Neural Information Processing Systems, 36, 2024.
10

---

tmax is not Enough (for Sharp Size Generalisation)
He, K., Zhang, X., Ren, S., and Sun, Deep residual learn- ing for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
Neumann, L., Zisserman, A., and Vedaldi, A. Relaxed softmax: Efficient confidence auto-calibration for safe pedestrian detection. 2018.
Olah, C., Cammarata, N., Schubert, L., Goh, G., Petrov, M., and Carter, S. 7oom in; An introduction to circuits. Distill, 2020, doi: 10.239 10tpistiii..002.00001.. URL
Heek, J., Levskaya, A., Oliver, A., Ritter, M., Rondepierre, B., Steiner, A., and van Zee, M. Flax: A neural network library and ecosystem for JAX, 2024. URL http:// github. com/google/flin
https://distin: htt.ddisilll.ubbbb//em wwb/2020/cireu tt/ zoom-in.
Olsson, C., Elhage, N., Nanda, N., N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., et al. In-context learning and induction heads. arXiv preprint arXiv:2209.1189: 2022.
Hendryeks, D. and Gimpel, K. Gaussian error linear units (gelus). arXiv preprint 1rXiv:160.008441 2016.
Jain, S. and Wallace, B. C. Attention is not explanation. arXiv preprint arXiv: / 19002.90219900 10186, 2019.
Ong. E. and Veličković, P. Learnable commutative monoids for graph neural networks. In Learning on Graphs Con- ference, pp. 43-1. PMLR, 2022.
Jha, N. K. and Reagen, B. Aero: Softmax-only ilms for effi- cient private inference. arXiv preprint arXiv:2410.13060 2024.
Peters, B., Niculae, V., and Martins, A. F. T. Sparse -qquence-t--sqqeence models. In Korhonen, A., Traum, D., and Mårquez, L. (eds.), Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 1504-1519, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10. 18653/v URL attps://aclaanthology org/P19-1146.
Johnson, D. D. Penzai + Treescope: A toolkit for interpret- ing, visualizing, and editing models as data. ICMI. 2024 Workshop on Mechanistic Interpretability. 2024.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In ICLR (Poster), 2015. URL http:// arxiv.org/abe hwrpg/abs/1412.. 01.11412.6999 6980.
Kissane, C., Krzyzanowski, R., Bloom, J. L, Conmy, A., and Nanda, N. Interpreting attention layer outputs with sparse autoencoders. arXiv preprint arXiv:2406. 17759, 2024.
Qiu, J., Tang, J., Ma, H., Dong, Y., Wang, K., and Tang, J. Deepinf: Social influence prediction with deep learning, In Proceedings of the 24th ACM SIGKDD international conference knowledge discovery & data mining, pp. 2110-2119 2018
Leviathan, Y., Kalman, M., and Matias, Y. Selective at- tention improves transformer. In The Thirteenth In- ternational Conference on Learning Representations, 2025. URL https: wtw.lppereiie..ne//fooooeee v0FzmPCd1e.
Radford, A., Narasimhan, K., Salimans, T., and Sutskever, I. Improving language understanding by generative pre- training, 2018. URL https: // www.m.ikeaaptain.. comp resources/pdf https//dff/PPTPPPPP /
Lieberum, T., Rahtz, M., Kramár, J., Nanda, N., Irving, G., Shah, R., and Mikulik, V. Does circuit analysis inter- pretability scale? evidence from multiple choice capa- bilities in chinchilla. arXiv preprint arXiv:2307.0945 009488, 2023.
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748-8763. PMLR, 2021.
Markeeva, L., McLeish, S., Ibarz, B., Bounsi, W., Ko- zlova, O., Vitvitskyi, A., Blundell, C., Goldstein, T., Schwarzschild, A., and Veličković, P. The CLRS-Text Algorithmic Reasoning Language Benchmark. arXiv preprint XrXiv:2406.04229 2024.
Ramapuram, J., Danieli, F., Dhekane, E., Weers, F., Bus- bridge, D., Ablin. P., Likhomanenko, T., Digani, J., Gu, 7.,, Shidani, A., and Webb, R. Theory, analysis, and best practices for sigmoid self-attention. arXiv preprint arXiv:2409.0443) 2024.
Martins, A. and Astudillo, R. From softmax to sparsemax: A sparse model of attention and multi-label classification. In International conference on machine learning, pp. 1614 1623. PMLR, 2016.
Ranzato, M. On learning where to look. arXiv preprint arXiv:1405.5485 2014.
Mnih, V., Heess, N., Graves, A., et al. Recurrent mod- els of visual attention. Advances in neural information processing systems, 27, 2014.
Schmidhuber, J. Learning to control fast-weight memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1):331-139 1992.

---

Softmax is not Enough (for Sharp Size Generalisation)
Tan, S., Yang, S., Courville, A., Panda, R., and Shen, Y. Scaling stick-breaking attention: An efficient implemen- tation and in-depth study, In The Thirteenth International Conference on Learning Representations, 2025. URL https://openrevia net/ seet/oorum2id=r81105544
Ye, T., Dong, L., Xia, Y., Sun, Y., Zhu, Y., Huang, G., and Wei, F. Differential transformer. In The Thir- teenth International Conference Learning Represen- tations, 2025. URL Ca5DvoCm1gGhh..
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, t.,, and Polosukhin, I. At- tention is all you need. Advances in neural information processing systems, 30, 2017.
Yehudai, G., Fetaya, E., Meirom, E., Chechik, G., and Maron, H. From local structures to size generalization in graph neural networks. In International Conference on Machine Learning, pp. 11975-11986 PMLR, 2021.
Veličković, P., Cucurull, G., Casanova, A., Romero, A., Liò, P., and Bengio, Y. Graph attention networks. In International Conference on Learning Representa- net/foruum
Zaheer, M., Kottur, S., Ravanbakhsh, S., Poezos, B., Salakhutdinov, R. R., and Smola, A. J. Deep sets. Ad- vances in neural information processing systems, 30, 2017.
tions, 2018. URL JdppIcC
Vitvitskyi, A., Araújo, J. G., Lackenby, M., and Veličković, P. What makes a good feedforward computational graph? arXiv preprint arXiv:2502.06755 2025.
Voita, E., Talbot, D., Moiseev, F., Sennrich, R., and Titov, I. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can pruned. arXiv preprint arXiv: 1905.0941.. 2019.
Wang, K., Variengien, A., Conmy, A., Shlegeris, B., and Steinhardt, J. Interpretability in the Wild: Circuit for Indirect Object Identification in GPT-2 small. arXiv preprint urXiv:2211.0059)) 2022.
Wu, W., Wang. Y., Xiao, G., Peng, H., and Fu, Y. Re- trieval head mechanistically explains long-context factu- ality, arXiv preprint arXiv:2404.15574 2024.
xjdr and doomslide. Entropix: Entropy Based Sampling and Parallel CoT Decoding, 2024. URL https:///ithh//
com/ xjdr-alt/eetrrpii....
Xu, K., Ba, Kiros, R., Cho, K., Courville, A., Salakhudi- nov, R., Zemel, R., and Bengio, Show, attend and tell: Neural image caption generation with visual attention. In Bach, F. and Blei, D. (eds.), Proceedings of the 32nd In- ternational Conference
Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 2048- 2057, Lille, France, 07-09 Jul 2015. PMLR. URL https: http::///proceedings. mlr. wwt.r//wrssseeeeee /xt/////xx1..../ 1. html.
Xu, K., Li, J., Zhang, M., Du S. S., ichi Kawarabayashi, K., and Jegelka, S. What can neural networks reason about? In International Conference on Learning Rep- resentations, 2020. URL. https://openrrevew..t ForumRid=-JxbJJeHFPP.
Yan, Y., Swersky, K., Koutra, D., Ranganathan, P., and Hashemi, M. Neural execution engines: Learning to execute subroutines. Advances in Neural Information Processing Systems, 33:17298-17300 2020.
12

---

tmax is not Enough (for Sharp Size Generalisation)
A. Experimental Details for the Maximum Entry Retrieval Task
As briefly described in the main paper, we leverage the max retrieval task over a single attention head as way to empirically validate our theory, as well as assess the benefits of adaptive temperature in controlled setting, In this section, we describe the various aspects of our experimental setup, for the purposes of clarity and reproducibility.
A.1. Motivation
We deliberately focus on single attention head environment and a simple selection function (max) to remove any confounders from our observations.
Since we are using exactly one attention head, whatever coefficients it outputs can be directly related to the network's 's belief in which items are most important for the downstream prediction. This allows to, e.g., correlate the coefficients with the ground-truti magnitude of the items.
Since we are looking for the maximal element's property, we are not requiring any complicated behaviour from the coefficients: when our target task is to approximate max, the softmax coefficients need to approximate ar argmax-which exactly what they are designed to be smooth approximation for. As such, this choice of target task exhibits high algorithmic alignment (Xu et al., 2020).
A.2. Data Generation
Let be the number of items in the set that we wish classify. For each item, i n, we need to define value, which is used to select the maximal entry. We sample these values from uniform distribution; Pr ~ (1, 1).
a priority
We would also wish our task to be classification rather than regression task, in order to leverage more robust accuracy metric. As such, let C' be the desired number of classes. We can now attach to each item a class, K2 ~ U((1...... C}2 sampled uniformly at random. We assume C = 10 fixed.
Then, for each input item, S S n, we consider its features to be x; E RC+1 to be defined as x; = ps ||||onn onnehot C), i.e. the concatenation of these two sampled pieces of data where Kę is represented as one-hot vector.
Lastly, since we will leverage dot-product attention, we also need query vector. In this particular task, the query is irrelevant and we initialise it to random uniformly-sampled value, q (1, 1).
Our task is to predict, given (xitsien and the class of the maximal item, i.c., Karg max, Pr*
A.3. Neural Network Architecture
The neural network model is designed to be simple set aggregation model (in the style of Deep Sets (Zaheer et al., 2017)), with single-head dot product attention as the aggregation function.
Its equations can summarised as follows:
h; (xx)
(9)
q (4((4)
(10)
e, (Qq) (Oa)((KO) (Kh;;
(11)
exp(e:/0)
0; Eisign exp(e;/0)
(12)
z Σ a,
(13)
15i5
y = ((z)
(14)
Equations 9-10 prepare the embeddings of the items and query, using two-layer MIPs th and vi, using the GeLU activation
function (Hendrycks & Gimpel, 2016) and an embedding size of 128 dimensions. Then, a single-head dot-producc attention (with query, key and value matrices Q K and V) is executed in Equations 11-13. Lastly, the output class logits are predicted from the attended vector using a two-layer GeLU MLP, & (Equation 14). Each component is a two-layer MLP to ensure it has universal approximation properties.
13

---

Softmax is not Enough (for Sharp Size Generalisation)
concise implementation of network using JAX (Bradbury al., 2018) and Flax (Heek et al., 2024) is follows:
import jax. numpy jnp from flax import linen nn from typing import Callable
class (ndel(nn.Moduue)) Module):
n_classes: int 10
n_feats: int 128
activation: Callable gelu
compact
def csel, q):
X = nn.
(x)
X self activvaiiiii))))
X nn. Dense Dew.eeeeaaures--eelf.....
(x)
self aclfvacctiatiion))
q = nn. Dewse www.Deeeefetuursssse.......oo
(q)
actfvacctivtiion()
= nn. (...Demnsffatturesssef.fnffffft) (q)
X = nn nn. MiLiileeddbbtrrddicttttenniii
num_heads=1,
qkv_features=self )
inputs_q=q,
inputs_kv=x)
Dww.DDeneeeeeteee Dffe ffffeatturrsssllffffeeatsts n_feats) (jnp. squeeze -2))
=
self ax)
= nn. Dense (ffateffeaures--seel wtwpseelffn.ccalsseeeeeeeeee n_classes) (x)
return x
A.4. Experimental Hyperparameters
We train our model for 100, 000 gradient steps using the Adam optimiser (Kingma & Ba, 2015) with initial learning rate of = 0.001. At each step, we present to the model a batch of 128 input sets. All sets within a batch have the same size, sampled uniformly from n ~ (1,,...., 16). The model is trained using cross-entropy, along with L2 regularisation with hyperparameter 0.001.
The mixed-size training is known tactic, designed better prepare the model for distribution shifts larger sets inference time. Similarly, the weight decay follows the recommendation in Proposition 3.1, as attempt to mitigate overfitting out-of-distribution as byproduct of sharpening the soft tmax coefficients.
Both methods prove to be effective deriving a stable baseline model.
B. Dispersion Harms Reasoning Performance
While it is intuitive that complete coefficient dispersion is an undesirable event, it may not be immediately obvious that its occurrence may have any bearing reasoning model's predictive power.
In this Appendix, we provide several corollaries and remarks stemming from Theorem 2.2 that concretise specific ways in which reasoning failures will occur as a consequence of dispersion.
Corollary B.1 (Dispersion induces reasoning failures). X is a finite set of possible values. Further, assume a strict total order on the elements of X. Assume we are solving reasoning task to find the rank of the highest-valued row x;
((n) E X be a matrix of input features for n items, where
(n) in ((n) (accordling to <), using a classifier over trained
single-head attention architecture: EIsisn an)) (x)")) where f and g are continuous functions implemented as
feedforward MLPss, and the coefficients o; (n) are computed using dot- dot-product self-attention with sof tmax normalisation (as
14

---

tmax is not Enough (for Sharp Size Generalisation)
in Appendix A). Further, assume there are no lies in the class confidences predicted by g when deciding how classify X (w), Then, assuming any floating- or fixed-point datatype with machine epsilon >> O is used to support the architecture 's data representation, it will necessarily start to make prediction errors beyond certain number of items due to the dispersion effect.
Proof: Let K be the size of the vocabulary X' = [vi..... .K ). The reasoning task presented here is effectively a K K-class classification problem, predicting the maximum rank in a set of values from X. Any prediction of the architecture must be of the form EisigK 5, (v)). with the constraints that {j 2 0. EiSisk B3 and B, = 00 & x
Now, consider two specific points v. and va, such that Va > Vb. The architecture, if trained properly, must classify (vv())) into the class, and g( (v)) into the b class.
Let ((x) be an input matrix formed such that x = va and x; vo for all < i S n. For such an input, the desired output class is a, and the prediction must be of the form 9 (af" (n) f(va) + af f(vo)).
Since the input features come from a fixed vocabulary and are processed only using feedforward networks and self attention layers, we can leverage the argument in Theorem 2.2 to conclude that there will be fixed spread in the trained architecture, 6, and further that 0; S 1 n exp 8
á for all i.
Using this we can see that, when n > 1 indistinguishable from zero, and the weighted sum will reduce g ((s(v))), due the assumed continuity of g around f(vo).
e exp é it must hold that (. < c. At this point, the value of a will be
Hence, by previous assumptions, and by the assumption that there are no ties in the class logits in g(f(v)))*, at least one of the following must be true once dispersion occurs:
The input (Va, Vo, Vbs..., ]o) of sufficiently large size will be misclassified into class 6;
The input (vo...., vo (for any size) will be misclassified.
In either case, the architecture had to have made an error.
While Corollary B.I concerns single attention heads, note that we can leverage the setting of Theorem 2.2 to prove that such failures will occur in deep Transformers as well. We sketch this intuition:
Remark B.2. Given the same task as in Corollary B.I, using a deep Transformer architecture as described in Theorem 2.2, dispersion in its attentional layers is sufficient to cause misclassifications to occur. To see why, first, assume that the models have no residual connections. The arguments for why such architectures must miselassify are subtly different depending on the Transformer model:
For BERT-style Transformers, since all attention heads are global, after one dispersed layer, any sufficiently large set (Va, Vo...., v&} will have identical embeddings to a set (V.,..., of the same size. After this, is impossible to classify them differently.
For GPT-style Transformers, to simplify the argument, we assume the Va element is at the end of the input: (Vo...., Vo: VaJ. In this setting, only the final token's attention head will receive the features from Va. It disperses, this set will once again be indistinguishable from a set (Vby..., vo) of the same size. This argument is inspired by Barbero et al. (2024).
Residual connections (He et al., 2016) allow for preserving the information contained in v. even across dispersed layers. However, as we have assumed all heads attending over V. have dispersed, no subsequent layer will be able to meaningfully integrate this information across the set, and eventually the computation will hit the final layers' attentional heads, where the final embeddings will once again be indistinguishable across these two different sets.
"This assumption is important in the case that (vv(v)) gives equal logits to classes a and b. As this is boundary condition for the classifier, if it occurred exactly on f(v)), we would not be able to guarantee that any two sets mapped to f(vs) will be classified identically without sacrificing local continuity around f(v.). Note that, due to floating-point rounding errors, this assumption is rarely broken in modern deep classifiers.
15

---

Sof tmax is not Enough (for Sharp Size Generalisation)
We note that the only condition the coefficients necessary for this breakdown oceur is that they decay towards zero-the failure on sets of the kind (Va, Vb, Vb....; } is not prevented even o: decays substantially more slowly than the other coefficients!
Remark B.3. If we assume dispersion setting where
log
i=1
Icišn
The failure described by Corollary B.1 still applies, following exactly the same proof, i.e. eventually (*) < for any machine epsilon value >> 0. Note that, as per Theorem 2.2, this situation is impossible in vocabulary-based Transformer architectures.
C. How does Dispersion Interact with Depth?
While Theorem 2.2 concludes that dispersion must eventually affect all global attention heads in Transformer architectures over vocabularies, not much is said about how rapidly the dispersion must affect heads at various depths.
Intuitively, if dispersion occurs at a particular layer, it will cause the outputs of the dispersed attention heads to converge to the average of all value vectors. This convergence, in turn, minimises the spread of logits, 5, that the subsequent layer will experience. As shown by Lemma 2.1, the value of the spread directly controls at which sizes dispersion will occur.
Using this argument, we can show that in BERT-style Transformers without residual connections, complete dispersion of all heads in particular layer leads all subsequent layers' attention heads to immediately disperse.
Remark C.I. Let (") =hh Jesgan be the input node embeddings for an intermediate layer BERT-style Transformer without residual connections. If all of this layer's attention heads have dispersed on that input, i.e. Oij
(n) where is
the machine epsilon, then all of that layer's output node embeddings will be equal to the average embedding, li(")
1 Eiisjen Vhl"), Since these constitute the inputs for the next layer's attention heads, we can conclude that of the next layer's key and query vectors will be identical, namely (for any feedforward layer f):
K(T0 K'S
vh
cLth
=
Σ Vh
71 15i50
I5iSa
As such, all logits of such a layer will themselves be equal to
) K'J
ēgg = Q'í
Σ Vh (m)
Σ vn(") Isssnn
and hence, the spread will converge to 8 0. Given Lemma 2.1, such layer can only compute averages for any input size which is equivalent behaviour to full dispersion. That is, dispersion in layer implies that all subsequent layers will output embeddings equivalent to fully dispersed ones.
Note that, we introduce residual connections in BERT-style Transformers, or leverage GPT-style Transformers, these kinds of conclusions are no longer applicable. This is because residual connections, as well as the more localised attention heads in GPT-style models, ensure that not all token embeddings will converge to the average embedding (even under dispersion). And whenever the output token embeddings of attentional layer are not fully converged, any intermediate transformations (such the K and Q matrices) can re-amplify 8 to less dispersed levels (see also Proposition 3.1).
Note this does not mean that any global attentional layer of Transformers over finite token vocabularies will escape dispersion-TTeen
2.2 proves it is inevitable it only means that we cannot tie the exact moment a particular layer's heads will disperse to a preceding layer's 's dispersion event. But the dispersion of layer will certainly play direct part in reducing the á value of subsequent layers, and this may well accelerate dispersion in subsequent layers.
16

---

Softmax is not Enough (for Sharp Size Generalisation)
D. Proof of Proposition 3.1, with Numerical Validation
Proposition 3.1 (Sharpness in Transformers necessitates large weights), Let (xx) IK be a collection of logits, computed using dot product attention mechanism; i.e. Ck are parameters. Let ò max 15i5a
(n) QQy, Kxx), where y E IRmm is query vector and Q K TRmm xm
c min 15350
(m) be their maximum difference. Then is upper bounded as:
8 S 20 mas max
lyll
max Isi5n
|Xx]||
(Q) (K) where o max maax E R are the largest singular values of Q and K.
Proof. We start by showing that the largest singular values of Q and K determine the maximum stretch due that matrix acting on E n. More precisely, wish to show:
| S
(Q) max
x||
Kx
S a
(K) max
|
where || || is the Euclidean norm. Since both inequalities have the same form, focus on Q w.l.o.g. Many of these statements can be derived from linear algebra textbooks (Axler, 2015). However, the proofs are short enough that we re-derive them here for clarity.
Consider the singular value decomposition (SVD) Q = UEVT, where E is a rectangular diagonal matrix of singular values af I. As U and V are orthogonal, |xx| Vx|| = ||x||. Therefore, ||x -IEEVV xx|| = |Ev||, where
V x, meaning that ||v| = |x||. Then we derive:
Ev|l |xx
Σ
(o (1
(Q)
Σ os
σ max
max
We now note that
% (Qy, Kxx) = Qy Kx cos 0
with 0 the angle between the arguments of the inner product. We can bound e (n) from above:
(n) < || S m
( |yy||||x
(
(K)
with a max: max being the maximum singular value of Q and K, respectively, and where the last step comes from the inequality shown above. Similarly, we obtain a lower bound, yielding:
(2) max max By | S ex
S <0(0) max max ly [X
This gives us the desired upper bound for 8:
8 = max e, ISiSn
(n)
18i5 jj
min (m)
< max o max Isi5n max
(K)
|yy |||
min
-0(9) o (K)
max
|ly |||]|||
ISi5n
o o (K)
max max
|yll
max
||x]|| + o max
o (K) |yll|
max 15j5
||>|||
Isisa
20
max |
max
||x]||
1575
completing the proof.
We remark that Proposition 3.1 lends itself simple numerical verification as well. Accordingly, in Figure visualise the evolution of the logit spread, as well as its predicted upper bound, as our single-head attentional model from Appendix A is trained for increasing numbers of steps.
Indeed, we find that the upper bound is valid, and reveal a key mechanism in which our single-head architecture gradually learns to sharpen its attention: the logit spread grows with training time, but so does the norm the relevant vectors and parameter matrices (in spite of our weight decay loss).
17

---

tmax is not Enough (for Sharp Size Generalisation)
Spectral norm of self-attention kernels vs. logit spread
1200
2 max ||ax||
1000
6
Spread Magnitude
800
600
400
200
0.0
0.2
0.6 Training steps
0.4
0.8
1.0
1e5
(Q) (K) |Iy max; ||x,]] for the Figure 9. A plot of the logit spread, 8, against its upper bound value predicted by Proposition 3.1, 20 a single-head attentional experiment described in Appendix A, with statisties computed across ten seeds. This numerically validates Proposition 3.1.
E. Proof of Proposition 3.2
Proposition 3.2 (Decreasing temperature decreases entropy). Let e(n) E Rn be a collection of n logits. Consider the Boltzmann distribution over these n. items, Pt x exp(-Be (m)) for 3 E RR and let = - E; P: log Ps be its Shannon entropy. Then, as 3's magnitude increases, H must monotonically decrease.
Proof. We start by briefly acknowledging the extremal values of 3: at 8 = (i.e., + oc), all logits are weighed equally, hence p; = (n) are uniform, and entropy is maximised. Similarly, at B + +o0 (i.e., f = 0), either the minimum or the maximum logit is given probability of 1, leading to a distribution with minimal (zero) entropy.
Now, consider the partition function = Σ, exp(-3e,") (*)), such that p. exp(-Be(")) We will take derivatives of log with respect to Starting with the first derivative:
log Z
Z Σ
(x)
exp(
-ppl-Bep BBe
(n))
=-E
c Pj -Einp, (c)
(n)
db
18

---

Softmax is not Enough (for Sharp Size Generalisation)
we recover the expected logit value sampled under the distribution. Now differentiate again:
d²
log Z -
(m)
d32
d3
=- - (n) exp( exp(-Se, (n))
c
dB
Z
- (n) exp( expl-Bee -3e, )2 exp(-BBe, (1) Ej -e (n)
exp(-Bee) (m))
Σ
(n)
722
E(e!") exp(ppp00 (Bpp (n)
(n) expl (xxp)
Σ c/
(n)
exp(-Bek
Σ (gm)
E(elml)pp. -
Σe J Pj &k Pk
Bimp (((()))) Ei-p, (c) (1)) = Variup (10)
and we recover the variance of the expected logit value. Now we turn our attention the entropy formula:
Σ Pt logpi = Σ Li((logep((Bee)) log 2)
=
Σ pt log Z Σ -3e
= log Z + BEipps (e) = log Z - 3 d
log Z
dB
To check the monotonicity of H as B3 varies, we now take the derivative of this expression w.r.t. 3::
dH d
d log 7
d2
log 7 -3 d²
log 7
log Z -aVar (c(
d3
d3
d3
d32
d332
Since variance can never be negative, we find that dn magnitude |3] grows, the value of H must monotonically decrease.
< 0 when 3 2 0. and
dll
S 0 when B S 0. As such, as the 23
2
F. Derivation of the Adaptive Temperature Scheme
In this Appendix we expand on the description of our adaptive temperature scheme in a way that provides a detailed step-by-step overview of the individual steps taken to arrive at the final JAX algorithm presented in Figure 4.
The core principle of our approach is to adapt the 0 parameter within sof tmax, in a way that is mindful of the entropy of its inputs; that is, we first compute:
Ô = ((oo
(15)
where e E R is the vector of logits inputted into sof tmax, 0 the initial temperature (here, we set f for simplicity). H [0, 1jm IR is the Shannon entropy of the distribution emitted by sof tmfx/zzee):
H(p) = Σ ps log
(16)
and f R > Ry0 is a (learned) function mapping the calculated entropy to an adapted temperature value, (which should be nonnegative). Once computed, our sampling scheme returns softmax, (e) instead of sof softmax,(( (e) as the final distribution. Therefore, the only remaining important element to understand is how the function constructed.
19

---

Softmax is not Enough (for Sharp Size Generalisation)
While one could approach the construction of / using a complex learned function, for the purposes of this work we focus on first learning simple inverse-polynomial degree-4 approximator, f:
Σ a&H +
S(H)
(17)
In order to derive the coefficients a,, need a dataset of (H, 1/pppt) values, mapping the observed entropy in particular sample to an optimal choice of temperature scaling, Oopt From such dataset, one can use standard functions (such as nolpfiffit) to fit the coefficients.
We have a straightforward way of computing H for any particular sampled training example- - simply record attentional logits at every Transformer layer and compute the Shannon entropy. How can we compute the optimal temperature value for this example least in the case of our single-head set Transformer from Appendix A)? It should be the value that would place the highest probability possible the item are looking for in this case, the item with the maximal value.
For starters, note that, for Sopt to even be meaningful in this case, the item with the maximal value must not have the largest logit already. In such cases it is easy to see the "optimal" temperature is zero, leading to an unwieldy target value of -o0 when learning f. Hence, we will discard all samples where this is the case.
For all other samples, we know that the probability of the target item will be 0 at (as will be overwhelmed by the maximal logit), and it will be 1/n at 0 + oc. Somewhere between these values lies the optimal choice. Rather than using elaborate methods search for Oopt we perform simple grid search attempting all values of between 0 and 10, in increments of 0.1. Empirically,, find that for all samples cared about, Pupt was definitely in this range and record the 0 that led to the highest observed probability as our estimate of Oopt.
Finally, we would not wish to "overcorrect" the temperature if the attention head is already confident (entropy below a fixed threshold we chose 0.5) or, given the target of our work towards sharpness, if the suggested temperature scale would disperse the coefficients further (i.e. if would result in a temperature greater than 1).
All together, we recover the following final form of our algorithm, implemented by Figure
1
H S 0.5
f(H) = min (i(M))
(18)
H > 0.5
G. An Algorithm for Streaming Attentional Entropy
Computing our proposed adaptive temperature requires computing the entropy of the attentional coefficients. A naive algorithm for doing so requires fully materialising the oiyy entries of the attention coefficient matrix, which requires Onn²) memory and poses scalability concerns. Fortunately, there exists online algorithm for computing the entropy that is not FLOP/s efficient but does not leverage any additional memory, allowing for a linear-space attention implementation in conjunction with Flash Attention (Dao et al., 2022). We present one such algorithm in this section. We have successfully implemented this algorithm and numerically verified that its outputs match the expected adaptive temperature amounts, allowing us to deploy layers with large context windows (up to 131, 072 tokens) on a single NVIDIA A100 node.
In order to compute the adaptive temperature, we need to first compute the attentional coefficient entropy for each row the attentional matrix. For convenience, let us define the exponentiated logit of token i's attention over token j, taking into account only the first 1 S N < items:
(N) exp
q. k,
max ks) KEN
where qi and k; the query and key vectors, respectively, for token i.
Now, we can rearrange the terms of the expression for the entropy, H. N). of each row of the corresponding matrix of
20

---

tmax is not Enough (for Sharp Size Generalisation)
attentional coefficients, taking into account the first N items, in form that will be more favourable for streaming:
(1
(N)
H = 11
Nig
(N)
(N)
((
log
(N)
Σx (N)
k Ask
(N)
log 11j log
Σ
(M))
Σx k ii
((N)
ii
k
(N)
(N)
Σ
log (N)
Σ
Aij
log Σ 1
Σx (2V)
Ajj
Σn A
(N)
j
E, 1 (N)
log Ai
(N)
Ej A
(N)
Σ ((N))
log
Σx 1(N)
Σx Aik
(N)
h
Σ, 1/3 (N)
log 2
-log
Σxie x
Σx A(k
k
Next, we define two cumulative quantities:
(N)
(M)
m;
(N)
= max qi kj
JEN
JKN
which allow us to further analyse the E3 11 (N)
log 1 (N)
term as follows:
Σ
(N)
log (07
Σ exp q qilk, k, - max a kk log exp q. k, max q Kk
JCN
JcN
Σ 1 (q. kj- - mi )
JCN
(N) q k, -m A (N)
JKN
Now we remark that incrementally compute A, using the following iterative formula, leveraging the same concepts as Flash Attention (Dao et al., 2022):
1 A(7) Σ
(N)
Σ exp
(g. k, - (mmmm -m
JeN
JEN
A (N+1) A
mm (N)
mm (((ttl) 1/11)
exp
and we can incrementally compute the remaining term, x =EjeN 11y (N)
=
q ky. using the following iterative formula:
kc! (N) = Σ (N)
Xij q. k,, = E] exp (a kg -mm ml) q kj
JKN
JKN
x- x(N+1) N!!N exp (m! ! - - (N+2)) A q KN
So our final result in terms of (n)) and k((") (fully streamed across all items) is: (n) A(7) (n)
x!m) - m
m1 m
log A (n)
A (m)
((())
m (n) - -log A! (m)
(m)
21

---

Softmax is not Enough (for Sharp Size Generalisation)
This expression can be computed with O(n) memory, never have materialise entire matrix of coefficients. Under this implementation, adaptive temperature easily scale large context windows (which we have validated empirically up to 131, 072 tokens).
22


Reviews:
[{"review_id": "IxnDW8avNv", "reviewer": "Reviewer_WiLd", "summary": "This paper studies the *sharpness* of the softmax function from a *size generalization* perspective. The authors regard a function as being **sharp** if its output can be expressed using a constant number of inputs. The authors refer as **size generalization** the study of what happens when the function is subject to a larger number of inputs. In this paper, the authors argue that using an adaptive temperature parameter can help preserve sharpness by lowering the temperature enough to reduce entropy while maintaining the trained model accurate. More generally, the authors argue in their main theoretical results that it is not possible to preserve the sharpness of the softmax function as the number of inputs grows arbitrarily large.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer WiLd,\n\nWe 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.\n\nTo that end, we address all of your points in order:\n\n### **On significance of our results**\n\nWe 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.\n\nTherefore, we will focus this part of the response on discussing whether our results are evident.\n\nDue 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:\n\n* [Reviewer n5x1] _The paper's observation is simple, but thought-provoking._\n* [Reviewer jdu5] _I appreciate that this work presents the problem well and raises awareness of this issue._\n* [Reviewer 5ZBY] _The limitation on softmax and adaptive temperature is relevant for the boarder audience, given the prevalence of softmax in modern ML system_\n\nWe 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!\n\nIt 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:\n\n* The overall prevalence of the Transformer architecture and the many expressivity results for it;\n* The elevated importance of the dataset for training AI models, and the diminished importance of architecture choice;  \n* Mechanistic interpretability research, which reverse-engineers sharp behaviours in trained LLMs.\n\nSuch 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.\n\nHowever, 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.\n\nWe 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!\n\n### **On thresholding contributions**\n\nWe 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.\n\n### **On the infinity dimensions and the ‘scaling up sharpness’ framing**\n\nTo 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.\n\n### **Improvements to `softmax`**\n\nOur 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.\n\nWe 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.\n\nWe are happy to add this discussion to the revised paper!\n\n### **Miscellaneous issues**\n\nWe are happy to correct all minor nits you pointed out, as well as avoid usage of terms like ‘reasoning device’ in a revised paper.", "claims_and_evidence": "My main concern with this paper is not correctness, but rather significance: the theoretical results claimed do not seem surprising or nontrivial to prove for someone working on that line of inquiry. From the examples at the top of Page 2 that the max function is sharp and the average function is not sharp, the fact that softmax is not sharp for an arbitrarily large number of inputs seems evident and proving that result seems in line with a doctoral-level homework exercise. \n\nMoreover, if we are to assume that softmax is sharper with smaller inputs or subject to a lower temperature parameter, then I believe that we lack a definition for proper theoretical discussion: there should be a threshold for the minimum contribution of an input to consider it relevant to the output of the function (and possibly how many inputs significantly contributing to the function output is too many). Otherwise, any contribution of an input should be counted, and then softmax is trivially not sharp.\n\nMore generally, any discussion about models when their dimensions tend to infinity changes the nature of the beast. For example, arbitrarily deep or wide neural networks may hold properties that a finite neural network architecture cannot promise. Hence, I believe that the proper framing here should have been about scaling up sharpness with respect to input size - and how to overcome the challenges associated with that.", "methods_and_evaluation_criteria": "See item above.", "theoretical_claims": "See two items above.", "experimental_designs_or_analyses": "See three items above.", "relation_to_broader_scientific_literature": "I am not sufficiently familiar with the line of work to which this paper contributes to make a comment on this.", "essential_references_not_discussed": "I am not sufficiently familiar with the line of work to which this paper contributes to make a comment on this.", "other_strengths_and_weaknesses": "Strength: the authors have a clear writing and frame well some interesting aspects of the theoretical work around attention. I am curious to read more about it following their description.\n\nWeakness: I personally find the terminology \"reasoning device\" speculative. I would recommend making the discussion more objective without such and related terms.", "other_comments_or_suggestions": "The use of rephrasing in theoretical statements (the \"That is\" in Lemma 2.1, Theorem 2.2, and Proposition 3.1) is not adequate. If needed, those can be added either before or after the formal statement (or after the proof), but not in it.\n\nDrawing conclusions inside a theoretical statement (the \"Thus\" in Proposition 3.2) is not adequate. If relevant, that part should have been a separate corollary after the parent result.\n\nAbstract:\n- ' \"circuits\" ': replace '' with `` before this word\n\nPage 1:\n- \"does not have a chance\": too informal\n\nPage 4:\n- In the proof of Theorem 2.2: \"for [some choice of] constants $m$ and $M$\"", "questions_for_authors": "By the last paragraph of Section 2, my impression is that the whole argument of this paper is that there are clear limits to the transformer architecture if taking arbitrarily large inputs. Is that the case, or do you believe that there is a better function than softmax for the purpose that it serves?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EnJLGGYDne", "reviewer": "Reviewer_5ZBY", "summary": "The authors argue that modern deep learning architectures are fundamentally incapable of learning sharp functions(for example, max) due to the disperse nature of the softmax function in out-of-distribution settings. In addition, the authors propose an adaptive temperature mechanism as a plug-in technique at inference time for improving the sharpness.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer 5ZBY,\n\nThank 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!\n\n> 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?\n\nThis is an excellent question! \n\nWe 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:\n\n* [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).\n* [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.\n* [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.\n\nThe 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.\n\nThe 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).\n\n> 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.\n\nThank 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.\n\nWe 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.", "claims_and_evidence": "The authors argue about the limitation of the softmax function and provide both theoretical and empirical evidence for it.\nThe authors also provide theoretical and empirical visualization to support adaptive temperature.", "methods_and_evaluation_criteria": "Both the argument around the softmax limitation and the proposed method make sense. \n\nHowever, 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.", "theoretical_claims": "I checked Lemma 2.1 and Theorem 2.2, which seem correct.", "experimental_designs_or_analyses": "I checked the settings for max retrieval and CLRS-Text, and the experimental design and analysis make sense.", "relation_to_broader_scientific_literature": "The limitation on softmax and adaptive temperature is relevant for the boarder audience, given the prevalence of softmax in modern ML systems.", "essential_references_not_discussed": "I am not an expert on theoretical work around softmax or out-of-distribution. However, the discussion on the background, primer on attention heads and transformer, and related work on adapting temperature seem complete to understand the problem and proposed method.", "other_strengths_and_weaknesses": "I think the primary significance of this paper lies in the discussion on softmax limitation.  My main concern is discussed in the evaluation criterion area.", "other_comments_or_suggestions": "None", "questions_for_authors": "1. 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?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mH2Ly2khm7", "reviewer": "Reviewer_jdu5", "summary": "The function softargmax (commonly referred to as softmax), which is used to create probability output vectors and attention heads within neural networks, becomes less like argmax (less sharp) as the number of elements over which the softmax is applied increases. This is detrimental for learnt circuits within transformer architectures that need sharp attention, especially when deployed at inference time on longer sequences than presented during training.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer jdu5,\n\nThank 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!\n\nWe provide detailed answers to your comments below, and we hope they will be to your liking.\n\n### **Bound looseness**\n\nWe 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!\n\n> 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.\n\nJust 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.\n\n> 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.\n\nIn 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:\n\n> 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.\n\nThis is a fantastic suggestion and we really appreciate it!\n\nTo 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:\n\n| **Model**    | **Average** $\\delta$      | **Maximum** $\\delta$ | **Minimum** $\\delta$ |\n| :------- | :----------------------: | :-------: | :--------: |\n| Gemma 2B | 5.69 ± 2.05    | 14.78 | 2.28   |\n| Gemma 7B | 5.82  ± 2.61    | 32.74  | 0.09 |\n\nThis 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!\n\n> Additionally, I think there could be more discussion on the factors which prevent the model from achieving an arbitrarily sharp distribution\n\nAnother excellent proposal – we fully agree with all of your specific factors, and will make sure to thoroughly discuss those in our revision.\n\n### **Specifics of Adaptive-$\\theta$**\n\nAs 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.\n\nWe 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.\n\n### **On Figure 5**\n\n* We acknowledge your point regarding `jet`. We agree that perceptually uniform colormaps like `viridis` or `plasma` offer better visual representation, and will switch!\n* We used power series (controlled by $\\lambda$) as a simple model to represent sequences with varying degrees of \"sharpness\" in their logit distributions.\n* 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.\n\n### **Miscellaneous issues**\n\nWe 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.", "claims_and_evidence": "The proofs given in Lemma 2.1 and Theorem 2.2 are rather weak with respect to the arguments made in the paper. For example, Lemma 2.1 assumes that logits are bounded below by $m$ and above by $M$, both of which are finite values, and provides a proof for the limiting case $n\\to\\infty$. This is an incredibly loose bound. 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. 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.\n\nThe authors also included empirical studies (Fig 2 and 3), where a model is tasked with identifying the maximal element in a sequence, which I appreciated. However, I think it would help support the paper if there were additional empirical measurements. 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$.\n\nAdditionally, I think there could be more discussion on the factors which prevent the model from achieving an arbitrarily sharp distribution (e.g. label noise prevents the model from learning arbitrarily large parameter values; the derivative to make off-target logits arbitrarily negative can vanish when the softmax output is already sufficiently sharp).", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "The proofs in the main paper (Lemma 2.1, Theorem 2.2, Prop 3.1, Prop 3.2) are correct.", "experimental_designs_or_analyses": "For the experiments comparing the adaptive-$\\theta$ model presented by the authors against the baseline model for the max task (Table 1), I think these results could be usefully expanded by including additional \"oracle\" measurements which use the optimal theta at inference time. This will help to indicate how much of the possible gains which could be attained by only changing the temperature were achieved by the adaptive temperature.\n\nThe discussion of the comparison for the adaptive-$\\theta$ method introduced in the paper (Fig 8.) could be more detailed. In particular, some tasks see the adaptive temperature model perform worse than the baseline (namely heapsort, mst kruskal, and bubble sort). I would appreciate if the authors could comment on whether this deficiency is meaningful, for instance are there some features these tasks have in common which makes adaptive-$\\theta$ perform poorly here? Is it just because adaptive-$\\theta$ was fit on the max task and does not generalize well to these tasks?", "relation_to_broader_scientific_literature": "The issue of softmax dispersion is already known within the community broadly speaking, but I appreciate that this work presents the problem well and raises awareness of this issue.", "essential_references_not_discussed": "Not that I am aware of.", "other_strengths_and_weaknesses": "None.", "other_comments_or_suggestions": "**Figure 5.**\n\nThe title and y-axis indicates that in half the domain you are dividing the logits by a negative number, which doesn't make sense w.r.t. the meaning of temperature. I suggest this is refactored. The implication is more like the sequence is changed by negation rather than the temperature adjusted. But the sequence is controlled by lambda on the x-axis, so it is odd to have two different sequences stacked one above the other and joined at the $\\theta=0$ line To be honest, I am not sure what I am supposed to learn from this figure anyway. Is this power series supposed to be indicative of real data? If so, then how is it related to real data?\n\nThe \"jet\" colormap used here is not perceptually smooth. Its use creates the visual appearance of discontinuities at the luminosity inflection points (e.g. yellow and cyan). I recommend using a more modern colormap instead. There is a good reason why the defaults in matplotlib and MATLAB have both moved away from jet several years ago. Seaborn throws an error if the user requests to use jet, rather than comply. There are many resources which have published which discuss the issues with jet, e.g.\n- https://www.youtube.com/watch?v=xAoljeRJ3lU\n- http://jakevdp.github.io/blog/2014/10/16/how-bad-is-your-colormap/\n\nIf you like jet because it breaks the data down into blocks of distinct colours and feel that the new colormaps with smooth luminosity changes (e.g. viridis) are too smooth making them hard to read, then you should be breaking the data up into discrete colours with intentionality, rather than as the side-effect of a poor colormap. In this case, you should use a contour plot, either with contour lines or with discrete blocks of ~24 colours, instead of a smooth colormap.\n\n**Typos and snags.**\n- L018 \"circuits\" incorrect quote mark\n- L055 right column, missing word \"A strong current [topic] in this space\"\n- L073 right column, it would be helpful to know how samples are OOD for this point (i.e. is it OOD in sequence length or not, as that is pertinent to the topic of the paper)\n- L105 right column, \"a collection of [n] nodes\" as n is the number of nodes, not the label of the collection.\n- L112 The notation is a bit confusing since k_i and v_i are vectors instead of elements within vectors.\n- L124 Eq 3, suffix variable `j` should be `i`\n- Eq 3, 4. Adding a comma between equations on the same line may help delineate between them.\n- Fig 1: It appears that the caption text colours are intended to indicate the block types within diagram. However, the token and MLP colours are too similar in the figure and the text colours are not quite the same shades as the figure colours. I thus recommend the authors add labels within the figure, and/or describe the colours in the caption to make the reference clear.\n- Fig 2: What's the number of items used in the base case? Please add this to the caption.\n- Fig 3: Add units for y-axis (presumably bits)\n- Fig 4: ~L222 \"see Figure 5\" should be \"see Figure 6\"\n- Tab 1, L382: There shouldn't be a space after the thousands separator, having this in makes the numbers harder to read.\n- Fig 8: Some tasks have the adaptive temperature model worse than the baseline (heapsort, mst kruskal, bubble sort) Can the authors comment on what features these tasks possess\n- Fig 8: Legend should be at the bottom of bridges or activity selector instead of the top of bfs so it doesn't obstruct data.\n- L692: Usually it is just called the Adam optimizer, not the Adam SGD optimiser. Referring to it as such may cause confusion.\n- L654: Hard-written \"Equations 2--3\" is incorrect and should be replaced with actual equation numbers as reference.\n\n**Citations:**\n- Be careful with protecting casing of acronyms.\n  - L462: \"Transformers need glasses! information...\"\n  - L497: \"clrs\"\n  - L583: \"gpt-2\"\n- Please consider adding links to references which currently don't have one. This makes it easier for a reader to navigate to the cited works. You appear to have your arXiv references set up in a consistent manner, so adding either a URL or a DOI field can be done automatically for these with a regex call.", "questions_for_authors": "None.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Wg0BzLAFu9", "reviewer": "Reviewer_n5x1", "summary": "This paper can be divided into 3 parts:\n1) The observation and proof that softmax-based architectures (such as Transformers) will have a \"dispersion\" phenomenon when tested on longer inputs than they are trained on.\n2) The observation that this dispersion phenomenon can degrade the length-generalization performance of transformers on simple algorithmic tasks such as finding the maximum element in a list.\n3) An ad hoc \"adaptive temperature sampling\" scheme that seeks to remedy the dispersion phenomenon and leads to performance improvements on the \"maximum element\" task and on several problems in the CLRS-text benchmark.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer n5x1,\n\nWe are highly pleased to read your review, and really appreciate your positive view on our results and their significance! \n\nIn what follows, we reply to all of the points you raised:\n\n> 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.\n\nWe 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.\n\n> I had some trouble following the explanation of the adaptive temperature sampling scheme.\n\nWe 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.\n\n> 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\".\n\nThis is a great suggestion, and we will tone down by removing the word ‘important’ here.", "claims_and_evidence": "The proofs are clear.", "methods_and_evaluation_criteria": "Yes, these make sense.", "theoretical_claims": "Yes, I checked the main theorem and its proof.", "experimental_designs_or_analyses": "Yes, I checked the experimental details for the maximum task and the CLRS-text task and they seemed ok to me. The adaptive temperature scheme is arguably very ad hoc, but it seems to give some small gains.", "relation_to_broader_scientific_literature": "The paper's observation is simple, but thought-provoking. The maximum-element task is a convincing illustration that this dispersion phenomenon in LLMs is real and could be one of the barriers to length generalization. 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.", "essential_references_not_discussed": "Not insofar as I know", "other_strengths_and_weaknesses": "I had some trouble following the explanation of the adaptive temperature sampling scheme.", "other_comments_or_suggestions": "The paper is written in a somewhat informal style, which I am fine with because it is mostly clear. 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\".", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: After a productive discussion, there is consensus among the reviewers that this paper provides a solid contribution to the ML community. I urge the authors to incorporate the feedback by the reviewers, particularly by Reviewer WiLd, into the final version.
Decision: Accept (poster)