💬 NLP

Masked Distillation: Internalizing the Chain-of-Thought in Language Models

This paper introduces "masked distillation," a framework that internalizes the reasoning capabilities of Large Reasoning Models into a student model's parameters by training it to predict only solution tokens while receiving feedback conditioned on the teacher's Chain-of-Thought, thereby enabling direct answer generation and reducing inference latency without sacrificing accuracy.

Original authors: Durgesh Kalwar, Vardhan Palod, Subbarao Kambhampati

Published 2026-07-28

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Durgesh Kalwar, Vardhan Palod, Subbarao Kambhampati

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

Imagine a world where your favorite AI assistant is like a brilliant but chatty student. When you ask it a tricky question, it doesn't just blurt out the answer; it scribbles pages of notes, draws diagrams, and talks through every single step of its thinking before finally writing down the solution. This "thinking out loud" process, known in the tech world as a "Chain of Thought," has made AI models incredibly smart at math and logic puzzles. However, there's a catch: those pages of notes take up a lot of time and energy to generate. It's like paying a taxi driver to drive you all the way to the grocery store, but also paying them to drive around the entire city block first just to "think" about which aisle you need. Scientists have long wondered: Can we teach the AI to do all that thinking inside its head? If we could, the AI would skip the long, chatty notes and just give you the answer instantly, saving massive amounts of time and computer power.

This paper, titled "Masked Distillation," dives right into that question. The researchers, working with Large Reasoning Models (LRMs), asked if they could take a smart "teacher" model that loves to think out loud and train a smaller "student" model to internalize that thinking process. They wanted to see if the student could learn to skip the notes entirely and just output the final answer, while still being just as smart. They tested this using a clever technique called "Masked Distillation." Imagine the teacher is solving a puzzle while the student is watching. The teacher sees the full question and its own detailed thought process, but the student is only shown the question. The goal was to see if the student could learn the teacher's secret logic and then solve the puzzle on its own without ever needing to write down the steps, even though the teacher had all the notes to guide the learning process.

The team set up two main experiments. In the first, they used the same AI model as both the teacher and the student, just asking it to think differently. In the second, they used a bigger, smarter model as the teacher and a smaller, simpler one as the student. They also introduced a "scaffold" concept. Think of this like training wheels on a bike. Sometimes, the student might need to see a little bit of the teacher's notes (the "scaffold") to get the answer right, rather than trying to do everything from scratch. They tested how much of the teacher's thinking the student needed to see to get the best results.

The results were a mix of "wow" and "not so fast." On a dataset of grade-school math problems called GSM8K, the student models were amazing. They successfully learned to internalize the teacher's thinking. The fully "masked" student (one that never wrote down any notes) got 76.0% of the answers right, which was a huge jump from its starting ability, and it did so using only 369.6 tokens (units of text) on average, compared to the teacher's 2,398.1 tokens. That's a massive 6.5x reduction in effort!

However, on a trickier number puzzle called Countdown, the story changed. Here, the student that tried to do everything in its head (fully masked) actually got worse, dropping to 41.7% accuracy. It turned out that for this specific type of puzzle, the student needed to see some of the teacher's notes to succeed. But here's the cool part: the researchers found a sweet spot. By letting the student see just a tiny bit of the teacher's thinking (about 30% of the notes, or an "alpha" of 0.3), the student's accuracy skyrocketed to 86.2%. This was almost as good as the teacher (87.3%) but still used about 1.3 times fewer tokens than the teacher. It was like giving the student just enough training wheels to balance perfectly without needing the whole bike.

The paper also looked at how well these students handled new, unseen puzzles (out-of-distribution tasks). On the math side, the students did well, but on the Countdown puzzles, the results were a bit more mixed. Interestingly, in some cases, the students trained with a little bit of scaffolding actually generalized better than the teacher itself when the puzzles got harder or changed slightly.

