The Softmax Bottleneck in Large Language Models: Mathematical Foundations, Matrix Rank Limits, and Mixture of Softmaxes

title: "The Softmax Bottleneck in Large Language Models: Mathematical Foundations, Matrix Rank Limits, and Mixture of Softmaxes" slug: "the-softmax-bottleneck-in-large-language-models-mathematical-foundations-matrix-rank-limits-and-mixture-of-softmaxes" status: "published" feature_image: "https://cms.llms.blog/content/images/2026/08/softmax-bottleneck-cover-1.png" excerpt: "A standard linear projection followed by Softmax caps the rank of predicted log-probability distributions to the hidden dimension plus one, creating an architectural ceiling on natural language expressivity." tags:

• Explainers
• Architecture
• Deep Learning
• LLMs

In modern neural language models, next-token generation is parameterized by mapping a contextual hidden state vector $h_c \in \mathbb{R}^d$ through a linear projection matrix $W \in \mathbb{R}^{V \times d}$ and normalizing the resulting logits with a softmax function. While this formulation is standard across autoregressive Transformers, it introduces an intrinsic mathematical constraint known as the softmax bottleneck.

First formalized by Yang et al. (ICLR 2018), the softmax bottleneck establishes that standard softmax parameterization is equivalent to a low-rank matrix factorization. Consequently, the rank of the predicted log-probability matrix across all contexts and vocabulary tokens is strictly bounded by $d + 1$, where $d$ is the hidden embedding dimension. When the true conditional probability distribution of natural language requires a higher effective rank to capture complex polysemy, syntactic branching, and domain nuance, standard softmax models provably fail to represent the target distribution, regardless of network depth.

Language Modeling as Matrix Factorization

To understand the origin of the bottleneck, language modeling can be formulated as approximating a ground-truth matrix of conditional probabilities.

Let $\mathcal{C}$ denote the set of all possible context sequences, and let $\mathcal{V}$ denote the discrete vocabulary of size $V = |\mathcal{V}|$. The true conditional language distribution defines a log-probability matrix $A^* \in \mathbb{R}^{|\mathcal{C}| \times V}$, where each entry is:

$$A^<em>_{c, v} = \log P^</em>(v \mid c)$$

In a neural language model with hidden dimension $d$:

1. The encoder or Transformer backbone produces a contextual hidden state vector $h_c \in \mathbb{R}^d$ for context $c$.
2. The output embedding (unembedding) matrix $W \in \mathbb{R}^{V \times d}$ associates each token $v$ with a vector $w_v \in \mathbb{R}^d$.
3. The unnormalized logit for token $v$ is computed as the inner product $z_{c, v} = h_c^T w_v$.
4. The softmax function converts logits into a normalized probability distribution:

$$P_\theta(v \mid c) = \frac{\exp(h_c^T w_v)}{\sum_{v'=1}^V \exp(h_c^T w_{v'})}$$

Taking the logarithm of the predicted probability yields:

$$\log P_\theta(v \mid c) = h_c^T w_v - \log \sum_{v'=1}^V \exp(h_c^T w_{v'})$$

In matrix notation, let $H \in \mathbb{R}^{|\mathcal{C}| \times d}$ denote the matrix of hidden states across all contexts, and let $Z = H W^T \in \mathbb{R}^{|\mathcal{C}| \times V}$ denote the logit matrix. Defining $g \in \mathbb{R}^{|\mathcal{C}|}$ as the context-dependent log-sum-exp normalization vector where $g_c = \log \sum_{v'=1}^V \exp(h_c^T w_{v'})$, the predicted log-probability matrix $A_\theta \in \mathbb{R}^{|\mathcal{C}| \times V}$ is expressed as:

$$A_\theta = Z - g \mathbf{1}_V^T = H W^T - g \mathbf{1}_V^T$$

where $\mathbf{1}_V \in \mathbb{R}^V$ is an all-ones column vector.

+-----------------------------------------------------------------------------------+
|               MATRIX FACTORIZATION FORMULATION OF SOFTMAX OUTPUT                  |
|                                                                                   |
|    Context Hidden States (H)      Output Weights (W^T)       Logit Matrix (Z)     |
|      [ |C| x d ]                    [ d x V ]                 [ |C| x V ]         |
|   +-----------------+           +-----------------+       +-----------------+     |
|   |  h_1  h_2  h_3  |     *     |  w_1  w_2  w_3  |   =   |  z_11 z_12 z_13 |     |
|   |  .    .    .    |           |  .    .    .    |       |  .    .    .    |     |
|   |  h_|C| ...      |           |  ...            |       |  z_|C|1 ...     |     |
|   +-----------------+           +-----------------+       +-----------------+     |
|         rank <= d                      rank <= d                 rank <= d        |
|                                                                                   |
|    Predicted Log-Probability Matrix: A = Z - g * 1_V^T                            |
|    Linear Algebra Bound: rank(A) <= rank(Z) + rank(g * 1_V^T) <= d + 1           |
+-----------------------------------------------------------------------------------+

The Rank Bound Theorem and Its Proof

