The Softmax Bottleneck in Large Language Models: Matrix Factorization Bounds, High-Rank Token Distributions, and Mixture of Softmaxes

Autoregressive language models predict probability distributions over vocabulary tokens conditioned on preceding text. In standard Transformer architectures, the model computes a hidden state vector $h_c \in \mathbb{R}^d$ for a given context $c$, projects it into vocabulary space using a linear unembedding matrix $W \in \mathbb{R}^{V \times d}$, and applies the softmax function to normalize the resulting logits into probabilities.

While computationally convenient, this formulation imposes a fundamental structural constraint known as the softmax bottleneck. First formalized by Yang et al. (2018), the softmax bottleneck establishes that parametric softmax layers restrict the expressiveness of neural language models by bounding the algebraic rank of the log-probability matrix by the hidden dimension $d$. When the vocabulary size $V$ and the linguistic complexity of natural language exceed $d$, a standard linear projection cannot represent arbitrary multinomial distributions over tokens across distinct contexts.

Language Modeling as Matrix Factorization

To understand the mathematical origin of the softmax bottleneck, language modeling can be framed as a matrix factorization problem.

Let $\mathcal{C}$ denote the set of all valid text contexts, and let $\mathcal{V}$ denote the vocabulary of tokens, where $|\mathcal{V}| = V$. The ground-truth language distribution defines an ideal conditional probability matrix $P^* \in \mathbb{R}^{|\mathcal{C}| \times V}$, where each entry $P^_{c, v} = P^(x = v \mid c)$ represents the true probability of token $v$ following context $c$.

Taking the logarithm of these probabilities yields the ground-truth log-probability matrix $A \in \mathbb{R}^{|\mathcal{C}| \times V}$:

A[c, v] = log P*(x = v | c)

In a parametric language model, the probability distribution over tokens for context $c$ is computed by:

P_theta(x = v | c) = exp(h_c^T w_v) / sum_{v' in V} exp(h_c^T w_{v'})

where:

• $h_c \in \mathbb{R}^d$ is the final hidden representation produced by the network for context $c$.
• $w_v \in \mathbb{R}^d$ is the unembedding vector for token $v$, corresponding to the $v$-th row of unembedding weight matrix $W \in \mathbb{R}^{V \times d}$.
• $d$ is the model hidden dimension (e.g., 2048, 4096, or 8192).

Taking the logarithm of the model's predicted probability produces the parameterized log-probability matrix $\hat{A} \in \mathbb{R}^{|\mathcal{C}| \times V}$:

A_hat[c, v] = log P_theta(x = v | c) = h_c^T w_v - g(h_c)

where $g(h_c) = \log \sum_{v' \in \mathcal{V}} \exp(h_c^T w_{v'})$ is the scalar log-partition function (the softmax normalizer) for context $c$.

In compact matrix notation:

A_hat = H W^T - g(H) 1_V^T

where:

• $H \in \mathbb{R}^{|\mathcal{C}| \times d}$ is the matrix of context hidden representations across all contexts.
• $W \in \mathbb{R}^{V \times d}$ is the vocabulary unembedding matrix.
• $g(H) \in \mathbb{R}^{|\mathcal{C}| \times 1}$ is the column vector of log-partition values.
• $1_V \in \mathbb{R}^{V \times 1}$ is a column vector of ones.

The Rank-d Constraint

The linear algebra of matrix factorization dictates strict bounds on the rank of the predicted log-probability matrix $\hat{A}$:

1. The inner product term $H W^T$ is the product of an $|\mathcal{C}| \times d$ matrix and a $d \times V$ matrix. By rank inequalities:

``text rank(H W^T) <= min(|C|, V, d) = d ``

1. The log-normalizer correction term $g(H) 1_V^T$ is the outer product of two vectors, which has an algebraic rank of exactly 1:

``text rank(g(H) 1_V^T) = 1 ``

1. By subadditivity of matrix rank ($\text{rank}(X + Y) \le \text{rank}(X) + \text{rank}(Y)$), the total rank of $\hat{A}$ is bounded:

``text rank(A_hat) <= rank(H W^T) + rank(g(H) 1_V^T) <= d + 1 ``

This is the formal definition of the softmax bottleneck: regardless of how deep the Transformer backbone is, how many attention heads it uses, or how many training tokens it consumes, the matrix of log-probabilities it produces over the vocabulary cannot have a rank greater than $d + 1$.

In modern large language models, the vocabulary size $V$ is often between 32,000 and 256,000 tokens (e.g., LLaMA 3 uses $V = 128,256$; Gemma 2 uses $V = 256,000$), while $d$ is typically 2,048 to 8,192. Because $d \ll V$, the model is forced to project probability distributions onto a strictly low-rank subspace of dimension $d + 1$.

