🤖 machine learning

Learning through Internalization

This paper investigates how neural networks, particularly transformers, internalize explicit computational procedures like chain-of-thought into their weights to facilitate learning complex tasks such as parity and semiautomata simulation, while also revealing a trade-off where this process leads to degraded out-of-distribution performance.

Original authors: Nikolaos Tsilivis, Nirmit Joshi, Marko Medvedev, Julia Kempe, Nati Srebro

Published 2026-06-23

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Nikolaos Tsilivis, Nirmit Joshi, Marko Medvedev, Julia Kempe, Nati Srebro

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Idea: From "Step-by-Step" to "Instinct"

Imagine you are learning to drive a car. At first, a teacher tells you exactly what to do for every single step: "Press the clutch, shift to first gear, release the clutch slowly while pressing the gas." You have to think hard about every intermediate action. This is slow, but it helps you learn the complex task.

Eventually, after lots of practice, you don't need to think about the steps anymore. You just "know" how to drive. You have internalized the process. The complex sequence of actions has been compressed into a single, automatic instinct.

This paper studies how Artificial Intelligence (specifically, a type of AI called a Transformer) does the same thing. It asks: Can we teach an AI to solve hard problems by first making it write out all its work (like a student showing their math homework), and then training it to "forget" the steps and just give the answer instantly?

The paper finds that yes, this works and can help AI learn things it otherwise couldn't. However, there is a catch: when the AI becomes "automatic," it sometimes gets worse at handling new, strange situations.

1. The Two Ways of Learning

The researchers compare two ways an AI can learn a task:

Direct Learning (The "Hard Way"): You show the AI a problem and the final answer, and ask it to figure out the connection immediately. For some very tricky logic puzzles (like calculating a specific type of math problem called "parity"), the AI often fails. It's like asking a beginner to drive a car without ever teaching them the steps.
Learning through Internalization (The "Smart Way"):
1. Phase 1 (The Teacher): You teach the AI to solve the problem by forcing it to write out every single step (called "Chain-of-Thought" or CoT). It generates a long list of intermediate answers before the final one. Because it has these hints, it learns the task easily.
2. Phase 2 (The Internalization): You slowly start removing the hints. First, you remove the last step, then the second-to-last, and so on. The AI is forced to rewire its brain (its "weights") to figure out how to get the answer without the crutches.
3. The Result: Eventually, the AI can solve the problem in one single step, without writing anything down. It has "internalized" the logic.

The Paper's Claim: This two-step process allows the AI to learn difficult tasks that it would fail to learn if you tried to teach it directly from the start.

2. The "Flip Side": Speed vs. Safety

The paper discovers a surprising downside to this "automatic" mode.

The Analogy: Think of a student who memorizes a specific route to school. If they practice driving that exact route every day, they become very fast and efficient. But if you ask them to drive to a new neighborhood they've never seen, they might get lost because they only learned the specific path, not the general rules of driving.
The Finding: When the AI internalizes a task (becomes fast and automatic), it often performs worse on Out-of-Distribution (OOD) data. This means if you give it a problem that looks slightly different from what it practiced on, it makes more mistakes than it did when it was still writing out its steps.
Why? The paper suggests that when the AI learns to be fast, it might find "shortcuts" or "tricks" that work for the practice data but aren't robust enough for new situations.

3. What Makes a Task Easy or Hard to Internalize?

The researchers tested this on different types of logic puzzles (called "semiautomata"). They found that not all tasks are equally easy to turn into an instinct.

Width vs. Depth: They found that making the AI "wider" (giving it more parallel processing power) helped it internalize tasks better than making it "deeper" (adding more layers of thinking). It's like having a wider highway is better for traffic than adding more toll booths.
Prime Numbers are Tricky: They discovered that tasks involving "prime numbers" (like counting in groups of 7 or 13) were much harder for the AI to internalize than tasks involving numbers with many factors (like 12 or 16).
- The Metaphor: Imagine trying to pack a suitcase. If you have items that fit together perfectly (like a 12-piece set), it's easy to pack them efficiently. If you have items that don't divide evenly (like a prime number), it's messy and hard to compress into a small space. The AI struggles to "compress" the logic of prime-number tasks into its internal memory.

