Why Small Models Get Stuck: Exploring the Softmax Bottleneck

The Hidden Math of the “Stupidity Wall”

From Microsoft’s Phi-3 to Google’s Gemma, the race is on to shrink intelligence down to a size that fits on your phone. The assumption is usually simple: if we just train these smaller models for longer, on more tokens, they should eventually “catch up” in reasoning capability. Here’s what nobody tells you: there’s a wall. Not a metaphorical skill ceiling or a “we need more data” kind of problem, an actual mathematical barrier where your model physically runs out of room to think.

This fascinating recent paper, Why do small language models underperform?, reveals what this wall looks like in practice. In this post, I want to break down the Softmax Bottleneck and why models with fewer than 1,000 hidden dimensions might be destined to fail, no matter how long you train them.

The Mystery of the Plateau

When training massive models (like Llama-3-70B), the rule is usually “more is better.” More data equals lower loss. But researchers noticed something strange happening with smaller models (under 500M parameters).

Instead of getting smarter, they hit a wall. In late-stage training, their performance on specific tasks doesn’t just stall - it often gets worse. This phenomenon is called Saturation. The model has seen enough data to be smarter, but its brain (its architecture) physically rejects the new knowledge.

Why? The answer lies in the tension between the complexity of language and the geometry of the model.

The Conflict: God’s Dictionary vs. A Tiny RAM

To understand the failure, you have to look at the final layer of the Transformer: the Language Modeling Head.

This is the layer that takes the model’s “thought” (the hidden state, $d_{model}$) and projects it onto the Vocabulary (often 50,000+ words).

• The Vocabulary Space: Imagine a high-dimensional cloud representing all of Natural Language. It is incredibly complex. To perfectly map the probability of the next word, you might need a rank (dimensionality) of 10,000 or more.

• The Hidden State: This is the model’s working memory. In a small model like GPT-2 Small or Pythia-160M, this dimension is small - often just 768.

This creates a Rank Mismatch. You are trying to project a 10,000-dimensional probability manifold down through a 768-dimensional funnel. It’s like trying to play a symphony using only a kazoo. You physically cannot represent the complexity of the music.

The Evidence: Watching the “Brain” Collapse

The authors of the paper didn’t just theorize this; they visualized it. They tracked the Singular Value Distribution of the model’s output layer during training.

Singular values essentially measure how many different “directions” or “features” the model uses to solve problems.

• In a healthy model: The distribution is smooth. The model uses thousands of subtle features to make decisions.

• In a saturated small model: The distribution collapses. It becomes “spiky”.

This collapse is known as Representation Degeneration or Anisotropy. As the model desperately tries to lower its loss, it realizes it can’t learn everything. So, it panics. It focuses 100% of its capacity on the “loudest” features (likely simple grammar and high-frequency words) and completely deletes the subtle, low-rank features that represent complex reasoning.

The model effectively lobotomizes itself to save space.

The Math: The Lossy JPEG of Intelligence

The paper quantifies this failure using the Eckart-Young-Mirsky theorem.

Without getting too bogged down in the algebra, this theorem is the “JPEG compression” rule of Linear Algebra. It proves that if you are forced to approximate a high-rank matrix (Language) with a low-rank matrix (Small Model), the error you introduce is equal to the sum of the singular values you threw away.

Because natural language is “heavy-tailed” (meaning the subtle details matter a lot), “throwing away” the bottom 90% of dimensions guarantees a high minimum loss. The math dictates that a model with $d_{model}<1000$ carries a permanent “stupidity penalty” that no amount of training data can erase.

---

Engineering Takeaways: How to Beat the Bottleneck

So, if you are building or fine-tuning models in the sub-1B parameter range, are you doomed? Not necessarily. The paper and its findings suggest three specific strategies to engineer your way around the bottleneck.

1. The “Pancake” Strategy (Width over Depth)

The most direct fix is topological. If the bottleneck is caused by a narrow hidden dimension ($d_{model}$), you must widen it, even if that means sacrificing other parts of the model to stay within your compute budget.