The central result established by Yang et al. (2018) is that the predicted matrix $A_\theta$ has a rank strictly bounded by $d + 1$.

Theorem (Softmax Bottleneck Rank Limit)

For any neural language model parameterized by hidden dimension $d$, vocabulary size $V$, and context space size $|\mathcal{C}|$, the rank of the model's log-probability matrix satisfies:

$$\text{rank}(A_\theta) \le d + 1$$

Proof:

1. The logit matrix $Z$ is the product of $H \in \mathbb{R}^{|\mathcal{C}| \times d}$ and $W^T \in \mathbb{R}^{d \times V}$.
2. By subadditivity of matrix rank under multiplication:

$$\text{rank}(Z) \le \min(\text{rank}(H), \text{rank}(W^T)) \le d$$

1. The normalizer term $g \mathbf{1}_V^T$ is an outer product of two vectors ($g \in \mathbb{R}^{|\mathcal{C}|}$ and $\mathbf{1}_V \in \mathbb{R}^V$), which is a rank-1 matrix:

$$\text{rank}(g \mathbf{1}_V^T) \le 1$$

1. Applying the subadditivity of matrix rank under addition:

$$\text{rank}(A_\theta) = \text{rank}(Z - g \mathbf{1}_V^T) \le \text{rank}(Z) + \text{rank}(g \mathbf{1}_V^T) \le d + 1$$

When the true log-probability matrix $A^$ has rank $r^ > d + 1$, there exists no parameter configuration $\theta = (H, W)$ that can exactly reproduce the true conditional distributions $P^*(v \mid c)$ for all contexts.

Geometric Limitations and Convex Hull Enclosures

Beyond matrix rank limits, the softmax parameterization enforces rigid geometric constraints on the output probability simplex.

As demonstrated by Demeter et al. (ACL 2022) in their geometric analysis of the softmax layer:

• Let $\mathcal{H}_W = \text{conv}(\{w_1, w_2, \dots, w_V\})$ denote the convex hull of the word embedding vectors in $\mathbb{R}^d$.
• If a token's embedding vector $w_i$ lies in the interior of the convex hull $\mathcal{H}_W$, the maximum probability that standard softmax can assign to token $w_i$ across any context vector $h \in \mathbb{R}^d$ is strictly bounded by the probability assigned to at least one token whose vector lies on the boundary or vertices of $\mathcal{H}_W$.

This geometric restriction prevents standard models from assigning high contextual probabilities to specific multi-facet words, rare tokens, or fine-grained domain terms when their embedding representations are pulled toward the interior of the representation space by general semantic similarity objectives.

Mixture of Softmaxes (MoS): Breaking the Rank Ceiling

To circumvent the low-rank restriction without exploding parameter counts across the entire network, Yang et al. (2018) introduced the Mixture of Softmaxes (MoS).

MoS formulates next-token generation as a latent-variable mixture model with $K$ components (or facets). Instead of predicting a single context state $h_c$, the model computes $K$ distinct context vectors $h_{c, k} \in \mathbb{R}^d$ and a set of dynamic mixture weights $\pi_{c, k}$:

$$\pi_{c, k} = \frac{\exp(h_c^T w_{\pi, k})}{\sum_{j=1}^K \exp(h_c^T w_{\pi, j})}$$

The conditional token probability is then computed as a convex combination of $K$ individual softmax distributions:

$$P_{\text{MoS}}(v \mid c) = \sum_{k=1}^K \pi_{c, k} \frac{\exp(h_{c, k}^T w_v)}{\sum_{v'=1}^V \exp(h_{c, k}^T w_{v'})}$$

+-----------------------------------------------------------------------------------+
|                     MIXTURE OF SOFTMAXES (MoS) PIPELINE                          |
|                                                                                   |
|                                Context State h_c                                  |
|                                        |                                          |
|                 +----------------------+----------------------+                    |
|                 |                      |                      |                   |
|                 v                      v                      v                   |
|           Facet Head 1           Facet Head 2           Facet Head K              |
|             h_{c, 1}               h_{c, 2}               h_{c, K}                |
|                 |                      |                      |                   |
|                 v                      v                      v                   |
|           Softmax Dist 1         Softmax Dist 2         Softmax Dist K            |
|             P_1(v|c)               P_2(v|c)               P_K(v|c)                |
|                 \                      |                      /                   |
|                  \                     |                     /                    |
|                   +-----------------> (X) <-----------------+                     |
|                                        ^                                          |
|                                        | Gating Weights                           |
|                                    [pi_{c, k}]                                    |
|                                        |                                          |
|                                        v                                          |
|                              Output Distribution                                  |
|                         P(v|c) = SUM pi_k * P_k(v|c)                              |
+-----------------------------------------------------------------------------------+

Why MoS Overcomes the Bottleneck

Taking the logarithm of the MoS output probability yields:

$$\log P_{\text{MoS}}(v \mid c) = \log \left( \sum_{k=1}^K \pi_{c, k} \frac{\exp(h_{c, k}^T w_v)}{\sum_{v'=1}^V \exp(h_{c, k}^T w_{v'})} \right)$$