4. The Mathematical Proof: Parity

To prove this isn't just luck, the authors did a rigorous mathematical proof using a specific math problem called "Sparse Parity."

The Problem: Imagine you have a long list of 0s and 1s, and you need to find out if the total number of 1s is odd or even. But you don't know which numbers in the list matter; you have to figure that out too. This is known to be extremely hard for computers to learn directly.
The Proof: They showed that if they first taught the AI to write out the answer step-by-step (calculating the parity of the first number, then the first two, then the first three, etc.), and then slowly removed those steps, the AI could mathematically guarantee that it would learn the solution.
The Takeaway: This proves that "learning through internalization" is a valid strategy to solve problems that are theoretically impossible to learn directly.

5. A New Shortcut: The "Mixture" Method

Finally, the paper suggests a way to skip the long, slow training process.

The Old Way: You have to train the AI with full steps, then remove one step, train again, remove another, and so on. This is like a strict curriculum.
The New Idea (Mixture): Instead of a strict order, you just throw all the different versions of the problem at the AI at once. Some examples have full steps, some have half steps, some have no steps.
The Result: The paper shows that training on this "mixed bag" of data works just as well as the strict step-by-step removal. It's like giving a student a practice test that has some easy questions, some medium ones, and some hard ones all mixed together, rather than starting with easy and slowly getting harder. This makes the AI robust enough to answer both directly and with steps, depending on what you ask it to do.

Summary

Internalization is teaching AI to think slowly (with steps) so it can eventually think fast (instinctively).
It works: It allows AI to learn hard logic puzzles that it couldn't learn otherwise.
The Cost: When AI becomes "fast," it sometimes becomes "brittle," meaning it fails more often on new, weird examples.
The Secret: Making the AI wider helps more than making it deeper, and mixing different types of training data is a simple, effective way to achieve this learning.

Technical Summary: Learning through Internalization

Problem Statement

The paper investigates the phenomenon of internalization in neural network-based systems, specifically focusing on how these systems absorb explicit computational procedures (such as Chain-of-Thought, or CoT) into their weights to perform tasks more efficiently. While neural networks often benefit from step-by-step supervision to learn complex tasks (e.g., simulating semiautomata or computing parities), the direct learning of these tasks without such scaffolding is often computationally hard or impossible for standard architectures.

The core problem addressed is twofold:

Mechanism: How do transformers transition from generating explicit intermediate reasoning steps (CoT) to directly computing the final output in fewer steps (ideally one forward pass) while maintaining accuracy?
Generalization Trade-off: Does this internalization process preserve the robustness of the original reasoning, or does it lead to "shortcut" solutions that degrade out-of-distribution (OOD) performance?
Theoretical Hardness: Can internalization serve as a viable strategy to bypass known computational lower bounds for learning certain hard functions (like parities) that are difficult to learn via direct end-to-end training?

Methodology

The authors employ a combination of empirical studies on autoregressive transformers and rigorous theoretical analysis of simplified transformer architectures.

1. Empirical Framework: Semiautomata Simulation

The authors study the simulation of semiautomata (specifically modular counter semiautomata $C_n$ and the permutation group $S_3$ ) over $T$ timesteps. They compare three training regimes:

End-to-End (E2E): Training only on the final state.
Chain-of-Thought (CoT): Training with all intermediate states explicitly provided.
Internalization via Curriculum: A two-stage process where the model is first trained with full CoT supervision, followed by a curriculum that progressively removes intermediate tokens (amputating them from the left, right, or inductively) to force the model to internalize the computation.

They analyze the impact of architectural choices (depth vs. width) and the specific properties of the semiautomata (e.g., the order of the group) on the success of internalization.

2. Theoretical Framework: Sparse Parity

To provide provable guarantees, the authors analyze a simplified one-layer transformer with linear attention and a fixed non-linear MLP. They study the task of learning a hidden parity function over an unknown support set.

