Learning through Internalization: How Transformers "Digest" Reasoning into Weights

TL;DR

Can a model learn a task it is fundamentally "too weak" to learn directly? This paper argues Yes, through a process called Internalization. By first learning to "think out loud" via Chain-of-Thought (CoT) and then gradually stripping those thoughts away, Transformers can absorb complex logic into their internal weights. While this makes them faster and capable of solving "hard" problems like Parity, it comes with a hidden cost: a significant drop in out-of-distribution (OOD) robustness.

The Motivation: Why Can't We Just Learn Directly?

In classical machine learning, we usually provide an input and a label ($xoy$). However, for many tasks—like modular arithmetic or the "Parity" problem—the path from $x$ to $y$ is a cryptographic-like hurdle. Standard Gradient Descent often fails to find the "needle in the haystack" solution.

The authors observe a human parallel: When we learn to drive, we initially verbalize every step ("clutch in, shift gear, gas out"). Eventually, the process becomes "internalized"—it happens in a single "forward pass" of our brain without active thought. This paper explores if Transformers can do the same.

Methodology: The Internalization Process

The authors define Internalization as the transition from a "Slow System" (e.g., 25 steps of CoT) to a "Fast System" (1 step of generation) that shares the same weights.

The Curriculum vs. The Mixture

1. Left-Amputation Curriculum: Training starts with full CoT. Gradually, tokens are removed from the beginning of the sequence.
2. Mixture Distribution: Instead of a strict schedule, the model is trained on a mix of all stages (some examples with full CoT, some with none). This is linked to the Positive Distribution Shift (PDS) framework.

Architecture Insight: Width vs. Depth

A common intuition is that if a model "internalizes" a 10-step reasoning process, it must spread those steps vertically through 10 layers. The researchers found the opposite: Wider, shallower models are often better at internalization than deep models.

Figure: Internalization success rates across isoparametric curves. Increasing width is far more beneficial than increasing depth for horizontal logic absorption.

Theoretical Breakthrough: Solving the Unsolvable

The paper provides a formal proof for the Sparse Parity Task. Learning parity ($y=x_1\oplus x_k$) is notoriously hard for Gradient Descent. However, they prove that a 1-layer Transformer with linear attention can:

1. First learn the parity by predicting the next token in a CoT sequence.
2. Successfully internalize the logic as CoT tokens are removed.
3. Eventually compute the final parity in a single step—a result that is "computationally hard" to achieve through direct training.

The "Flip Side": The OOD Performance Pitfall

The most striking finding is the "Internalization Tax." As a model becomes faster and more "automatic" at a task, its ability to handle shifted data distributions plummets.

Figure: As more CoT tokens are dropped (dashed lines), the out-of-distribution accuracy progressively degrades, even if in-distribution performance remains high.

This suggests that internalization encourages the model to learn shortcuts. Like a human driving "on autopilot," the model stops looking at the fine details of the logic and starts relying on baked-in patterns that fail when the environment changes.

Critical Analysis & Conclusion

The "Prime" Difficulty

One fascinating discovery made by the authors is that Prime-order Semiautomata ($C_n$ where $n$ is prime) are significantly harder to internalize than those with composite orders. They theorize that Transformers leverage the Chinese Remainder Theorem to parallelize math; when a number is prime, this "parallelization shortcut" is blocked, making the internal representation much more complex.

Final Takeaway

Learning through internalization is a double-edged sword:

• Pro: It enables models to solve problems that are otherwise mathematically "unlearnable" via standard supervised learning.
• Con: It risks creating "brittle" models that lack the robustness of their explicit-reasoning (CoT) counterparts.

For the future of AI Scaling, this implies that "System 2 to System 1" distillation (making models faster/smarter) might require training on a mixture of reasoning lengths rather than pure compression to prevent catastrophic OOD collapse.