The Softmax Bottleneck Does Not Limit the Probabilities of the Most Likely Tokens
Authors: Ronen Basri
David Jacobs
Venue: ICLR 2026 Poster (2026)
Keywords: Softmax Bottleneck+
Transformer+
Output Projection Matrix+
Large Language Models+
Primary area: foundation or frontier models, including LLMs
TL;DR: We show that randomly initialized or trained output projection matrices can successfully produce exact probabilities for the top m tokens for rather large values of m.

Abstract: In many popular transformer architectures, an output projection matrix linearly maps lower-dimensional embeddings into a higher-dimensional space of logits. It has been shown that this leads to a softmax bottleneck that prevents the production of arbitrary probability distributions.  It has been argued that this limits large language models (LLMs) in their ability to express next token probabilities that perfectly align with the statistics of natural language.  We focus on the ability of such models to produce accurate probabilities for just the top-$m$ tokens.  We provide theoretical bounds that show that even a randomly initialized projection matrix can successfully do this for rather large values of $m$, supported by empirical results on both random and trained matrices. This raises questions about whether the softmax bottleneck significantly limits the capabilities of LLMs. We also derive bounds on the maximum number of probabilities that any trained output projection matrix can specify.

Published as a conference paper at ICLR 2026
THE SOFTMAX BOTTLENECK DOES NOT LIMIT THE PROBABILITIES OF THE MOST LIKELY TOKENS
Ronen Basri
David W. Jacobs
Faculty of Mathematics and Computer Science Weizmann Institute of Science ronen. ronen.basr! i @we zmann. ac
Department of Computer Science University of Maryland d djacobsé cs umd. edu
ABSTRACT
In many popular transformer architectures, an output projection matrix linearly maps lower-dimensional embeddings into a higher-dimensional space of logits. It has been shown that this leads to softmax bottleneck that prevents the production of arbitrary probability distributions.
has been argued that this limits large language models (LLMs) in their ability to express next token probabilities that perfectly align with the statistics of natural language. We focus on the ability of such models to produce accurate probabilities for just the top-m tokens. We provide theoretical bounds that show that even a randomly initialized projection matrix can successfully do this for rather large values of m, supported by empirical results on both random and trained matrices. This raises questions about whether the softmax bottleneck significantly limits the capabilities of LLMs. We also derive bounds on the maximum number of probabilities that any trained output projection matrix can specify.
Introduction
Transformers 25 have had a huge impact on machine learning, forming the basis for most large-scale foundation models. Understanding the virtues and limitations of the components of transformers is, therefore, an important challenge. Recent works have identified the softmax bottleneck produced by the final linear layer and softmax operation, which limits the output space transformers. paper, provide new analysis of the softmax bottleneck and propose new practical ways to assess the expressiveness of transformers in the context of next token prediction. We show that, under this lens, limitations in the expressiveness of transformers may be less significant than previously thought.
this
The softmax bottleneck concerns the output projection matrix (also called the unembedding or readout matrix), which provides linear map from relatively low-dimensional embeddings into the higher-dimensional space of logit for each token. Large Language Models (LL.Ms) apply softmax to these logits produce probability distribution over the next token in sequence. Work the softmax bottleneck points out that the output projection matrix (OPM) maps embeddings to low-dimensional linear subspace in the space of all logits, limiting the possible distributions over the argue that this limitation handicaps the performance of LLMs.
next token. 26 13
6
key question which probability distributions cannot be represented by LLMs due to this bottleneck. and how significant are they to performance? When LLMs use nucleus sampling, only the most likely tokens affect the next token probability at inference time (TT). When text is generated using methods like beam search 8 5] it is also the case that owwprobabiliity tokens may not affect generation. This suggests that it may be most important for LLMs to represent the probability of the most likely tokens as accurately possible, with the exact probability of very unlikely tokens mattering less.
For this reason, we consider the following question. Given any arbitrary set m tokens, does the OPM allow us to produce probability distribution in which they are the most likely tokens? And given specific probabilities for those m tokens that sum to nearly 1, how closely can we generate those probabilities in an LLM? How do the answers to these questions vary with m??

---