The paper explicitly notes that performance drops noticeably when the hidden dimension falls below 1,000. This suggests a hard floor for architecture design.

• The Old Way: A 100M param model might use $d_{model}=768$ and 12 layers. This hits the wall.

• The New Way: You should prioritize keeping $d_{model} \ge 1024$ or even $2048$. To keep the parameter count low, you significantly reduce the number of layers (depth).

While the authors warn that reducing depth can hurt generalization, modern “Small Language Models” (like Phi-3 or Gemma) have largely proven that this trade-off is worth it. They are essentially “Pancakes” (wide and shallow) specifically to maintain a high rank in the residual stream.

2. Specialize to Lower the Data Rank

If you can’t increase the model’s rank (because you are constrained by hardware), you must decrease the data’s rank.

The Softmax Bottleneck is only a problem because “General English” is high-rank (10k+ dimensions). But specific sub-domains are often much simpler.

• Action: Don’t train a small model to be a generalist. Train it to be a specialist.

• Why it works: A model dedicated solely to SQL generation or Medical Summarization deals with a much smaller effective manifold. The “tail” of singular values in a specialized dataset is lighter, meaning a small $d_{model}$ can capture the domain without “shaving off” critical information. This explains why small coding models often outperform massive generalist models on coding benchmarks.

3. Break the Linearity (Mixture of Softmaxes)

The bottleneck exists because the final layer is a linear projection ($H \times W^T$). A linear map cannot easily represent a non-linear, high-rank manifold.

The paper suggests exploring “expressive softmax alternatives”.

• The Research Fix: Instead of one single projection, you can use a Mixture of Softmaxes (MoS). This is effectively having $K$ different “heads” at the output, each calculating a probability distribution, which are then combined non-linearly.

• The Intuition: It’s like having three different low-rank painters work on the same canvas. One paints the “Grammar,” one paints the “Facts,” and one paints the “Tone.” Combined, they create a high-fidelity image that a single painter couldn’t achieve.

Conclusion: Beyond the “Mini-GPT”

For the last few years, the strategy for on-device AI has been simple: take a state-of-the-art Large Language Model and shrink it until it fits on a phone. The Softmax Bottleneck suggests that this strategy is hitting a mathematical dead end.

We are seeing the “Stupidity Wall” not because small models are incapable of intelligence, but because we are forcing them to play a game they can’t win. We are asking low-rank matrices to approximate the high-rank manifold of the entire internet.

The future of efficient ML likely isn’t a single 2B parameter model that knows everything. It is a constellation of high-rank specialists - one shaped for coding, one for creative writing, one for logic - swapped in and out of memory as needed.

The Softmax Bottleneck is a warning sign that the “Scaling Laws” work differently in the basement. To break through, we have to stop optimizing for parameter count and start optimizing for rank. The next great breakthrough in small models won’t come from more data; it will come from better geometry.

---

References

1. The Main Paper (The source of the “Saturation” findings)

Godey, N., de la Clergerie, É., & Sagot, B. (2024). Why do small language models underperform? Studying LM Saturation via the Softmax Bottleneck. arXiv preprint arXiv:2404.07647.

2. The Origin of the “Softmax Bottleneck” Theory

Yang, Z., Dai, Z., Salakhutdinov, R., & Cohen, W. W. (2018). Breaking the Softmax Bottleneck: A High-Rank RNN Language Model. International Conference on Learning Representations (ICLR).

3. On Representation Degeneration (Anisotropy)

Gao, J., He, D., Tan, X., Qin, T., Wang, L., & Liu, T. (2019). Representation Degeneration Problem in Training Natural Language Generation Models. International Conference on Learning Representations (ICLR).

4. On the Geometry of Word Embeddings (The “Cone Effect”)

Ethayarajh, K. (2019). How Contextual are Contextualized Word Representations? Comparing the Geometry of BERT, ELMo, and GPT-2 Embeddings. Proceedings of EMNLP-IJCNLP.