Curriculum: The model is trained on a sequence of inputs where CoT tokens (partial parities) are progressively removed from the left.
Algorithm: They define a specific training algorithm (Algorithm 2) involving gradient descent with specific initialization and rounding steps to demonstrate that the model can learn the explicit CoT first and then internalize the computation to output the parity directly.

3. Connection to Positive Distribution Shift (PDS)

The paper frames internalization within the Positive Distribution Shift (PDS) framework (Medvedev et al., 2026). They argue that learning a "slow" system (with CoT) on a distribution that includes hints, and then internalizing it, constitutes a PDS. This allows the model to learn a "hard" target distribution by first learning an "easier" shifted distribution. They further explore whether a sequential curriculum can be replaced by a single mixture distribution containing examples of varying CoT lengths.

Key Contributions

1. Empirical Findings on Semiautomata

Width vs. Depth: Successful internalization is more strongly correlated with model width than depth. Wider, shallower architectures (e.g., depth 4) outperform deeper ones (e.g., depth 25) at fixed parameter budgets. This suggests internalization occurs "horizontally" across the model's width rather than "vertically" across layers.
Prime Order Difficulty: Semiautomata with prime orders (e.g., $C_{13}, C_{17}$ ) are significantly harder to internalize than those with composite orders (e.g., $C_{14}, C_{15}$ ). The authors conjecture this relates to the ability of transformers to compactly represent counter semiautomata whose orders have small prime-power factors (Theorem 1).
The OOD Trade-off: Internalization often leads to a degradation in Out-of-Distribution (OOD) performance. While CoT-trained models generalize well to biased input distributions, internalized models tend to learn "shortcut" solutions that fail under OOD conditions (Figure 1, right).
Curriculum Direction: Left-curriculum (removing tokens from the start) is generally more effective than right-curriculum for internalization, likely due to the stability of positional embeddings in the transformer.

2. Theoretical Proofs for Parity

Provable Internalization: The authors provide the first provable analysis showing that a simplified one-layer transformer can learn a hidden parity with explicit CoT supervision and then provably internalize the computation as CoT tokens are removed (Theorem 2).
Computational Hardness: They demonstrate that learning the parity representation directly from data without CoT internalization is computationally hard (in the Statistical Query model), whereas the internalization path is tractable.

3. Mixture Distributions as an Alternative

The paper demonstrates that for both MLPs and autoregressive transformers, the sequential curriculum for internalization can be replaced by training on a single mixture distribution containing examples with varying CoT lengths (and answer length hints). This mixture approach achieves similar success in learning hard tasks and, notably, can recover better OOD performance in the full-CoT regime compared to curriculum-trained models.

Results

Semiautomata: Transformers trained with CoT achieve near-perfect In-Distribution (ID) and OOD accuracy. However, after internalization via a left-curriculum, ID accuracy remains high, but OOD accuracy drops significantly (often to random chance for prime-order automata).
Parity Task: The theoretical model successfully learns the parity function. After the curriculum, the model outputs the correct parity in a single forward pass, having internalized the multi-step logic.
Mixture Learning: Experiments show that training on a mixture of CoT lengths (with length hints) allows a single model to both answer directly and generate full CoT, often outperforming curriculum-trained models in OOD robustness when full CoT is allowed.

Significance and Claims

The paper claims that learning through internalization is a powerful strategy that enables neural networks to solve tasks that are otherwise computationally intractable for direct learning. By first learning a "slow" system with explicit reasoning and then compressing this into weights, models can bypass hardness barriers (as seen in the parity task).

However, the authors also highlight a critical caveat: internalization is not a panacea. The process often leads to a loss of robustness, where the model learns brittle shortcuts that fail under distributional shifts. This suggests that while internalization reduces inference costs and enables learning, it may compromise the generalization capabilities that explicit reasoning provides.

Finally, the work connects these empirical observations to the Positive Distribution Shift framework, suggesting that the success of curriculum-based learning and internalization can be understood as leveraging a sequence of easier distributions to learn a hard target. The finding that a single mixture distribution can replace complex curricula offers a more practical and robust pathway for training models on hard reasoning tasks.