Why Natural Language Requires High Rank

Natural language exhibits complex semantic, syntactic, and contextual dependencies that cannot be captured by low-rank linear projections.

1. Polysemy and Context-Dependent Word Substitutions

A single word can have disparate meanings depending on context (e.g., "bank" in financial vs. hydrological vs. aeronautical contexts). In each domain, "bank" co-occurs with a distinct cluster of vocabulary tokens.

In a low-rank projection, assigning high logits to "river", "water", and "flow" alongside "bank" in a geography context requires aligning $h_{\text{geo}}$ with the direction of $w_{\text{bank}}$. However, in a finance context, $h_{\text{fin}}$ must align with $w_{\text{bank}}$ while simultaneously aligning with $w_{\text{deposit}}$, $w_{\text{interest}}$, and $w_{\text{loan}}$, without spilling probability mass onto the geographical tokens. When thousands of polysemous words interact across millions of topic spaces, the geometry of a $d$-dimensional sphere becomes over-constrained, causing geometric interference.

2. Multi-Modal Context Manifolds

Natural language contexts often require multi-modal probability surfaces where two contexts agree on the probabilities of one set of words, disagree completely on a second set, and invert their preferences on a third set. Representing such combinatorial conditional distributions requires the log-probability matrix $A$ to have an effective rank approaching the full vocabulary dimension $V$.

When forced into a rank-$(d+1)$ parameterization, the model suffers from capacity truncation: it must smooth out fine-grained conditional variances across words to preserve general semantic coherence, leading to higher test perplexity on tail tokens.

Mixture of Softmaxes (MoS)

To break the low-rank restriction without inflating the hidden dimension $d$ across every layer of the network, Yang et al. (2018) introduced the Mixture of Softmaxes (MoS) architecture.

Instead of computing a single softmax distribution from a single context vector $h_c$, MoS computes a weighted mixture of $K$ separate softmax distributions:

P_MoS(x = v | c) = sum_{k=1}^K pi_{c, k} * [ exp(h_{c, k}^T w_v) / sum_{v'} exp(h_{c, k}^T w_{v'}) ]

where:

• $K$ is the number of mixture components (typically 3 to 15).
• $\pi_{c, k}$ is the prior mixture weight for component $k$, computed via a routing softmax:

``text pi_c = softmax(W_pi h_c) ``

• $h_{c, k} \in \mathbb{R}^d$ is the $k$-th contextual representation vector, generated by a component-specific projection:

``text h_{c, k} = tanh(W_{h, k} h_c) ``

• $w_v \in \mathbb{R}^d$ is the shared unembedding vector for token $v$.

Why MoS Breaks the Rank Constraint

In MoS, the total probability is a linear combination of exponentials rather than the exponential of a linear combination:

log P_MoS(x = v | c) = log [ sum_{k=1}^K pi_{c, k} * exp(h_{c, k}^T w_v - g_k(h_{c, k})) ]

Because the logarithm of a sum of exponentials (LogSumExp) is a non-linear operation, the resulting log-probability matrix is no longer expressible as a simple product of two low-rank matrices shifted by a rank-1 vector. The effective rank of the log-probability matrix scales with the number of components $K$, reaching up to $\min(|\mathcal{C}|, V, K \cdot d)$.

Empirical evaluations on benchmarks such as Penn Treebank and WikiText-2 demonstrated that adding MoS to language models produced substantial perplexity reductions without increasing the depth or sequence processing cost of the core recurrent or self-attention layers.

Architectural Alternatives and Modern Solutions

While MoS proved the theoretical limitation of standard softmax, computing $K$ separate full-vocabulary softmax normalizations introduces significant latency during training and inference. Subsequent research and modern LLM designs explore several alternative pathways to mitigate the bottleneck:

+-----------------------------------------------------------------------+
|                       Output Projection Paradigms                    |
+-----------------------------------------------------------------------+
|  Standard Softmax:                                                    |
|  h_c (d) ---------> [ W (V x d) ] ---------> Softmax ---------> Rank d|
+-----------------------------------------------------------------------+
|  Mixture of Softmaxes (MoS):                                          |
|  h_c (d) --+--> [ Proj 1 ] --> Softmax 1 \                            |
|            +--> [ Proj 2 ] --> Softmax 2 --> [ Weighted Sum ] -> High |
|            +--> [ Proj K ] --> Softmax K /                       Rank |
+-----------------------------------------------------------------------+
|  Expanded Non-Linear Unembedding:                                     |
|  h_c (d) ----> [ MLP Expansion: d -> 2d/4d ] ----> [ W ] ----> Softmax|
+-----------------------------------------------------------------------+
|  Untied Large Embeddings:                                             |
|  h_c (d) ----> [ W_out (V x d) independent of W_in ] ---------> Softmax|
+-----------------------------------------------------------------------+