One key finding was that the method of teaching mattered. When they tried to teach the student just by showing it the teacher's notes and asking it to copy them (a method called Supervised Fine-Tuning), it didn't work as well. The "Masked Distillation" method, which uses a more complex way of comparing the student's guesses to the teacher's actual thoughts, was much more effective at transferring the "thinking" skills.

In the end, the paper suggests that we can't just magically delete all the thinking steps for every problem. Whether an AI can "think silently" depends heavily on the type of problem and how much the student model already knows about that topic. For familiar tasks like basic math, the AI can learn to think silently and save a ton of time. For harder, unfamiliar puzzles, it might still need a little help—a small scaffold of notes—to get the job done right. The researchers conclude that this "scaffold" is a powerful tool we can tune, allowing us to find the perfect balance between how smart the AI is and how fast it can give us an answer.

Technical Summary: Masked Distillation – Internalizing the Chain-of-Thought in Language Models

Problem Statement

Large Reasoning Models (LRMs) achieve high performance on complex tasks (math, code, planning) by generating long, explicit chains of intermediate tokens (Chain-of-Thought or CoT) before producing a final answer. While these traces improve accuracy, they incur significant costs in inference latency, memory usage (KV-cache footprint), and serving energy, which scale linearly with trace length.

Recent findings suggest that the correctness of these intermediate traces is not causally linked to the correctness of the final answer, nor does trace length reliably correlate with problem complexity. This raises a critical question: Can the computation expressed in these intermediate tokens be internalized into the parameters of a language model, allowing it to produce correct answers directly (or with significantly shorter traces) without explicitly generating the reasoning steps at inference time?

Methodology: Masked Distillation

The authors propose Masked Distillation, a knowledge-distillation framework designed to internalize the reasoning process of a "thinking" teacher model into a "non-thinking" student model.

Core Framework

In this framework, the student model is trained to predict only the solution tokens conditioned on the input question. The teacher model, however, provides feedback based on both the question and its own generated CoT trace. The objective is to minimize the divergence between the student's distribution (conditioned only on the input) and the teacher's distribution (conditioned on the input and the intermediate reasoning).

The training objective uses on-policy reverse-KL divergence:
$L_{MD} = D_{KL}(\pi_S(\cdot | x, y_{<t}) \parallel \pi_T(\cdot | x, \text{ITs}, y_{<t}))$
Where $\pi_S$ is the student policy, $\pi_T$ is the teacher policy, $x$ is the input, $y_{<t}$ are previously generated tokens, and $\text{ITs}$ are the teacher's intermediate tokens.

Experimental Settings

The authors instantiate this framework in two settings:

Self-Distillation: The same model acts as the teacher (in thinking mode) and the student (in non-thinking mode).
Dual-Model Setting: A larger reasoning model serves as the teacher, supervising a separate, smaller non-thinking student.

The $\alpha$ -Suffix Scaffold

To explore the trade-off between full internalization and inference cost, the authors introduce an $\alpha$ -suffix masked distillation variant.

$\alpha = 0$ (Fully Masked): The student internalizes the entire intermediate trace and emits only the solution.
$\alpha = 1$ (Non-Masked): The student reproduces the full trace before the answer (equivalent to standard CoT generation).
$0 < \alpha < 1$ : The student internalizes the $(1-\alpha)$ prefix of the teacher's trace and is trained to emit the remaining $\alpha$ -fraction suffix of the trace at inference, alongside the solution.

This allows the authors to treat intermediate tokens as a "scaffold" that can be partially retained or fully internalized.

Key Contributions

Proposed Framework: Introduced Masked Distillation to internalize CoT reasoning into model parameters, enabling direct answer generation.
Scaffold Analysis: Defined and evaluated the $\alpha$ -suffix parameter, establishing a controlled axis to tune the accuracy vs. inference-cost trade-off.
Comprehensive Evaluation: Conducted extensive experiments across two domains:
- Math: GSM8K (In-Distribution), MATH-500, and AIME-25 (Out-of-Distribution).
- Countdown: A number-puzzle search task with variations in target range and search depth (OOD splits).
Ablation Studies: Compared Masked Distillation (using reverse-KL) against Supervised Fine-Tuning (SFT) on teacher traces and the use of explicit "teaching tokens" as scaffolds.