Because the logarithm of a weighted sum of exponentials is a non-linear operation, the resulting log-probability matrix can no longer be decomposed into $H W^T - g \mathbf{1}^T$. By introducing $K$ components, the theoretical rank of the predicted matrix expands to $\mathcal{O}(K \cdot d)$, allowing the model to fit high-rank conditional distributions.

Empirically, on standard language modeling benchmarks (Penn Treebank and WikiText-2), MoS achieved state-of-the-art perplexity reductions while demonstrating singular value spectra with significantly higher effective ranks than standard softmax baselines.

Trade-Offs and Modern LLM Engineering

While Mixture of Softmaxes resolved the theoretical low-rank bottleneck, frontier autoregressive Transformers (such as Llama 3, Gemma 2, and Qwen 2.5) generally do not implement MoS in production. The reasons stem from computational efficiency, vocabulary scaling, and architectural choices.

| Architecture Choice | Theoretical Rank Bound | Output Projection FLOPs | GPU Memory Overhead | Production Adoption | | :--- | :--- | :--- | :--- | :--- | | Standard Linear Softmax | $d + 1$ | $\mathcal{O}(B \cdot S \cdot d \cdot V)$ | Baseline | Universal standard in frontier LLMs | | Mixture of Softmaxes (MoS) | $\mathcal{O}(K \cdot d)$ | $\mathcal{O}(K \cdot B \cdot S \cdot d \cdot V)$ | $K \times$ logit tensor memory | Rare in large LLMs due to memory bandwidth | | Direct Output Connection (DOC) | $\mathcal{O}(K \cdot d)$ | $\mathcal{O}(K \cdot B \cdot S \cdot d \cdot V)$ | Multi-layer hidden concatenation | Specialized compact models | | Multi-Token Prediction (MTP) | Independent rank-$d$ heads | $\mathcal{O}(M \cdot B \cdot S \cdot d \cdot V)$ | Parallel head buffers | DeepSeek-V3, specialized code models |

1. Output FLOPs and Memory Bandwidth at Scale

In modern LLMs, vocabulary sizes range from $V = 128{,}000$ to $V = 256{,}000$. Computing the final unembedding projection is memory-bandwidth intensive. In MoS, computing $K$ separate logit matrices requires multiplying the final GEMM FLOPs and output memory footprint by $K$. In serving runtimes where decoding is strictly memory-bandwidth bound, multiplying unembedding memory traffic by $K = 3$ or $K = 4$ introduces substantial latency penalties.

2. Hidden Dimension ($d$) Scaling

When Yang et al. published MoS in 2018, standard recurrent models operated with $d \in [512, 1024]$. At $d = 512$, the rank ceiling of 513 posed an immediate bottleneck against complex linguistic distributions.

Modern frontier LLMs scale hidden dimensions significantly:

• 8B models: $d = 4096$
• 70B models: $d = 8192$
• 400B+ models: $d = 12288$ to $16384$

With $d \ge 4096$, the linear rank capacity of standard softmax is often sufficient to capture next-token conditional distributions without hitting severe bottleneck limits during pre-training.

3. Impact on Small Language Models (SLMs)

Recent empirical studies, such as Godey et al. (2024), confirm that the softmax bottleneck remains a primary factor in the performance degradation and representation saturation of compact models ($< 1\text{B}$ parameters with $d \le 1024$). For resource-constrained edge deployments, techniques that expand effective output rank (such as low-rank bottleneck bypasses, output facet routing, or non-linear prediction heads) remain active areas of optimization.

Summary and Core Takeaways

1. Rank Bound: A standard linear output head followed by Softmax caps the rank of the predicted log-probability matrix to $d + 1$, where $d$ is the model's hidden dimension.
2. Expressivity Limit: If the true context-conditional language distribution has an intrinsic rank $r^* > d + 1$, standard softmax models cannot represent the distribution, regardless of network depth.
3. Geometric Constriction: Dot-product softmax prevents interior tokens in the vocabulary embedding convex hull from receiving maximal probability over boundary tokens.
4. Mixture of Softmaxes: MoS breaks the linear rank ceiling by expressing output probabilities as a convex combination of $K$ softmax distributions, scaling effective rank to $\mathcal{O}(K \cdot d)$.
5. Modern Engineering Trade-Off: In frontier LLMs, scaling $d \ge 4096$ provides sufficient rank capacity while avoiding the $K \times$ compute and VRAM bandwidth tax of multi-softmax output layers over large vocabularies ($V \ge 128\text{k}$).

Sources

• Yang et al. (ICLR 2018): Breaking the Softmax Bottleneck: A High-Rank RNN Language Model
• Demeter et al. (ACL 2022): Softmax Bottleneck Makes Language Models Unable to Represent Multi-facet Words
• Godey et al. (2024): Why Do Small Language Models Underperform? Studying LM Saturation via the Softmax Bottleneck
• Takase et al. (EMNLP 2018): Direct Output Connection for a High-Rank Language Model
• Kanai et al. (AAAI 2021): On the Softmax Bottleneck of Recurrent Language Models