1. Pointwise Non-Linearities and Mixtape

Kanai et al. (2019) demonstrated that the softmax bottleneck can be addressed by applying monotonic pointwise non-linear transformations directly to logit vectors before softmax normalization.

Building on this, Yang et al. (2019) developed Mixtape, which combines logit-space vector gating with sigmoid tree decomposition. Mixtape achieves high-rank expressiveness while reducing the computational overhead of MoS by 3.5x to 10.5x.

2. Embedding Tying vs. Untying

Early Transformer implementations (including Vaswani et al. (2017) and GPT-2) tied the input token embedding weights to the output unembedding matrix ($W_{\text{in}} = W_{\text{out}}$). While weight tying reduces parameter counts for large vocabularies, it forces input representations and output logit distributions to share the exact same low-rank metric space.

Modern frontier LLMs increasingly untie input and output embeddings (e.g., LLaMA 3, Mistral, and DeepSeek architectures). Untying allows the unembedding matrix $W_{\text{out}}$ to specialize exclusively in optimizing decision boundaries and negative log-likelihood across the vocabulary, partially mitigating geometric distortion.

3. Logit Soft-Capping and Dimensional Scaling

In architectures with very large vocabularies ($V \ge 128\text{k}$), models like Gemma 2 incorporate logit soft-capping:

logits = cap * tanh( (h_c^T w_v) / cap )

Soft-capping prevents extreme logit magnitudes from collapsing the gradient space and ensures that multiple competing tokens remain active in the normalization pool, avoiding premature rank degeneration during training.

4. Expansion Projections Before Unembedding

Rather than projecting directly from hidden state $h_c \in \mathbb{R}^d$ to $V$, some architectures apply an intermediate expansion layer $f: \mathbb{R}^d \to \mathbb{R}^{d_{\text{out}}}$ where $d_{\text{out}} = 2d$ or $4d$ prior to matrix multiplication with $W$. This raises the maximum rank bound to $d_{\text{out}} + 1$ without incurring the full quadratic attention cost of running wide hidden dimensions throughout the entire Transformer stack.

Parallels in Multi-Head Attention and MoE Routing

The mathematical principles of the softmax bottleneck extend beyond next-token vocabulary prediction to other components of modern AI architectures:

• Self-Attention Normalization: The scaled dot-product attention map $A = \text{softmax}(Q K^T / \sqrt{d_k})$ computes token-to-token attention weights. For a sequence of length $T$, the attention matrix $A \in \mathbb{R}^{T \times T}$ is generated from projections of dimension $d_k$ (where $d_k = d / h$). When sequence length $T \gg d_k$, a single attention head is rank-constrained. Multi-Head Attention (MHA) resolves this by computing $h$ parallel low-rank attention heads, analogous to a Mixture of Softmaxes over spatial positions.
• Sparse MoE Gating: Mixture-of-Experts routers use softmax gating to allocate tokens among $E$ available experts ($E \in [16, 256]$). Standard linear gating networks can suffer from routing bottlenecks when assigning complex, multi-task tokens across fine-grained expert pools, motivating multi-head or hierarchical gating mechanisms in advanced MoE architectures.

Summary

The softmax bottleneck is a structural constraint arising from the low-rank matrix factorization inherent in standard parametric softmax layers. When the hidden dimension $d$ is smaller than the vocabulary size $V$, a single linear unembedding cannot represent arbitrary high-rank token distributions.

By analyzing this constraint through the lens of matrix rank, methods such as Mixture of Softmaxes, logit non-linearities, untied embeddings, and dimensional expansion demonstrate how the final output interface of a language model dictates its representational limits.

Sources

• Breaking the Softmax Bottleneck: A High-Rank RNN Language Model (Yang et al., ICLR 2018 / arXiv:1711.03953)
• Breaking the Softmax Bottleneck via Learnable Monotonic Pointwise Non-linearities (Kanai et al., ICML 2019 / arXiv:1902.08077)
• Mixtape: Breaking the Softmax Bottleneck Efficiently (Yang et al., NeurIPS 2019 / Google Research)
• Softmax Bottleneck Makes Language Models Unable to Represent Multi-Modal Distributions (Ganea et al., ACL 2022)
• Attention Is All You Need (Vaswani et al., NeurIPS 2017 / arXiv:1706.03762)
• The Softmax Bottleneck in Language Modeling (Tibo Vanleke, 2026)