Results

1. Task-Dependence of Full Internalization

The ability to fully internalize reasoning ( $\alpha=0$ ) is highly dependent on the task and the student's prior exposure to the domain during pretraining:

GSM8K (Math): Full internalization is successful. In the self-distillation setting, the student achieved 76.0% accuracy (vs. teacher's 87.9%) while reducing inference tokens by 6.5×. In the dual-model setting, accuracy improved from 15.69% (base) to 45.26% with a 5× token reduction.
Countdown: Full internalization fails without scaffolding. In self-distillation, the fully masked student dropped to 41.7% accuracy (below the base model's 56.2%). In the dual-model setting, it improved from 0% to 34.1% but remained far below the teacher (87.3%) and the non-masked variant (81.1%).

2. The Efficacy of $\alpha$ -Suffix Scaffolding

Providing a small amount of scaffolding (emitting a fraction of the trace) significantly recovers performance, particularly on tasks where full internalization fails:

Countdown (Self-Distillation): Setting $\alpha = 0.3$ raised accuracy from 41.7% to 86.2%, essentially matching the teacher (87.3%) while reducing inference tokens by ~1.3× compared to the non-masked variant.
Countdown (Dual-Model): The $\alpha = 0.5$ variant achieved 84.2% accuracy, outperforming the non-masked baseline (81.1%) while using fewer tokens, establishing a Pareto-optimal point.
GSM8K: Small scaffolds ( $\alpha=0.3$ ) recovered most of the accuracy gap (83.2% vs. 87.9%) with roughly half the token cost of the non-masked variant.

3. Generalization to Out-of-Distribution (OOD)

Target-Range Shift: Models trained with various scaffold paradigms transferred cleanly to OOD tasks where the target number range changed.
Search-Depth Shift: In the self-distillation setting on Countdown, scaffold-trained students (specifically $\alpha=0.5$ ) generalized better than the non-masked variant under increased search depth (32.8% vs. 25.7%).
Dual-Model Limitations: In the dual-model setting on harder math benchmarks (MATH-500, AIME-25), the capacity gap between the teacher and student was too large; neither full internalization nor scaffolding recovered OOD performance, with all variants collapsing to the base student's floor.

4. Training Objectives

Reverse-KL vs. SFT: The paper finds that Supervised Fine-Tuning (SFT) on teacher traces fails to exploit scaffolding effectively. In Countdown, SFT with scaffolds yielded ~38% accuracy, whereas Reverse-KL Masked Distillation with the same scaffolds reached ~67%.
Teaching Tokens: Providing explicit "teaching tokens" as scaffolds improved performance over full masking but was less effective than the $\alpha$ -suffix approach.

Significance and Claims

The paper claims that suffix scaffolding provides a controllable mechanism to balance accuracy and inference efficiency. The optimal operating point is determined by the specific task and the capacity gap between the teacher and student.

Internalization is possible but not universal: Models can internalize reasoning for tasks they are familiar with (like GSM8K) but struggle with novel search tasks (like Countdown) without retaining some explicit trace.
Cost-Performance Trade-off: A small amount of explicit reasoning (e.g., $\alpha=0.3$ ) can recover nearly all the performance of full CoT generation while offering substantial inference cost savings (1.3× to 2× reduction in tokens).
Generalization: Scaffolding can improve generalization in specific OOD scenarios (e.g., search-depth shifts in self-distillation), suggesting that retaining a partial trace helps the model navigate complex search spaces better than a fully internalized model or a full trace in certain contexts.

The authors conclude that while full internalization is an attractive goal for efficiency, the $\alpha$ -suffix approach offers a practical, tunable alternative that maintains high performance with reduced latency.