Published as a conference paper at ICLR 2026
We show that for realistically sized models, there exist OPMs such that any set of tokens can be assigned any specific probabilities that sum to near for rather large m. Moreover, this can be done using a random matrix initialized with i.i.d. Gaussian entries. For example, for the embedding and vocabulary size used in GPT-2, we derive a theoretical lower bound that shows that with probability near one, can generate any desired probabilities (that sum to nearly 1) for up to any Ri 26 tokens using a random OPM. Experimentally, we show that, in fact, this is possible for m up to about 90. For larger models, like GPT3-175B the lower bound rises to over 400, while for Llama-7B the empirical bound rises to over 1,000. We conjecture, and support with empirical evidence (see Appendix (A), that inference time the probabilities of tokens beyond the top 1,000 will not significantly affect performance when nucleus sampling is used. This further raises questions about whether the softmax bottleneck will affect performance significantly at training time.
We then consider the potential capabilities of trained OPMs. First, we examine the following question. What is the largest m for which there exists an OPM such that, for any collection of m tokens, there exists an embedding that the matrix will map to distribution in which those are the most probable tokens? We use results on the sign rank of matrices to theoretically bound this m, showing that is significantly higher than even the large values of m produced by random matrices.
Finally, we experimentally examine the trained OPMs for GPT-2, TinyLlama, T5-Large, GPT2-XL, and Llama2. We find that their ability to specify probabilities for random collections of tokens is similar to that of random matrices.
Our primary contribution is show that although the softmax bottleneck limits the set of probability distributions that an LLM can produce, it does not do so in a way that significantly restricts its ability to accurately determine the probability of the most likely next tokens in a sequence.
2
Prior Work
26 provided the initial discussion of the softmax bottleneck. They show that models that linearly map a low-dimensional embedding space to a high-dimensional logit space can only represent a low-dimensional linear subspace in the space of log probabilities of all tokens. They conjecture that in natural language, the space of log probabilities is of high dimension and cannot be captured by this subspace. They propose mixture of softmaxes to overcome this. [13] provides a detailed analysis of the bottleneck. They propose a softmax function that maps the logits to a non-linear subspace, making it possible cover wider range of probability distributions. (2) show further limitations. They point out, for example, that if the words king, queen, woman, man form a parallelogram, it will be impossible to assign the highest probabilities to both woman and king. They address this with a multi-facet softmax.
These works suggest that model performance is limited by the effective rank of the OPM. 20 raises questions about the relationship between rank and performance, and suggests that improvements from these methods may be due to other factors, such as implicit regularization. More recently, (7) presents empirical evidence that when the rank of the OPM is reduced, model accuracy drops and training loss increases, along with further analysis.
(6) proposes the arbitrary distributions and those expressible with OPMs, and measure the ability of models to correctly predict the top (m 1) token. (9) further studies the m case, building on results of (3. They develop a linear program that can determine whether a trained OPM can select a specific token as the most probable one, and show that in many trained models, virtually all tokens can be selected. We use similar linear program to evaluate this question for the top-m tokens for m 2 1.
They also theoretically analyze the distance between
Overall, number of papers have proposed methods mitigate the effects of the softmax bottleneck. showing some improvements in performance as a result. At the same time, the work of (20) points out that these approaches may be leading to improvement gains for other reasons. Our work raises further questions about whether the softmax bottleneck is truly limiting performance, especially at inference time. We suggest that remains an open question whether the softmax bottleneck is a problem in need of solution.
2

---

Published as a conference paper at ICLR 2026
Our work uses the sign rank of matrices to prove lower bounds the embedding size needed to represent the probabilities the top m tokens. Sign rank has surfaced in machine learning in various contexts [ Of some relevance is the recent work of [15]. which uses the sign rank to highlight the limitations of graph representations using attention layers.
More broadly, transformer architectures have been analyzed from wide range of perspectives.
(7 28 12] prove that transformers are universal permutation equivariant approximators over compact domains. (16) analyzes the approximation of sparse attention matrices.
Other papers explore the computational power of transformers, particularly with Chain of Thought [18]
10 21). Finally, (19) shows how transformers can implement gradient descent learn in-context tasks at inference.
3 Preliminaries
We consider an LLM that, in its penultimate layer, has produced d-dimensional embedding, for each token in the context. We consider the output projection matrix, before and after training. This produces N logits, where N is the vocabulary size, so A is N x d. On initialization, A will be random with entries drawn from an i.i.d. Gaussian distribution with zero mean. For simplicity, we assume a variance of 1, but changing the variance does not affect our results. produces logits: y Ax. Letting a, denote the j th row of A, we can also say that yj = aạ x. Softmax then produces probabilities as:
Pj
Ei=1 CV.
We assume temperature of here, but our results easily generalize to arbitrary temperatures in the softmax. Without loss of generality, we assume that we desire the first m tokens have the highest logits, and denote their corresponding logits ym and probabilities rows of as Am, so that ym = x. In the next section, we consider the question: given specific values for Pm and random matrix A, what is the probability that there exists embedding, that will produce the specified values of Pm. along with Σ E/=m+1 Pj 8 for some small 8.
Pm- We denote the first m
4 Random Matrices
In this section, we consider the initialization of A as random matrix, A 1), We begin in Section 4.1 by considering the probability that there exists an x so that > Mý Vi,j such that i S m < J. This ensures that the first m logits are the largest, which means that the corresponding tokens have the highest probability. We point out that the probability that such an exists is maximized when all y/i are equal and positive, for i S m. Then in Section 4.2 we show that with a slight variation, we can ensure that the p:, S m, all have any desired probabilities that sum to nearly while J i > m can be made arbitrarily small.
4.1 Specifying the top-m logits with equal value
For random output projection matrix A, we derive an approximate lower bound on m², the maximal number of tokens for which A is nearly guaranteed to produce higher probabilities for any m-tuple. Proposition I states that for random OPM, the probability that random set of m tokens has an embedding that makes them more probable than all other tokens (Ihs) is bounded by an exponential of cumulative Gaussian distribution (rhs). We prove this bound by constructing a simple embedding that provides logits for the top-m tokens (WLOG the first m), and then estimating the probability that any other logit will be smaller than the largest of these. Corollary [ then shows that this bound is maximized (and simplified) when the top-m logits are equal.
3

---

Published as a conference paper at ICLR 2026
Proposition 1. Let A be an N d matrix whose entries are drawn from i.i.d. standard normal
distribution. For an arbitrary embedding vector x, we let y = Ax Let m < < N. Then,
[minggjm)) (m)] Vi
P(3x, s.t. (ym+1...., UN < 9l: .... vm)) 4
(2)
y ll
where is the variance of the elements of a Am (A A m + S s N, and is given
approximately by
d(d - 1)
2 (d- 1) -3)
(3)
with the corollary: Corollary 1.
1 me
max P(3x, s.L. Ym = Am X, (ym+1,.., yN < yl. Vm)) 2 Nm
Ym
and the maximum is obtained for ym. = clm for any > 0.
Here mm = (1,.... 1)T TR
Proof. To derive bound, will consider only embeddings that lie in the row space of Am- We denote such an embedding by x So X|| = A, w for some weight vector w E R"" and therefore
Ym Amx AmA"...
(5)
and w is determined by w = (Am - ymm Am AT is invertible with probability 1. Thus, there is unique X| in the row space of A that produces ym. Next, the bottom N logits given by
yj = = = a. A w a7 ATh (Am Ym. < N
(6)
Next, we determine the distribution of y/ and express the probability that it obtains smaller values than y; for all i [m]. Since the elements A are i.i.d. (0, 1). the m x m matrix (Am4"") )-1 follows the Inverse Wishart distribution with d degrees of freedom, denoted (I, [17] (Sec.
3.8), whose mean is μ --m-I (Im denotes the m identity matrix) and its variance, V, is
Im
V Var Am A -1 (d + 1) (d m m
(d (d-m)( (d m 1) - m
Next, note that the elements of az A E T roughly as a? AT ~ (0, d). For simplicity, we approximate by assuming these values and the en- tries of (Am4T) are statistically independent. We experimentally verify that the resulting approxi- mation is accurate in Appendix B We obtain that the elements of the vector a A Am) R have zero mean, and their variance is given by d(V + n²)) Plugging in the values for μ and V, we notice that all the entries are i.i.d. with mean and variance given by
of d normal products, so they distribute
d(d 1)
(8)
(d - (d - m 1) m - 3)
Approximating this distribution with Gaussian (0, v). we obtain
yj ~ (O, 11.
(9)
This implics that for in the space of Am such that ym Am. have
(minse(m) 34 Il VVu
P( UN S min yi) xe N-
(10)
vElm
where denotes the CDF of the standard normal distribution. Note that VN- is the CDF of the Gumble distribution with location parameter (AN-m = --
x- and scale dN-m =
N N-m The bound in implied. It can be readily verified that this probability is largest when 1/1 ym 0, and is invariant scaling of y, proving (4 n
[N
4

---

Published as a conference paper at ICLR 2026
In Table I (Appendix Cj. show the value of our lower bound for 22 common models. We see values ranging from m = 26 for GPT2 to m = 418 for GPT3-175B. This shows that for larger current models, the number of tokens that can be specified is quite high.
Note finally that this analysis restricts x to lie in the row space of Am. and so it provides only a lower bound on the top-m probability. Using linear programming. show empirically in Section 6 that adding a component Xi, orthogonal to the row space of Am. leads to significantly higher probabilities.
4.2 Specifying the top-m probabilities
Suppose now that we want to select x to exactly specify the relative probability of the top-m tokens. This is more complex, since softmax causes the probability of a token depend on the logits of all tokens. We first show how to construct an embedding in which the ratios of the top-m probabilities achieve specified values, but the sum of these probabilities, and so their exact values, are not fully specified. Then we show how to scale and adjust this embedding so that the sum of the top mn probabilities approaches I. When these values sum close to one, the probability that such an embedding exists for a random OPM is given by our previous results in Prop.
Formally, say that E Xp........ if and only if o Pi for 1 S S m, and we want to select e Pm
x from this set. Without loss of generality, we always assume that Pm S Pt for i < m. Further, we suppose that only these top-m tokens should be generated, so we wish to assign low probabilities to all other tokens. We will consider the event that a random matrix can generate probabilities so that the top-m tokens have specific ratios of values while the remaining probabilities collectively sum to d or less, for some small ò. Our analysis below derives a bound on the probability of this event, as a function of using the conditions of Prop. D i.e., when the embedding x is restricted to lie in the row space of Am. As goes to 0 this converges to the bounds in 44. We formalize this as:
Proposition 2. Let A be an N d matrix whose entries are drawn from i.i.d. standard normal and 0 < c, with
distribution. Let Pt. be any specified set of ratios for < m. Given 0 < 8 < Pmo
y = -Ax:
P(3x E Xpr.... s.t. (ym+1...., YN S. min (I - e)) 2 N- (minge(m) Vi(1 -c) y V
(11)
[E[m]
with the further condition on x that if gm+ S minse(m) v1 E), then EN
Σf=m+1 Pj S 6.
We have the corollary:
Corollary 2. As c gooes to 0, the conditions of Proposition 2 can be met, and:
lim P(3x &pr.....m c00
s.l. (ym+1) UN
min (1(1 ))) 2 $ N-m
(12)
4 [m]
Proof. We create series of values for x of increasing scale. We arrange this so that E pp but as the scale increases, the sum of their probabilities approaches 1. Let x, denote a choice for that is determined by scale factor, in fashion described below. We decompose x, sX,, X, is chosen to produce the desired ratios of probabilities for Pl. Pm. while its scale is fixed by requiring that x, 1. We further define y, A, That is, y, represents the logits that would be produced by the unscaled X,. So the logits for scale 8 are y. = sy.. We denote by P. the probabilities produced by applying softmax y.. We will refer to the i' elements of y., X, and p. as: y. (i). (i) and (i).
Given these definitions, specified ratios of the top-m probabilities, and s, we now show how derive 8, and y, For a given value of s, we have:
(1)
(9)
Ps
(13)
Pm (m)
6
since y. (m) = 1. So we have
log
Pi
= s (1) - 1),
and (i(1)= log
1
Pj Pm
+
1
(14)
Pm
5

---

Published as a conference paper at ICLR 2026
for i m. So, given and the ratios PL we can determine the first values of y.. As goes to infinity, the yz(1) become arbitrarily close while for every value of the yx(i) will produce (i) such that P=(t) Pi, for S m. As noted in the previous section, because Am has full rank, we [m(m) Pm
Pin
can compute x, so that Am x, = y, m)) where y, m) the vector of (1.()).....(mm)))
To summarize, we have defined x, so that the vector x = sx, produces logits that lead to the desired ratio of probabilities for the first m tokens, satisfying the first part of the proposition. As 8 increases the first m elements of y. become more uniform. We will see that the sum of probabilities for the remaining tokens shrinks, provided that their corresponding logits are smaller than the first m logits.
We note that y (i) m (i)) if and only if y, (i) < ym 1, that is, that scaling all the logits does not affect their order. So the probability that the first m logits are larger than all remaining ones, with gap of in the unscaled logits, is given by (II). As goes to infinity, the first m values of y. become arbitrarily close to 1. So in the limit, the probability is given by (4), which applies
when the first m logits equal to 1. This tells us that the probability that the first logits the largest is independent of the ratios of these probabilities for large 8.
We need show that, becomes large, the collective probability of the top-m tokens approaches 1. The value of 8 needed to achieve this depends on the gap between the max Y.( and
1. So we will suppose that all {=( < c, for some small e. As e goes to zero, the probability that this supposition is true will be given by
Suppose we want Em+1 Pj 8, for some small ö. If we have:
Σ N
Čj=m+1 p=(1)
Ey=m+1 (f))
<
(15)
- (1(
Di=1 (v).(1)
this ensures that E)=++1 Pj < 6. We have:
Σ
Σj= jjm+1 evy.
N
Ne)(1-e)
N
c 0
(16)
m i=1 e e
me
m
We can ensure that Equation
15 holds if we choose s so that:
N
which implies
s $
log
(17)
m
Therefore, choosing X, described above, and sufficiently large, ensure that the top-m tokens have probabilities with the desired ratios, while the sum of the remaining ones is arbitrarily small. O
These results show that using a random OPM, if an LLM can specify the top m tokens it can also speciffy their exact probabilities, if they sum close to 1. When using nucleus sampling during inference, this means that for fairly large values m the LLM can specify the exact probabilities of all tokens that might be sampled.
5 Minimal embedding dimension
In this section, instead of showing lower bounds on performance using random OPMs, consider the upper bound on the possible performance of trained OPMs. We ask: for what value m can the best possible OPM specify the top-m next tokens, for any set of m tokens? We formalize this as:
Problem statement: Let m E [N]. What is the minimal embedding dimension d" that allows us to assign maximal probabilities all subsets of tokens of size m??
This tells us the limits of the best possible trained matrix. We start with a simple construction for m that shows that any token be assigned the highest probability when d 2 2.
Proposition 3. For embedding dimension d = 2 and N > 2 there exists matrix, A, of dimension N 2 such that, for any i, there exists a 2d vector x such that af x: > a² x, Vj +

---

Published as a conference paper at ICLR 2026
Proof. Let the rows of A correspond to distinct points on the unit circle, i.., |a,ll = 1. Then if we set x = a, we have a?" x = 1 and a x < for any j }+ i.
We next generalize this result to larger values of m. We begin by showing that this problem is tightly related to sign-rank problem for a matrix that reflects the combinatorial structure of the problem. Given matrix, S, containing values of +1 or 11, the sign rank of S is defined as the minimum rank of any matrix of the same size as S whose elements have the same sign as the corresponding elements of This connection will enable us to determine tight upper and lower bounds on d". For our first
proposition, we define a matrix S E (-1,+1) Nx (x(,,, which we call an (() -inddccatoor matrix. S is defined such that there are exactly m entries in each column with value +1 and N - m entries with value -1 and each column S +11, has positive values corresponding to unique subset Tj C [N] of size [Ty] = m.
Proposition 4. Let S denote an (cr) -indicaatoo matrix. Then,
ld - -ignrank(SS) < 1.
Proof. Consider a collection of (+) logit vectors, Y1, ... Y (()) E RR-V such that each vector yy. a different subset of size m. logits are largest. Arrange these logits in a matrix Y of size N x (M) We call Y a complete top-m logit matrix. Suppose rank(Y) = d, then there exist N d and x (M) matrices A and X such that =AxX, Here, A serves as an output projection matrix and the columns of include the respective embeddings. Consider next the set of all complete top-m logit matrices, denoted )). We therefore ask, what is the minimal rank of Y E y/?
To relate this question to sign rank problem, note that for every logit vector y, in Y, there exists a threshold ty such that the top-m entries are larger than i and the remaining entries are smaller than l. Therefore, y, I is positive only in the top-m entries, where 1, is a vector of N ones. Let
be defined through the outer product = 1/ (1))))()))) then clearly the matrix - T has the sign pattern of S (with the corresponding permutation of the columns). By definition, the minimal rank of matrices with this sign pattern is given by the sign rank of S. Finally, since rank(T) = 1, then d' = minyey rank(Y) cannot differ from the sign rank of S by more than 1.
Our next proposition uses results by Alon et al. [ to determine the sign rank of S exactly.
Proposition 5. Assuming N 2 3m + 1, then signrank(S) = 22 + 1.
Proof. We begin by proving that sigurank(S) S 2m + We will use Lemma 19 from Given a sign matrix 5, let SC(S) denote the maximum number sign changes along any of its columns, and let SC ($) = min SC(M) where the minimum is taken over all matrices M obtained from Š by a permutation of its rows. Then according to this lemma, signrank(S) < SC* ($) + 1. Applying this lemma to our (NN -indicatoor matrix S, we notice that every column of S has up to 2m changes of sign. Moreover, already when N 2 2m + 1, there exists a column with exactly 2m changes of sign, and such a column exists with any permutation of the rows of S. Therefore, SC* (S) 2m, implying that sigurank(S) S 2m + 1.
We next prove that signrank(S) 2 2m + 1. For this, we use the notion of dual sign rank, as defined in I and in Claim 15 therein. To match our definitions, we exchange the roles of rows and columns that claim. We say that a set of rows C is antipodally shattered in sign matrix S if for each vector v e {+1} CI, either v or v appear as a column the restriction of S the rows in C. Claim 15 in establishes that the set of rows C is antipodally shattered in Š if and only if the sign rank of is at least C). The dual sign rank of Š is the cardinality of the maximal such set of rows C.
Returning to our (*) -indicator matrix S, we claim that its dual sign rank is exactly 2m + 1. To see this, notice that with N > 3m + every set of 2m + rows is antipodally shattered, while no set of 2m + 2 rows is antipodally shattered. Specifically, the columns of S include all possible sign configurations with exactly plus signs, and therefore the restriction of S to the set of rows C with |C| = 2m + 1, denoted Sc. includes columns with all possible sign configurations with I m plus signs. The remaining vectors E [+1]2m+1 include > plus signs and <= minus signs. For
7

---

Published as a conference paper at ICLR 2026
these vectors, -V -Sc. Consequently, the union of columns in Se and - include all vectors in [=1]2m-1, implying that C' is antipodally shattered. On the other hand, with |C] >= 2m + 2. there exist vectors in {=1][C] that have both > m plus signs and > m minus signs. Such vectors are neither in Sc nor in -Sc. implying that the dual sign rank of S is 2m + 1. Since the dual sign rank lower bounds the sign rank, we have proved our proposition.
Propositions 4 andd iimmediatelyy imply the following bounds on d".
Corollary 3. Assuming N 2 3m + 1, then 2m S d' 1 2m + 2.
In conclusion, Corollary 3 tells us that for a desired m, there exists an N x d matrix with 2m S d S 2m + that enables assigning the highest probabilities any m desired tokens. Note, however, that a higher embedding dimension may be needed to allow prescribing any desired probabilities to the top-m tokens, or even just to assign them any desired order. Our results show that there exist OPMs that significantly exceed the performance of random matrices analyzed earlier. For example, our lower bounds show that a random OPM in a GPT2-sized model can specify roughly the most likely 26 tokens, while our bounds in this section show that there exists an OPM that can specify the top 383 tokens such a model.
6 Experiments
We perform a number of experiments to provide more precise information about the behavior of output projection matrices In each experiment, select m random tokens, varying m. We then evaluate the ability of the OPM to generate appropriate logits and probabilities for these tokens.
In
We experiment using the trained OPMs for GPT-2 and GP [22] TinyLlama-1. 1a..lammm1.1 IB [29]. T5- Large (23) and 1BB (24), and also with random matrices of the same size as these. GPT-2 has N 50, 257 tokens and embedding size of d 768 and has N = 32, 000, d = 4, 096, while the other models are of intermediate size. We note that the weights of the OPM are tied to the embedding weights in the GPT models and T5, but not in the Llama models.
Our derived lower bound (Section 4.1) assumes each embedding x lies in the row space of Am- Suppose instead we allow
contain component orthogonal to the space spanned by the rows of An. We denote by Xl| the component of x that lies in the row space of Am and by x I the component orthogonal to this space, so that x X|| + x X will not affect the yı, S m. However, will affect y > m, and can be chosen so that it reduces the largest of these logits. Given either a randomly generated or trained and a choice of m tokens (without loss of generality, the first m tokens), we use Linear Programming (LP) to find feasible vector x such that A,, Ym. where ym may be (1,..., 1)" when we verify whether we can specify the desired top-m tokens, or specific values when we want to generate specific probabilities. We further require that for all 1 < S m < af x <x<88 a x.
Finally, we measure the ability of matrices to generate precise probabilities for the top-rn tokens. We generate the desired probabilities from a uniform distribution between 0. and 1.1 (this limits the maximum ratio between the largest and smallest probability to 11). In Appendix D we show similar results using probabilities generated by an LLM. We have shown theoretically that, in the limit, and 5 approach 0, the chances that we can generate precise probabilities should approach the probability that we can generate top-m logits with equal probability. In experiments, we use 0.05, 5 0.1.
In Figure I top), for a GPT-2-sized random matrix, we see (dotted purple curve) that our theoretical bound predicts that can generate any combination of z 26 top tokens with probability near 1. In Appendix B we show empirical results that match this closely. Figure aalso shows that for a TinyLlama-sized OPM with a larger embedding and fewer tokens, the probability is close to I for m s 70 (purple also). We show similar results for three additional models, T5-Large, GPT2-XL and Llama2. Llama2, whose embedding size of 4, 4,096, can generate any m 143 as the top tokens. In we graph this bound for several more models of different sizes.
Figure 2
!Code to reproduce these experiments is available at httpss://gitte com wtt./o//eooarri////o// T'he- S---Sooff-xxottoooo 1 Not Not Probaa-rrbbbbi
8

---

Published as a conference paper at ICLR 2026
64
Fandom 12--abad
0.2
120
04
0.4
0.4
www.amm----aaaaaaaaao Largeraaed Ioser brandt
CP LaL--atc
340
360
180
100
350
200
250
500
TinyLiame
Liamm2-770
mone Ioundd
50
300
350
200
250
300
5500
4.00
430
500
200
400
600
500
1.000
1700
Figure Top: Top-m probability for the GPT-2 and random GPT2-sized (50257 x 768) output projection matrices. Curves show the fraction of times that the matrix satisfied desired constraints for m randomly chosen tokens as m increases. We show our lower bound for a random matrix (4) in dotted purple, along with the trained GPT-2 (solid blue) matrix. Also in this panel, we show LP simulations with prescribed top-m probabilities (random matrices in dotted red and the trained GPT-2 matrix in solid red). Other: Similarly. for T5-Large, GPT-XL, Tiny Liam,, and Llama2 and corresponding random matrices the same size KinyLlama-sized simulations for random matrices (dotted blue) and the trained matrix (solid blue).
simulations for random matrix (dotted blue) and
show lower bound 4 in dotted purple along with linear program
When we incorporate Linear Programming fully optimize x, we find (Figure top)) that GPT-2 (solid blue) and same sized random matrix (dotted blue) both have probability close one of being able to specify m 95 tokens, while TinyLl for example, (Figure [bottom right)) maintains 400 (blue). and Llama2 (bottom right) does this for up we specify exact probabilities for these tokens (Figure[[top)) The largest m with probability close to 1 drops slightly, which is expected given our choice of and å (solid and dotted red). In comparison, using Corollary 3 the best possible trained GPT-2 and Tinyl mnylllam--size matrices can assign the 383 and
this up
1070. Using GPT-2,
highest probability to all configurations of up
1023 tokens, respectively.
In an additional experiment (Figure 3), we show the value of m as a function of d and when LP applied, for random matrices. With < N. we observe that m d/5, and m decreases slowly with N as N becomes large. This is an empirical observation, but it provides simple formula that predicts how large an embedding is needed if one knows how many tokens one needs to represent accurately.
Overall, we conclude that standard OPMs can express fairly large high-probability tuples already at initialization. It seems that training neither improves nor worsens this ability. Many prior works

---

Published as a conference paper at ICLR 2026
Bounds
Model
N
GPT2
50257
768
T5-Large
32128
1024
GPTS-XT.
TMyLLLMMA
Qwen3
152064
3584
113
LLaMAZ-7R
32000
4036
143
DeapSedk-V3
129280
0
200
400
600
830
L1aMA2-708
32000
287
GPT3-1758
50257
12288
418
Figure 2: For the models in the table (right), we show the probability of assigning the highest probabilities to the top-m logits according to the bound (4) (left, graphs are ordered from left to right d increases). For each model, the table the right shows the vocabulary size N, the embedding dimension d, and the value of m where our bound falls below 0.99. These values for more models are shown in Appendix C
Maximal m(d), N=32000
Maximal m(N), d-1000
500
1000
400
300
mtid)
500
200
100
500
1000
1500
2000
5000
10000 15000 20000 25000 30000
N
Figure 3: Top-m achieved by random matrices of size 32000 x d with 100 < d < 2050 (left) and matrices of size N 1000 with 2000 N S 32000 (right). The solid blue curves show the values of m for which 0.99 probability is achieved, computed via linear- programming simulation. Our bound shown by the dotted red curves.
have analyzed the statistical regularities in embeddings and OPMs that emerge through training. In Appendix we explicitly compare first and second-order properties of trained and random matrices.
Conclusions
The softmax bottleneck means that LLMs can only represent a measure-zero subset of the set of all probability distributions for the next token. However, have shown, both theoretically and empirically, that reasonably sized models can represent the exact probabilities of a large number of the most probable tokens. This calls into question the degree which the softmax bottleneck imposes limitations on LLMs in most realistic settings.
Acknowledgements
RB was supported in part by the Israeli Council for Higher Education (CHE) via the Weizmann Data Science Research Center, by the MBZUAI-WIS Joint Program for Artificial Intelligence Research. and by research grants from the Estates of Tully and Michele Plesser and the Anita James Rosen and Harry Schutzman Foundations.
was supported by a National Science Foundation (NSF) grant #2213335. We thank an anonymous reviewer for suggesting the simplified proof of Proposition used in this paper.
10

---

Published as a conference paper at ICLR 2026
References
[1] Noga Alon, Shay Moran, and Amir Yehudayoff. "Sign rank versus VC dimension". In: Annual Conference Computational Learning Theory. 2015.
[2] Haw-Shiuan Chang and Andrew McCallum. "Softmax bottleneck makes language models unable to represent multi-mode word distributions". In: Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics. Vol. 1. 2022.
[3] David Demeter, Gregory Kimmel, and Doug Downey. "Stolen Probability: A Structural Weakness of Neural Language Models". In: Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics. Ed. by Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel Tetreault. Online: Association for Computational Linguistics, July 2020, pp. 2191- 2197.
[4] Kawin Ethayarajh, David Duvenaud, and Graeme Hirst. "Towards understanding linear word analogies". In: Proceedings of the 57th annual meeting of the Association for Computational Linguistics. 2019, pp. 3253-3262.
[5] Markus Freitag and Yaser Al-Onaizan. "Beam Search Strategies for Neural Machine Transla- tion". In: Proceedings of the First Workshop Neural Machine Translation. Ed. by Thang Luong, Alexandra Birch, Graham Neubig, and Andrew Finch. Vancouver: Association for Computational Linguistics, Aug. 2017, pp. 56-60.
[6] Octavian Ganea, Sylvain Gelly, Gary Bécigneul, and Aliaksei Severyn. "Breaking the softmax bottleneck via learnable monotonic pointwise non-lincaritiee". In: International Conference on Machine Learning. PMLR. 2019, pp. 2073-2082.
[7] Nathan Godey, fiii Villemonte de la Clergerie, and Benoit Sagot. "Why do small language models underperform? Swdying Language Model Saturation via the Softmax Bottleneck". In: First Conference on Language Modeling. 2024.
[8]
Alex Graves. Sequence Transduction with Recurrent Neural Networks. 2012. arXiv: 1211
3711
[cS. NE URL:
[9]
Andreas Grivas, Nikolay Bogoychev, and Adam Lopez. "Low-Rank Softmax Can Have Unargmaxable Classes in Theory but Rarely in Practice". In: Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pupers). Ed. by Smaranda Muresan, Preslav Nakov, and Aline Villavicencio. Dublin, Ireland: Association for Computational Linguistics, 2022, pp. 6738-673 7738-6758.
[10] Michael Hahn. "Thcoretical Limitations of Self-Attentio in Neural Sequence Models". In: Transactions of the Association for Computational Linguistics 8 (2020). Ed. by Mark Johnson, Brian Roark, and Ani Nenkova, pp. 156-171.
[11] Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. "The Curious Case of Neural Text Degeneration". In: International Conference on Learning Representattinss 2020.
[12] Tokio Kajitsuka and Issei Sato. " "Ae Transformers with One Layer Self-Attention Using Low- Rank Weight Matrices Universal Approximators?" In: The Twelfih International Conference on Learning Representations. 2024.
[13] Sekitoshi Kanai, Yasuhiro Fujiwara, Yuki Yamanaka, and Shuichi Adachi, "Sigsoftmax: Re- analysis of the softmax bottleneck". In: Advances in Neural Information Processing Systems
31 (2018).
[14] Goro Kobayashi. Tatsuki Kuribayashi, Sho Yokoi, and Kentaro Inui. "Transformer Language Models Handle Word Frequency in Prediction Head". In: Findings of the Association for Computational Lingnistics: ACL 2023. Ed. by Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki. Toronto, Canada: Association for Computational Linguistics, July 2023, pp. 4523- 4535.
[15] Su Hyeong Lee, Qingqi Zhang, and Risi Kondor. "Sign Rank Limitations for Attention-B.ao Bassed Graph Decoders". In: CoRR (2024).
[16] Valerii Likhosherstov. Krzysztof Choromanski, and Adrian Weller. "On the Expressive Flex- ibility of Self-Attention Matrices". In: Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 37. 7. 2023, pp. 8773-8781.
[17] Kanti V. Mardia, John T. Kent, and John M. Bibby. Multivariate Analysis. London: Academic Press, 1979.
11

---

Published as a conference paper at ICLR 2026
[18] William Merrill and Ashish Sabharwal. "The Expressive Power of Transformers with Chain of Thought". In: International Conference on Learning Representations. 2024.
[19] Johannes von Oswald, Niklas Niklasson, Ettore Murabito, Felix Hill, Timothy P Lillicrap, James Ritchie, and Chrisantha Fernando. "Transformers learn in-context by gradient descent".
In: International Conference on Learning Representations. 2023.
[20] Dwarak Govind Parthiban, Yongyi Mao, and Diana Inkpen. "On the Softmax bottleneck of recurrent language models". In: Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 35. 15. 2021, pp. 13640-153 13607.
[21] Jorge Pérez, Pablo Barceló, and Javier Marinkovic. " "Attention is of Machine Learning Research 22 (2021), pp. 1-35.
In: Journal
[22] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. "Lan- guage Models are Unsupervised Multitask Learners.. In: (2019).
[23] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. "Exploring the limits transfer learning with a unified text-to-text transformer". In: Journal of machine learning research 21. 140 (2020), pp. 1-67.
[24] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, et al. "Llama 2: Open Foundation and Fine-Tuned Chat Models". In: arXiv preprint arXiv:2307. 239809987888888 (2023). [25] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. " "Attention is all you need", In: Advances neural information processing systems 30 (2017).
[26] Zhilin Yang, Zihang Dai, Ruslan Salakhutdinov. and William W. Cohen. "Breaking the Softmax Bottleneck: Riggh-Rank RNN Language Model". In: 6th International Conference Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 MMy 3, 2018, Conference Track Proceedings. OpenReview.net, 2018.
[27] Chulhee Yun, Srinadh Bhojanapalli, Sanjiv Reddi, and Sanjiv Kumar. "Are transformers universal approximators of sequence-to-sequence functions?" In: International Conference on Machine Learning. PMLR. PMLR, 2019, pp. 6776-6786.
[28] Chulhee Yun, Soomin Kim, Jaewoo Lee, Hyoungjin Lee, Do-Kyum Kim, and Jeongwook Kim. "O(n) Connections are Expressive Enough: Universal Approximability of Sparse Trans- formers". In: Advances in Neural Information Processing Systems. Vol. 33. 2020, pp. 10166- 10178.
[29] Peiyuan Zhang, Guangtao Zeng, Wenbin Ma, Weizhong Tu, Jing Ma, Yizhi Li, Shiyi Li, Mingqi Gao, and Wei Lu. TinyLlama: An Open-Source Small Language Model, 2024, arXiv: 2401. 02385
[cs. C
12

---

Published as a conference paper at ICLR 2026
Appendix
A
Cumulative probabilities of most likely tokens
Figure 4 shows the average cumulative probabilities of the m likely tokens in ten stories from the New York Times computed using GPT2. Roughly 200 tokens are needed to predict the next token with 0.9 probability. threshold of 0.9 commonly used for nucleus sampling. We note that larger models such as Llama2-7B can empirically produce precise probabilities for about m 1000 tokens.
Cumulative Probebilities
Probobiity m most likely tokens
0.85
0.8
0.75
0
200
400
800 800
1000
Figure 4: For different values of m we determine the average cumulative probabilities of the m most likely tokens. This is averaged over all embeddings produced by ten stories from the New York Times, across different subject areas, Probabilities are produced by analyzing these stories using GPT2. Error bars show the standard deviation of average values produced the ten different stories. Note that we have previously shown that even small models, such as Llama2 can empirically produce precise probabilities for about 1000 tokens.
B Independence assumption
Figure compares the bound derived in Equation 4 to numerical simulation. The probabilities obtained with our bound are slightly smaller than those observed in the simulation due to the statistical dependence of the factors Equation (6).
C Bounds for additional models
We include the lower bound results in Equation 4 for additional models in Table 1
Approximation quality
5 Simulation
Wishart
20 40
60
100
120
Figure 5: For a random matrix of the size of the GPT2 OPM (50257 768), the figure shows top-m probability for the optimal embedding the row space of Am obtained with simulation (solid blue) compared to inverse- inverse-Wishart simulation. (A 0.99 probability is achieved at m = 27 with the closed form expression compared to 28 with the simulation.)
bound (dashed purple). Our bound slightly undershoots the
13

---

Published as a conference paper at ICLR 2026
Model
N
d
GPT2
50257
768
26
BERT-Basc
30522
768
26
Base
32128
768
XLNet-Bass
32000
768
26
BERT-Large
30522
1024
35
T5-Large
32128
1024
35
XLNet-Large
32000
1024
GPT2-XL
50257
1600
Olmo2-1B
100352
2048
66
Gamma2-2B
256000
2304
70
TinyLLaMA
32000
2048
71
Phi3-Mini
32064
3072
107
Qwen3
152064
3584
113
LLaMA3-8B 8B
128256
4096
131
MPT-7B
50176
4096
139
LJaMA2-7B
32000
4096
143
Falcon-40B
65024
4544
152
LLaMA2-133
32000
5120
179
DeepSeek-V3
129280
7168
230
Falcon-40B
65024
8192
274
LLaMA2-70B
32000
8192
287
GPT3-175B 50257
12288
418
Table 1:
This table lists number of models, along with their vocabulary size and embedding size. We use our theoretical results to provide a lower bound on the probability that, for a given m, random model of this size designate any m tokens the most likely ones. The last column indicates the value of m for which this probability drops below 0.99.
D Using natural language probabilities
Our last experiment compares using the uniform distribution set the target probabilities in Figure T to probabilities drawn from the distribution of natural language. For this experiment, we applied GPT2 contexts from ten New York Times articles, as in Appendix A Then, for specific choice of m, randomly selected token and extracted the top probability values for the next token after We then used these probabilities. Note that for the LP, assign the extracted probabilities to random tokens. Figure 6 shows that the bounds with these probabilities and with probabilities drawn from uniform distribution are almost identical.
determine whether a model could find an embedding that would predict
E Statistical properties of OPM rows
Many prior works have analyzed the statistical properties of embeddings and OPMs (c.g.. In Figure ? we explicitly compare the statistics of trained and random matrices.
3
14
14

---

Published as a conference paper at ICLR 2026
GPT2-probabilitts from natural language
0.5
G TT: Probabilitiss from NYT
*- TT: Uniform probabilities
20
40
60
80
100
120
Figure 6: In Figure we showed that GPT2 can produce precise probabilities for m tokens, with up about 85 (red curve). These probabilities were drawn from uniform distribution. In this figure, we add a second version of the experiment (blue curve) that uses probability distributions of next-token probabilities generated by GPT -2 on news articles. We can there is little difference.
GPT2
TinyLlama
4000
8000
landom
Random Trained
Trained
6000
2000
4000
1000
2000
0.5
1.5
Row Magnitude
Row Magnitude TinyLlama
GPT2
2000
1000
Random
Trained
1500
Trained
500
1000
500
-1
-0.5
0.5
-0.5
0.5
Cosine Similarity of Rows
Cosine Similarity of Rows
Figure 7: Examining first and second order statistics between rows, we do see significant difference between random and trained models. In each figure, we compare trained model to random model of the same size. In the top row of the figure, we plot a histogram of the magnitudes of the rows. The random matrix is scaled so that its rows have the same mean magnitude as the trained model. We that the trained models have much greater variance in their magnitude. In the bottom row, compare the cosine similarity of randomly chosen rows. While the random models have mean-zero cosine similarity, the trained models have almost all positive values, with greater variance. 3 have previously described similar empirical results. Note that GPT-2 is trained with tied weights between the encoding and unencoding matrices, whereas TinyLlama uses untied weights.
15


Reviews:
[{"review_id": "fQGbjfiRTY", "reviewer": "Reviewer_kpFi", "summary": "The paper investigates both theoretically and empirically the impact of the softmax bottleneck on the ability of large language models to correctly represent the probabilities of the m most probable tokens. Their conclusion is that the softmax bottleneck does not seem to \"provide any limitations to LLMs in most realistic settings\".", "questions": "Suggestions rather than questions:\n- Please double-check that the papers you cite as arXiv papers have not been published (by that I mean \"really\" published, in the proceedings of a conference or in a journal). Whenever it is the case, the proper publication should be cited, not the arXiv pre-print. There are multiple cases of this issue in your bibliography.\n- Although the paper is generally well written, please refrain from using abbreviations such as \"WLOG\" (used twice).", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- An interesting contribution to the \"softmax bottleneck\" literature that goes beyond previous work (notably Demeter et al. 2020 and Grivas et al. 2022).\n- The paper combines theoretical analyses and empirical results.\n- The paper is generally clear and well written.", "weaknesses": "- The motivation for the research question is not convincing enough. Why is it so important that models can assign the correct probabilities, summing to (almost) 1, to the m best tokens? If it is possible (and it is for fairly large m's), what does it tell us about language models? If it hadn't been possible, why would that have been an issue, beyond very niche scenarios such as choosing a number or a US state at random?\n- The discussion in section 6 could be more detailed. In particular, how do your results complement previously published papers that found that the softmax bottleneck *is* an issue for language models (e.g. Parthiban et al. 2021, Godey et al. 2024)?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "cgRmflaXU6", "reviewer": "Reviewer_SQX5", "summary": "This paper tackles the softmax bottleneck problem, i.e. the study of the expressivity of usual neural language models that use a hidden dimension that is smaller than the vocabulary size $N$. In that setup, it has been shown that there exist probability distributions in $\\Delta^N$ that cannot be predicted. In this paper, the authors state that language model outputs can match probability distributions with a relatively large support, either by predicting the token set that belong in that support, or even by predicting exactly the probabilities for the tokens in such supports. As a result, they argue that the softmax bottleneck is not a dramatic issue for language models, which they assess through a short experimental analysis.", "questions": "- Could you report the results that show independence for the rows of $(A_mA_m^T)^{-1}$?\n- By my understanding, Proposition 3 could be proven more easily by setting all a_j at distinct positions on the unit sphere in 2d and setting x_j = a_j?\n- Do OPMs actually train to account for larger m values? Is it possible that the learned A_m do not have the used invertibility properties?\n- To what extent does token rank and specified ratios would affect the results in section 5?", "limitations": "", "rating": 4, "confidence": 5, "soundness": 1, "presentation": 3, "contribution": 4, "strengths": "This paper studies a very interesting topic and conducts a novel and relevant theoretical analysis.\n- **Theoretical results**: The propositions presented in this paper are novel and provide a more profound understanding of this issue. They shed light on the complexity of reachable distributions, and probabilize (a part of) the set of reachable distributions, which is both a challenging and exciting outcome. Even though the proof in section 4.1 is not perfectly rigorous, it makes very reasonable assumptions and uses the Inverse Wishart distribution in this context, which is a new and very insightful contribution in the context of random matrices for the softmax bottleneck problem. I appreciate that the authors extended their work to fixed top-m probability. Proposition 4 is very elegant. The use of results on the signrank is also very novel and could open new venues for this topic. It leads to a clean result on the minimal dimensions needed to accurately match the supports of probability distributions.", "weaknesses": "Although I deeply appreciate the core theoretical results of this work, I quite disagree with some of the claims and conclusions made in the abstract, introduction, and empirical sections. I also believe that this paper could be enhanced by deeper empirical experiments.\n- **Claims about broader implications**: It is mentioned in the abstract and introduction that \"the softmax bottleneck does not significantly limit the capabilities of LLMs\", or that \"limitations to the expressiveness of transformers are not really that significant\". My understanding of what is proven in the paper is 1- for a single given target probability vector and a random OPM matrix, there exists an $x$ for which the predicted probability matches the top-$m$ target probabilities; 2- if one only cares about the support but wants to match any set of targets on the top-$m$ ranking of tokens, then in most setups $d \\approx m/2$ is sufficient. Hence, it is unclear whether matching any set of targets for the correct probabilities (let alone the order of such probabilities) is possible. This is a crucial distinction, as the complexity of matching any permutation of token order is much higher than matching the support as a set. Moreover, even when there would exist output representations that would give the desired probabilities, they might be configured in ways that are very difficult to reach during training, which would limit the applicability of these results to usual training setups. What can thus be concluded from this paper is that the softmax bottleneck phenomenon does not strongly limit the prediction of the next-token probability supports, and that individual \"low-entropy\" probability vectors should be truthfully matched with non-negligible probability.\n- **Lack of clarity in some proofs**: The proof of Proposition 2 is a bit hard to follow and could be made clearer. The first paragraph is a bit confusing and it seems like it could be summarized to convey the idea more directly. Moreover, the statistical independence of the entries of $(A_mA_m^T)^{-1}$ is never verified empirically in the paper.\n- **Experimental design**: The experimental section is less appealing than the theoretical section. Figure 1 verifies the lower bound given in Propositions 1 and 2 (and incidentally shows that it is not particularly tight). Figure 2 computes the bound curve for several setups. Figure 3 explores a question that seems loosely related to the topic at hand, and that was covered many times by the anisotropy literature (see works of Ethayarajh et al., among others). A lot of questions remain unanswered, from the necessary experiments to extensions that would have been relevant: given an actual token distribution taken from natural language, how much of the probability supports can a random\\trained matrix cover? What are the types of x and s that are observed? In a synthetic controlled setup, e.g. data generated from a bigram with known supports, can a model be trained to properly recover the supports up to $m$ tokens? In cases where the supports are of different sizes, ie where m should be different across target probabilities, what types of solutions are found and with what performance? The experimental section also ignores the second theoretical section of the paper.\n\nIn the current state of the paper, the claims and conclusions that are made about the relevance of the softmax bottleneck are not strongly supported by the theoretical section and the experiments. It is particularly important as this article is quite technical, which implies that readers may take these claims for granted without thoroughly reading the paper and understanding its potential limitations. Hence, it is crucial that the claims accurately reflect the results that are presented to avoid any misinterpretation.\n\nOverall, I am excited by the topic and theoretical part (which I would rate 8/10), but I am underwhelmed by the conclusions and experimental sections (which I would rate 2/10).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "7iTqkvmNUV", "reviewer": "Reviewer_B1qY", "summary": "The work addresses the question of the softmax bottleneck in transformers -- that is, the output projection limiting the production of arbitrary probability distributions. The authors ask whether the softmax bottleneck restricts the LLM from representing the probabilities of the top-m most probable tokens, arguing that the exact probabilities of unlikely tokens are less important. The authors provide theoretical results showing that there exist OPMs that can represent any specific probabilities over top-m tokens for large m and empirical results on GPT-2 and TinyLlama.", "questions": "- What is the role of weight tying here? \n- Are there settings (e.g., RL, exploration, or calibration tasks) where representing low-probability tokens accurately would matter, and how would your framework apply there? Do the values of m you found seem sufficient?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Extends on prior work on softmax bottleneck, reframing the problem around top-m probabilities\n- The research question is clear as well as the writing\n- The derivations give nice lower bounds for random matrices and the theoretical insights are supported by experiments on GPT-2 and TinyLlama.", "weaknesses": "- The results show the existence of embeddings that can realize given top-m probabilities, but does not address whether these are learned by real transformers\n- Assumes low-probability tokens are irrelevant, but there may be domains (e.g., RL fine-tuning or exploration) where coverage over rare tokens is important", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "ufFE7CS9Di", "reviewer": "Reviewer_3zJk", "summary": "This paper offers a new look at the so-called \"Softmax Bottleneck,\" which is said to limit the ability of a model as it can struggle to go from inner embeddings to logits over the whole vocabulary: the authors claim that what matters in practice is to predict accurately the probabilities for top tokens, especially given that sampling is usually done from a truncated distribution. The paper first gives theoretical lower bounds showing that indeed any chosen set of m tokens can be the most likely ones, with high probability, and even reach, collectively, a specific probability. It then validates this on GPT‑2 and TinyLlama: in practice this works for somewhat surprisingly larges m.", "questions": "repeating what I see as a weakness:\n- what do you think of the reachability question? Your results do establish existence of an embedding $x$ for a given $A$ but I wonder whether a network can be trained to produce that $x$...\n- are we sure that same (trained) network would produce the valid embeddings, whatever the context? This seems important to fully grasp the extent of your work.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- this is a really inspiring take on the softmax bottleneck, as indeed given the prevalence of truncated sampling, what matters is often the ability to produce valid probabilities for the top tokens of a vocabulary ;\n- the authors show the theoretical possibility to produce the embeddings that would result in a given m-subset getting the most probabilities, for  a randomly initialized projection matrix, with some bounds, and even to reach a given probabilities for this subset ;\n- the experiments validate their theory for a trained GPT-2 and TinyLlama, far exceeding the lower bounds ;\n- the paper is convincingly written, with a well formulated problem, sound proofs and related experiments.", "weaknesses": "- the proofs seem sound but this is really to the best of my understanding and I had to trust the authors for this theoretical work. Similarly, it's not clear what to make of such a sentence \"since this derivation involves some approximation, we have empirically confirmed that simulations match our theoretical predictions (not shown)\";\n- although this work proves the existence of \"embedding that the OPM will map into the appropriate probability distribution\", there is no guarantee that the model will produce it ;\n- furthermore, although this is interesting, it does not say anything about the ability of a given network to produce the embeddings for \"desired\" m-subsets in many contexts. This is not the point of this work, but it does take a practical look at the softmax bottleneck, so my _practical_ concern might be warranted...", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)