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

# Masked Distillation: Internalizing the Chain-of-Thought in Language Models
Source: [https://arxiv.org/html/2607.22629](https://arxiv.org/html/2607.22629)
Durgesh Kawlar SCAI, Arizona State University dkalwar@asu\.edu &Vardhan Palod∗ SCAI, Arizona State University vpalod@asu\.edu &Subbarao Kambhampati SCAI, Arizona State University rao@asu\.edu Equal contribution \- Joint first Author Accepted to FoGen 2026: Foundations of Deep Generative Models: Understanding Memorization, Generalization, and Reasoning, an ICML 2026 workshop\.

###### Abstract

Large Reasoning Models \(LRMs\) produce long, explicit chains of intermediate steps before generating a final answer at inference time\. These intermediate traces dominate latency, memory usage, and serving cost, even though the final answer correctness is not causally related to the trace correctness and the trace length is not a reliable indicator of the problem complexity\. This raises a natural question: can the computation expressed in these intermediate tokens be internalized into the parameters of a language model, enabling it to produce answers directly \(or with much shorter intermediate traces\)? We introducemasked distillation, a knowledge\-distillation framework in which a student LLM is trained to predict only the solution tokens conditioned on the question, while a reasoning teacher provides feedback on the student’s responses after conditioning on the question and its own CoT trace\. We instantiate this framework in two settings: \(i\) aself\-distillationsetting, in which the same model serves as the teacher in thinking mode and as the student in non\-thinking mode, and \(ii\) adual\-modelsetting, in which a larger reasoning teacher supervises a separate smaller non\-thinking student over the solution tokens\. By treating intermediate tokens as a scaffold which reasoning models use to fit over the solution tokens, We additionally vary the length of intermediate\-token scaffolding the student is supervised on, interpolating between full internalization \(the student emits only the solution\) and no internalization \(the student emits the full trace before the answer\)\. We evaluate the framework through controlled experiments on two reasoning domains: GSM8K \(grade\-school arithmetic\) and Countdown \(a number\-puzzle search task\)\. Our results show that the success of full internalization is sharply task\-dependent and tracks the student’s prior exposure to the task during pretraining: it works on GSM8K but fails on Countdown without inference\-time scaffolding\. However, providing the student models with a small scaffold closes the gap to self\-distillation on Countdown\. The task performance increases from41\.7%41\.7\\%\(fully masked\) to86\.2%86\.2\\%, essentially matching the teacher \(87\.3%87\.3\\%\) at∼\\sim1\.3×1\.3\\timesfewer inference tokens than the non\-masked variant\.

## 1Introduction

Large Reasoning Models \(LRMs\), post\-trained with RL, achieve strong performance on math, code, and planning benchmarks by generating long chains of intermediate tokens before producing a final answer\. While these traces improve task performance, they also significantly increase inference cost\. LRMs spend most of their generation budget on intermediate tokens rather than the answer, and inference latency, KV\-cache footprint, and energy consumption all grow roughly linearly with trace length\. At deployment scale, this is the main cost of serving them\.

A broad line of work aims to reduce this cost for LRMs by adding length\-control objectives during RL post\-training \[L1\[[2](https://arxiv.org/html/2607.22629#bib.bib2)\], CoT\-Valve\[[17](https://arxiv.org/html/2607.22629#bib.bib17)\], O1\-Pruner\[[16](https://arxiv.org/html/2607.22629#bib.bib16)\], Kimi\-1\.5\-style length\-penalty rewards\[[25](https://arxiv.org/html/2607.22629#bib.bib25)\], GFPO\[[23](https://arxiv.org/html/2607.22629#bib.bib23)\],\[[3](https://arxiv.org/html/2607.22629#bib.bib3)\]\]\. These methods reward shorter traces, but the model still produces an explicit trace at inference, and the savings are bounded by how short a trace can be\. While the fact of the performance increase with intermediate traces is well\-known, the reasons for it are less clear\. Our recent work has shown that there is no causal connection between trace correctness and final\-answer correctness\[[26](https://arxiv.org/html/2607.22629#bib.bib26),[13](https://arxiv.org/html/2607.22629#bib.bib13),[4](https://arxiv.org/html/2607.22629#bib.bib4)\], and the trace length is not connected to the computational complexity of the problem instance being solved\[[19](https://arxiv.org/html/2607.22629#bib.bib19)\]\. These findings raise the question that if model’s reasoning is not causally connected to its solution, is it necessary for the model to explicitly produce these tokens and can the computation expressed in these intermediate tokens be internalized into the parameters of a language model?

A recent line of work related to self\-distillation has shown that language models can be trained to internalize additional context using knowledge distillation\[[12](https://arxiv.org/html/2607.22629#bib.bib12),[21](https://arxiv.org/html/2607.22629#bib.bib21)\]\. The general idea of distilling from a larger, more capable teacher into a smaller student is itself well established\[[9](https://arxiv.org/html/2607.22629#bib.bib9),[1](https://arxiv.org/html/2607.22629#bib.bib1)\];\[[12](https://arxiv.org/html/2607.22629#bib.bib12)\]and\[[21](https://arxiv.org/html/2607.22629#bib.bib21)\]take this further by showing that even when the teacher and the student are the same model, the teacher’s access to additional context \(expert demonstrations or environment feedback\) can be transferred into the student’s parameters\. Prior to the emergence of LRMs, there were works which investigated whether Chain\-of\-Thought \(CoT\) tokens could be internalized into a model’s parameters\.\[[24](https://arxiv.org/html/2607.22629#bib.bib24)\]introduced context distillation and showed that a T5\-small can be distilled from a CoT\-prompted version of itself into a direct\-answer model that solves arithmetic problems\.\[[7](https://arxiv.org/html/2607.22629#bib.bib7)\]generalized this with implicit chain\-of\-thought via knowledge distillation \(ICoT\-KD\), and\[[8](https://arxiv.org/html/2607.22629#bib.bib8)\]later proposed stepwise internalization \(ICoT\-SI\), which removes intermediate tokens stage by stage so that the teacher’s reasoning is gradually compressed into the student’s hidden states\. A common pattern across these methods is that in\-distribution accuracy can usually be matched increased inference efficiency, provided the training setup is right\. However, the trade\-off between extra training cost and inference savings is not well\-studied\. More importantly, whether the resulting students generalize is also an open question\.

Building on\[[10](https://arxiv.org/html/2607.22629#bib.bib10),[21](https://arxiv.org/html/2607.22629#bib.bib21),[12](https://arxiv.org/html/2607.22629#bib.bib12)\], we use knowledge distillation to internalize the information carried by intermediate tokens into the student’s parameters\. We proposemasked distillation, a knowledge\-distillation framework in which a student model is trained to predict only the solution tokens conditioned on the question, while the teacher provides feedback on the student’s responses after being conditioned on the question as well as its CoT trace\. We instantiate this framework in two settings: \(i\) aself\-distillationsetting, where the same model serves as the teacher in thinking mode and as the student in non\-thinking mode, and \(ii\) adual\-modelsetting, where a reasoning model supervises a separate non\-thinking model over the solution tokens\. By treating intermediate tokens as a scaffold that reasoning models use to fit over the solution tokens, we additionally study anα\\alpha\-suffix masked distillationvariant where the student internalizes the\(1−α\)\(1\-\\alpha\)*prefix*of the teacher’s intermediate trace and is trained to produce the remainingα\\alpha\-fraction*suffix*of the trace at inference, together with the solution tokens\.

In this study, we explore the following research questions:

- •RQ1\.Through masked distillation, can a student model internalize the information present in the teacher’s reasoning tokens and achieve similar task performance with improved inference efficiency?
- •RQ2\.If the masked\-distillation student does not achieve similar task performance as the teacher, does providing additional scaffolding \(measured in intermediate tokens emitted at inference\) improve the task performance and what is the tradeoff against inference cost?
- •RQ3\.Does the internalization of intermediate tokens lead to degradation of performance on out\-of\-distribution \(OOD\) tasks, and does providing additional scaffolding help the student maintain OOD performance?
- •RQ4\.Does the training objective matter? Specifically, can we provide different scaffolds to the student and can we use supervised fine\-tuning on teacher rollouts to achieve similar results?

We conduct extensive experiments across two domains, Math \(GSM8K, with MATH\-500 and AIME\-25 as out\-of\-distribution splits\) and Countdown \(with target\-range and search\-depth shifts as out\-of\-distribution splits\) under both*self\-distillation*and*dual\-model*settings\. Sweeping the suffix\-scaffold parameterα∈\{0,0\.3,0\.5,0\.7,1\.0\}\\alpha\\in\\\{0,0\.3,0\.5,0\.7,1\.0\\\}, our results show that the ability to fully internalize the information present in intermediate tokens is task\-dependent: it succeeds on GSM8K, where the base student has prior exposure to the domain, but fails on Countdown\. However, our results demonstrate that theα\\alpha\-suffix scaffold closes this gap at modest cost; for example, applyingα=0\.3\\alpha=0\.3to self\-distillation on Countdown raises accuracy from41\.7%41\.7\\%\(fully masked\) to86\.2%86\.2\\%, essentially matching the teacher \(87\.3%87\.3\\%\) at a∼\\sim1\.3×1\.3\\timesreduction in inference tokens relative to the non\-masked variant\. In out\-of\-distribution tasks, the models trained under various scaffold paradigms transfer cleanly under target\-range shift, and in the self\-distillation setting they generalize better than the non\-masked variant under search\-depth shift\. These findings establish suffix scaffolding as a controlled axis along which the accuracy vs\. inference\-cost trade\-off can be tuned, with the optimal operating point determined by the task and the teacher\-student capacity gap\.

The rest of the paper is organized as follows\. Section[2](https://arxiv.org/html/2607.22629#S2)provides the background on knowledge\-distillation techniques\. Section[3](https://arxiv.org/html/2607.22629#S3)presents the masked\-distillation framework\. Section[4](https://arxiv.org/html/2607.22629#S4)details the experimental setup\. Section[5](https://arxiv.org/html/2607.22629#S5)reports results on GSM8K and Countdown and discusses their implications for the scaffolding–generalization tradeoff\.

## 2Background

### 2\.1Knowledge Distillation

Traditionally, Knowledge distillation is a framework where there is a student\-teacher pair and the student model is being trained to mimic the behaviour of the teacher by minimizing the divergence between their output distributions\. This framework was first introduced by\[[10](https://arxiv.org/html/2607.22629#bib.bib10)\]for transferring knowledge from an ensemble or from a large highly regularized model into a smaller, distilled model\. In the context of autoregressive models, knowledge distillation has been extensively studied for training a smaller Language model to mimic the outputs of a larger teacher LLM\. Letθ\{\\theta\}denote the student model’s parameters, andpSθp\_\{S\}^\{\\theta\}denote the student model’s policy, differentiable w\.r\.tθ\{\\theta\}\. Let us denote the dataset of input\-output pairs as\(X,Y\)\(X,Y\)\. For a divergenceDD, we define the discrepancy between token\-level distributions ofpTp\_\{T\}andpSp\_\{S\}as

DpT∥pSθ\(y∣x\):=1Ly∑n=1LyD\(pT\(⋅∣y<n,x\)∥pSθ\(⋅∣y<n,x\)\),D\_\{p\_\{T\}\\parallel p^\{\\theta\}\_\{S\}\}\(y\\mid x\):=\\frac\{1\}\{L\_\{y\}\}\\sum\_\{n=1\}^\{L\_\{y\}\}D\\\!\\left\(p\_\{T\}\(\\cdot\\mid y\_\{<n\},x\)\\;\\middle\\\|\\;p^\{\\theta\}\_\{S\}\(\\cdot\\mid y\_\{<n\},x\)\\right\),\(1\)
There are many variants of KD studied in the literature and the distinction between the variants is based on whether the teacher’s probability distributionpTp\_\{T\}is accessible during training and from where the targets,YY, are obtained\.

- •Sequence\-level KD:\[[14](https://arxiv.org/html/2607.22629#bib.bib14)\]IfpTp\_\{T\}is not accessible during training, and the outputsYYare sampled from the teacher on inputsXX, then with supervised finetuning, the student model is trained to maximize the likelihood of high\-probability teacher sequences\.
- •Supervised KD:\[[10](https://arxiv.org/html/2607.22629#bib.bib10)\],\[[20](https://arxiv.org/html/2607.22629#bib.bib20)\]The student model is trained to mimic the next\-token probability distribution of the teacher model\. The lossLSD​\(θ\)L\_\{\\mathrm\{SD\}\}\(\\theta\)over the dataset\(X,Y\)\(X,Y\)is calculated as minθLSD\(θ\):=𝔼\(x,y\)∼\(X,Y\)\[DKL\(pT∥pSθ\)\(y∣x\)\]\.\\min\_\{\\theta\}\\;L\_\{\\mathrm\{SD\}\}\(\\theta\):=\\mathbb\{E\}\_\{\(x,y\)\\sim\(X,Y\)\}\\left\[D\_\{\\mathrm\{KL\}\}\\\!\\left\(p\_\{T\}\\;\\middle\\\|\\;p^\{\\theta\}\_\{S\}\\right\)\(y\\mid x\)\\right\]\.\(2\)
- •On\-policy KD:MinilLM\[[9](https://arxiv.org/html/2607.22629#bib.bib9)\], GKD\[[1](https://arxiv.org/html/2607.22629#bib.bib1)\], performs on\-policy knowledge distillation by generating responses from the student model\. LGKD​\(θ\):=𝔼x∼X,y∼pSθ\(⋅∣x\)\[DKL\(pT\(⋅∣y<n,x\)∥pSθ\(⋅∣y<n,x\)\)\]\\begin\{split\}L\_\{\\mathrm\{GKD\}\}\(\\theta\)&:=\\mathbb\{E\}\_\{x\\sim X,\\;y\\sim p^\{\\theta\}\_\{S\}\(\\cdot\\mid x\)\}\\left\[D\_\{\\mathrm\{KL\}\}\\\!\\left\(p\_\{T\}\(\\cdot\\mid y\_\{<n\},x\)\\;\\middle\\\|\\;p^\{\\theta\}\_\{S\}\(\\cdot\\mid y\_\{<n\},x\)\\right\)\\right\]\\end\{split\}\(3\) In\[[1](https://arxiv.org/html/2607.22629#bib.bib1)\], authors have shown that on\-policy KD works well with different variants of KL divergence such as forward KL, reverse KL and JSD\. Sequence\-level KD and Supervised KD have the problem of distribution mismatch between output sequences seen during training and those generated by the student during inference\. The on\-policy training in GKD addresses this drawback effectively\.

## 3Masked Distillation

![Refer to caption](https://arxiv.org/html/2607.22629v1/figures/masked_distillation_framework_latest.drawio.png)Figure 1:Overview of the masked distillation framework\.Phase 1, intermediate\-token collection \(top\):for each input questionx∼𝒟x\\sim\\mathcal\{D\}we sample a response from the teacherπT\\pi^\{T\}\. The response has the form<<think\>\>\(ITs\)<<\\\\backslashthink\><\><answer\>\>\(STs\)<<\\\\backslashanswer\>\>, from which we extract the intermediate tokens \(ITs\) and the solution tokens \(STs\)\.Phase 2, training \(bottom\):the non\-thinking studentπθS\\pi^\{S\}\_\{\\theta\}is conditioned only onxx, while the frozen teacher is conditioned onxxtogether with the first\(1−α\)\(1\-\\alpha\)fraction of the ITs from Phase 1\. The parameterα\\alphacontrols how much of the intermediate trace the student is supervised to emit at inference: the student is trained to reproduce the lastα\\alphafraction of the teacher’s intermediate trace together with the solution tokens\. The student’s next\-token distributionPPis matched against the teacher’s distributionQQvia the on\-policy reverse\-KL loss𝔼y∼P​\[log⁡P/Q\]\\mathbb\{E\}\_\{y\\sim P\}\[\\log P/Q\]computed on student\-sampled tokens; gradients update only the student\. Atα=0\\alpha=0the variant reduces to fully masked distillation, and atα=1\\alpha=1to non\-masked distillation\. When the teacher and the student are the same model, this is the*self\-distillation*setting; when the teacher is a larger thinking model and the student is a smaller non\-thinking model, this is the*dual\-model*setting\.We present a knowledge\-distillation framework, which we call*masked distillation*, that internalizes the intermediate\-token \(IT\) generation process \(or the so\-called reasoning process\) of a reasoning modelπT\\pi^\{T\}\(the teacher model\) into a non\-thinking student modelπS\\pi^\{S\}\. The goal is to push the teacher’s reasoning into the student’s parameters so that, at inference time, the student directly generates the solution tokens conditioned only on the input problem, without emitting any intermediate tokens of its own\. This makes inference substantially more efficient and cheaper, since the intermediate computation has already been internalized into the student’s internal representations during distillation\. Concretely, the student is trained to mimic the teacher’s conditional distribution over the input problemxxand the teacher’s own generated intermediate tokens,πT\(⋅∣x,ITs\)\\pi^\{T\}\(\\,\\cdot\\mid x,\\mathrm\{ITs\}\)\.

We consider two settings of*masked distillation*:*self\-distillation*and the*dual\-model*setting\. In the*self\-distillation*setting, the same model acts as both the teacher and the student: the teacher operates in thinking mode and generates intermediate tokens, while the student operates in non\-thinking mode and is trained to directly produce the final solution\. In the*dual\-model*setting, a larger reasoning model serves as the teacher and supervises a separate, smaller non\-thinking student on the generation of solution tokens\. Unless otherwise specified, we use the term*masked distillation*to refer generically to the framework, and explicitly distinguish between the two settings only when necessary\.

Figure[1](https://arxiv.org/html/2607.22629#S3.F1)shows an overview of the masked distillation framework\. InPhase 1, we collect intermediate tokens for each questionx∈𝒟x\\in\\mathcal\{D\}by sampling a response from the teacher model and extracting the segment between the<think\>and</think\>tags as the ITs; InPhase 2, we train the student model by minimizing the divergence between the student distributionπθS\(⋅∣x\)\\pi^\{S\}\_\{\\theta\}\(\\,\\cdot\\mid x\)and the teacher distributionπT\(⋅∣x,ITs\)\\pi^\{T\}\(\\,\\cdot\\mid x,\\mathrm\{ITs\}\), where the teacher is conditioned onxxtogether with the ITs it generated in Phase 1, while the student is conditioned only onxx\. The masked\-distillation objective is the reverse\-KL divergence between the student and teacher next\-token distributions, computed on student\-sampled responses\. The loss is defined as

ℒMD\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{MD\}\}=DKL\(πθS\(⋅∣x,y<t\)∥πT\(⋅∣x,ITs,y<t\)\)\\displaystyle=D\_\{\\mathrm\{KL\}\}\\\!\\left\(\\pi^\{S\}\_\{\\theta\}\(\\,\\cdot\\mid x,y\_\{<t\}\)\\,\\big\\\|\\,\\pi^\{T\}\(\\,\\cdot\\mid x,\\mathrm\{ITs\},y\_\{<t\}\)\\right\)\(4\)=𝔼yt′∼πθS\(⋅∣x,y<t\)​\[log⁡πθS​\(yt′∣x,y<t\)πT​\(yt′∣x,ITs,y<t\)\]\\displaystyle=\\mathbb\{E\}\_\{\\,y^\{\\prime\}\_\{t\}\\,\\sim\\,\\pi^\{S\}\_\{\\theta\}\(\\,\\cdot\\mid x,y\_\{<t\}\)\}\\\!\\left\[\\log\\frac\{\\pi^\{S\}\_\{\\theta\}\(y^\{\\prime\}\_\{t\}\\mid x,y\_\{<t\}\)\}\{\\pi^\{T\}\(y^\{\\prime\}\_\{t\}\\mid x,\\mathrm\{ITs\},y\_\{<t\}\)\}\\right\]=∑yt′∈𝒱πθS​\(yt′∣x,y<t\)​log⁡πθS​\(yt′∣x,y<t\)πT​\(yt′∣x,ITs,y<t\)\.\\displaystyle=\\sum\_\{y^\{\\prime\}\_\{t\}\\,\\in\\,\\mathcal\{V\}\}\\pi^\{S\}\_\{\\theta\}\(y^\{\\prime\}\_\{t\}\\mid x,y\_\{<t\}\)\\;\\log\\frac\{\\pi^\{S\}\_\{\\theta\}\(y^\{\\prime\}\_\{t\}\\mid x,y\_\{<t\}\)\}\{\\pi^\{T\}\(y^\{\\prime\}\_\{t\}\\mid x,\\mathrm\{ITs\},y\_\{<t\}\)\}\.
where𝒱\\mathcal\{V\}is the vocabulary\. Since we use the same family of models for the teacher and the student, their vocabularies are identical\.

We compare masked distillation against a*non\-masked*variant in which both the teacher and the student are conditioned only on the input questionx∈𝒟x\\in\\mathcal\{D\}, so that the student also learns to reproduce the teacher’s intermediate\-token generation process before emitting the solution\. For non\-masked distillation the loss is

ℒNon−MD=DKL\(πθS\(⋅∣x,y<t\)∥πT\(⋅∣x,y<t\)\)\.\\mathcal\{L\}\_\{\\mathrm\{Non\-MD\}\}=D\_\{\\mathrm\{KL\}\}\\\!\\left\(\\pi^\{S\}\_\{\\theta\}\(\\,\\cdot\\mid x,y\_\{<t\}\)\\,\\big\\\|\\,\\pi^\{T\}\(\\,\\cdot\\mid x,y\_\{<t\}\)\\right\)\.\(5\)In the*self\-distillation*setting, since the teacher and the student are the same model, the non\-masked variant reduces to running the teacher in thinking mode at inference\.

Masked and non\-masked distillation lie at the two extremes of a spectrum that controls how much of the teacher’s reasoning process the student reproduces at inference\. Under masked distillation the student fully internalizes the intermediate\-token generation process and emits only the solution tokens at inference; under non\-masked distillation the student, like the teacher, first produces intermediate tokens and only then the final answer\. To probe the effect of intermediate\-token scaffolding between these extremes, we additionally study anα\\alpha\-suffix masked distillationvariant where the student internalizes the\(1−α\)\(1\-\\alpha\)*prefix*of the teacher’s intermediate trace and is trained to produce the remainingα\\alpha\-fraction*suffix*of the trace at inference, together with the solution tokens\. The teacher is conditioned on input question and prefix that the student has to internalized, while the student is conditioned only on input question:

ℒMD​\-​S​f​xα=DKL\(πθS\(⋅∣x,y<t\)∥πT\(⋅\|x,ITs1:⌊\(1−α\)​\|ITs\|⌋,y<t\)\),\\mathcal\{L\}^\{\\alpha\}\_\{\\mathrm\{MD\}\\text\{\-\}Sfx\}=D\_\{\\mathrm\{KL\}\}\\\!\\left\(\\pi^\{S\}\_\{\\theta\}\(\\,\\cdot\\mid x,y\_\{<t\}\)\\,\\big\\\|\\,\\pi^\{T\}\\\!\\left\(\\,\\cdot\\,\\big\|\\,x,\\;\\mathrm\{ITs\}\_\{1:\\lfloor\(1\-\\alpha\)\\,\|\\mathrm\{ITs\}\|\\rfloor\},\\;y\_\{<t\}\\right\)\\right\),\(6\)whereITs1:⌊\(1−α\)​\|ITs\|⌋\\mathrm\{ITs\}\_\{1:\\lfloor\(1\-\\alpha\)\\,\|\\mathrm\{ITs\}\|\\rfloor\}denotes the first\(1−α\)\(1\-\\alpha\)tokens of the teacher’s intermediate trace\. Atα=0\\alpha=0the variant reduces to fully masked distillation \(Eq\.[4](https://arxiv.org/html/2607.22629#S3.E4)\); atα=1\\alpha=1it reduces to non\-masked distillation \(Eq\.[5](https://arxiv.org/html/2607.22629#S3.E5)\)\. In practice we reportα∈\{0\.3,0\.5,0\.7\}\\alpha\\in\\\{0\.\\text\{3\},0\.\\text\{5\},0\.\\text\{7\}\\\}to map the suffix\-budget axis at a few representative operating points\.

## 4Experimental Setup

We conducted our experiments in the following domains \-

1. 1\.Math \-We usedGSM8Kas our In\-distribution dataset\. This dataset, introduced by\[[5](https://arxiv.org/html/2607.22629#bib.bib5)\], consists of grade school math problems, and is designed to evaluate the reasoning capabilities of large language models\. It contains 8\.5K problems, each paired with a question and an answer\. The dataset is divided into 7\.5K training problems and 1K test problems\. To test generalization in this domain, we test on the AIME\-25\[[18](https://arxiv.org/html/2607.22629#bib.bib18)\]and MATH\-500\[[15](https://arxiv.org/html/2607.22629#bib.bib15)\]datasets as OOD splits that share the same arithmetic substrate but exercise it at substantially higher difficulty\.
2. 2\.Countdown \-This domain\[[6](https://arxiv.org/html/2607.22629#bib.bib6)\]is a generalized version of the classic 24 Game\[[29](https://arxiv.org/html/2607.22629#bib.bib29)\], where the objective is to combine a set of input numbers using basic arithmetic operations \(\+, –, ×, ÷\) to reach a specified target number\. In our ID dataset, each problem consists of 3 to 4 two\-digit input numbers, with the target number also being a two\-digit number\. The dataset contains 9K examples, split into 8K training instances and 1K test instances\. For OOD testing, we have three datasets, each perturbing a different axis\.*OOD\-1*keeps the input cardinality and digit count fixed but shifts the target to a one\-digit number \(*target\-range shift*;678678problems\)\.*OOD\-2*and*OOD\-3*keep the digit counts fixed but increase the input cardinality to55and66numbers respectively \(*search\-depth shift*;522522and478478problems\)\.

Training details:We run all experiments on two NVIDIA H100 GPUs with 80 GB of VRAM each, and our implementation builds on theverllibrary\[[22](https://arxiv.org/html/2607.22629#bib.bib22)\]\. For our experiments in the Math domain, in theself\-distillationsetting, we use the Qwen3\-1\.7B\[[28](https://arxiv.org/html/2607.22629#bib.bib28)\]base model where the model in non\-thinking mode serves as the student, while the same model in thinking mode acts as the teacher\. For ourdual\-modelsetting, we use the Qwen3\-1\.7B base model in thinking mode as the teacher and Qwen2\-0\.5B\-Instruct\[[27](https://arxiv.org/html/2607.22629#bib.bib27)\]as the student model\. Similarly in the countdown domain, we use Qwen3\-4B base model as the teacher and Qwen2\-0\.5B\-Instruct model as the student\. The masked\-distillation loss is the reverse\-KL divergence between the student and teacher next\-token distributions, computed over the full vocabulary\. We provide additional hyperparameter details in the appendix section[A](https://arxiv.org/html/2607.22629#A1)\.

## 5Results and Discussion

Figures 2–5 show the accuracy and average response length for the different values ofα\\alphain our suffix\-scaffold masked\-distillation variants, evaluated on GSM8K and Countdown across in\-distribution \(ID\) and out\-of\-distribution \(OOD\) splits\. Referring to these, we address our first three research questions: RQ1 \(the effect of fully internalizing the teacher’s intermediate tokens on in\-distribution performance\) in subsection[5\.1](https://arxiv.org/html/2607.22629#S5.SS1); RQ2 \(whether additional scaffolding budget improves task performance, and its trade\-off against inference cost\) in subsection[5\.2](https://arxiv.org/html/2607.22629#S5.SS2); and RQ3 \(Does internalizing intermediate tokens degrade OOD performance, and whether scaffolding helps to recover it\) in subsection[5\.3](https://arxiv.org/html/2607.22629#S5.SS3)\. Finally, we compare on\-policy reverse\-KL distillation with supervised fine\-tuning on teacher rollouts \(RQ4\) in subsection[5\.4](https://arxiv.org/html/2607.22629#S5.SS4)\.

![Refer to caption](https://arxiv.org/html/2607.22629v1/figures/pareto_gsm8k_in_d_alltokens_self_distill_setting.png)\(a\)GSM8K \(in\-distribution\)
![Refer to caption](https://arxiv.org/html/2607.22629v1/figures/pareto_math_ood_alltokens_self_distill_setting.png)\(b\)MATH\-500 \(OOD\)
![Refer to caption](https://arxiv.org/html/2607.22629v1/figures/pareto_aime25_ood_alltokens_self_distill_setting.png)\(c\)AIME25 \(OOD\)

Figure 2:Accuracy vs\. average response length on math domain \(self\-distillation setting\) for suffix\-scaffold masked distillation variants withα∈\{0,0\.3,0\.5,0\.7,1\.0\}\\alpha\\in\\\{0,0\.3,0\.5,0\.7,1\.0\\\}Teacher: Qwen3\-1\.7B \(thinking\); Student: Qwen3\-1\.7B \(non\-thinking\)\.### 5\.1Effects of full internalization of teacher’s intermediate tokens on in\-distribution performance

In this section we analyze whether student can internalize the information present in the teacher’s reasoning tokens while still matching teacher\-level performance at lower inference cost\. The degree to which full internalisation succeeds depends on \(i\) distillation setting \(self\-distillation vs\. dual\-model\) and \(ii\) task\. We discuss the two settings in turn\.

Self\-distillation setting:On GSM8K \(Figure[2\(a\)](https://arxiv.org/html/2607.22629#S5.F2.sf1)\), the base student \(Qwen3\-1\.7B\-Base in non\-thinking mode\) attains61\.0%61\.0\\%accuracy on the GSM8K test set, while the teacher \(Qwen3\-1\.7B\-Base in thinking mode\) reaches87\.9%87\.9\\%\. Training using the fully masked distillation \(α=0\\alpha=0\) paradigm closes this gap: the student reaches76\.0%76\.0\\%accuracy while emitting only369\.6369\.6tokens on average against the teacher’s2398\.12398\.1, a6\.5×6\.5\\timesreduction in inference cost\.

On Countdown \(Figure[4\(a\)](https://arxiv.org/html/2607.22629#S5.F4.sf1)\), the base student \(Qwen3\-4B\-Base in non\-thinking mode\) attains56\.2%56\.2\\%accuracy on the Countdown ID test set and the teacher \(Qwen3\-4B\-Base in thinking mode\) reaches87\.3%87\.3\\%\. Unlike on GSM8K, fully masked distillation here does not recover this gap with the fully masked student reaching only41\.7%41\.7\\%accuracy\. The student does emit far shorter responses on average \(96\.396\.3tokens against the teacher’s2499\.62499\.6, a26×26\\timesreduction\), but its in\-distribution accuracy is now lower than the student’s base model\.

![Refer to caption](https://arxiv.org/html/2607.22629v1/figures/pareto_gsm8k_in_d_alltokens_dual_distill_setting.png)\(a\)GSM8K \(in\-distribution\)
![Refer to caption](https://arxiv.org/html/2607.22629v1/figures/pareto_math_ood_alltokens_dual_distill_setting.png)\(b\)MATH\-500 \(OOD\)
![Refer to caption](https://arxiv.org/html/2607.22629v1/figures/pareto_aime25_ood_alltokens_dual_distill_setting.png)\(c\)AIME25 \(OOD\)

Figure 3:Accuracy vs\. average response length on the math domain \(dual\-model setting\), for suffix\-scaffold masked distialltion variants withα∈\{0,0\.3,0\.5,0\.7,1\.0\}\\alpha\\in\\\{0,0\.3,0\.5,0\.7,1\.0\\\}\. Teacher: Qwen3\-1\.7B \(thinking\); Student: Qwen2\-0\.5B\-Instruct\.Dual\-model setting:On GSM8K \(Figure[3\(a\)](https://arxiv.org/html/2607.22629#S5.F3.sf1)\), the student \(Qwen2\-0\.5B\-Instruct\) is distilled from a larger Qwen3\-1\.7B teacher that reaches89\.23%89\.23\\%accuracy\. The base student already attains15\.69%15\.69\\%on GSM8K without any distillation, indicating that GSM8K style problems are likely well\-represented in its pretraining distribution\. Full internalization substantially improves the student over its base model, raising accuracy from15\.69%15\.69\\%to45\.26%45\.26\\%at a5×5\\timesreduction in average response length relative to the teacher\. The performance of the fully\-masked student is comparable to the performance of the model trained with non\-masked distillation with a significant gain in inference efficiency\.

On Countdown \(Figure[5\(a\)](https://arxiv.org/html/2607.22629#S5.F5.sf1)\), the student is distilled from a larger Qwen3\-4B teacher that reaches87\.3%87\.3\\%accuracy\. Notably, the base student has0%0\\%accuracy on Countdown which suggests that the Countdown puzzle format is likely not well represented in its pretraining data\. Fully masked distillation increases the accuracy of the student from0%0\\%to34\.1%34\.1\\%, but unlike the GSM8K dual\-model setting, there is a significant gap in performance of the student and of the non\-masked variant \(81\.1%81\.1\\%\) which is trained to reproduce the full thinking trace\. The fully masked student does emit much shorter responses on average \(403\.8403\.8tokens against the non\-masked variant’s3052\.63052\.6, a7\.6×7\.6\\timesreduction\)\.

![Refer to caption](https://arxiv.org/html/2607.22629v1/figures/pareto_in_d_alltokens_self_distill_setting.png)\(a\)In\-distribution
![Refer to caption](https://arxiv.org/html/2607.22629v1/figures/pareto_ood1_targetshift_alltokens_self_distill_setting.png)\(b\)OOD: target shift
![Refer to caption](https://arxiv.org/html/2607.22629v1/figures/pareto_ood_5num_alltokens_self_distill_setting.png)\(c\)OOD: 5 numbers

Figure 4:Accuracy vs\. average response length on Countdown \(self\-distillation setting\), for suffix\-scaffold masked distillation variants withα∈\{0,0\.3,0\.5,0\.7,1\.0\}\\alpha\\in\\\{0,0\.3,0\.5,0\.7,1\.0\\\}\. Teacher: Qwen3\-4B \(thinking\); Student: Qwen3\-4B \(non\-thinking\)\.
### 5\.2Effect of providing additional scaffolding,α\\alpha\-suffix, on task performance and inference cost

The RQ1 results show that full internalization does not always reach the same performance as the model in thinking mode or the non\-masked variant\. A natural follow\-up is whether allowing the student to emit a partial scaffold at inference, a fractionα\\alphaof the teacher’s intermediate trace can recover this accuracy, and at what cost in inference budget\. In addition to the fully masked and non\-masked variants, we sweepα∈\{0\.3,0\.5,0\.7\}\\alpha\\in\\\{0\.3,0\.5,0\.7\\\}on both tasks and trace the resulting accuracy vs\. response\-length frontier\.

Self\-distillation setting:On GSM8K \(Figure[2\(a\)](https://arxiv.org/html/2607.22629#S5.F2.sf1)\), adding even a small suffix scaffold delivers most of the recovery\. Withα=0\.3\\alpha=0\.3the student reaches83\.2%83\.2\\%accuracy, a77\-point gain over the fully masked76\.0%76\.0\\%baseline, while emitting1,1221\{,\}122tokens on average\. Further increasing the suffix budget leaves accuracy in a similar range \(79\.0%79\.0\\%atα=0\.5\\alpha=0\.5and75\.2%75\.2\\%atα=0\.7\\alpha=0\.7\), and only the full non\-masked variant \(α=1\.0\\alpha=1\.0\) reaches87\.9%87\.9\\%at2,3982\{,\}398tokens\. The most efficient operating point isα=0\.3\\alpha=0\.3: it produces∼\\sim2×2\\timesshorter responses than the non\-masked variant \(1,1221\{,\}122vs\.2,3982\{,\}398tokens\) at the cost of about55accuracy points \(83\.2%83\.2\\%vs\.87\.9%87\.9\\%\)\.

On Countdown \(Figure[4\(a\)](https://arxiv.org/html/2607.22629#S5.F4.sf1)\), scaffolding has a much sharper effect\. The fully masked baseline reaches only41\.7%41\.7\\%, below the base student butα=0\.3\\alpha=0\.3jumps to86\.2%86\.2\\%, essentially matching the teacher’s87\.3%87\.3\\%\. The sweep then plateaus:87\.8%87\.8\\%atα=0\.5\\alpha=0\.5,86\.6%86\.6\\%atα=0\.7\\alpha=0\.7, and87\.3%87\.3\\%atα=1\.0\\alpha=1\.0, all within a22\-point range\. The most efficient operating point is thereforeα=0\.3\\alpha=0\.3, which recovers teacher\-level accuracy at a∼\\sim1\.3×1\.3\\timesreduction in inference cost relative to the non\-masked variant \(1,8711\{,\}871vs\.2,5002\{,\}500tokens\)\.

Dual\-model setting:On GSM8K \(Figure[3\(a\)](https://arxiv.org/html/2607.22629#S5.F3.sf1)\), the fully masked variant reaches44\.9%44\.9\\%accuracy and the non\-masked variant reaches50\.3%50\.3\\%\. The narrow5\.45\.4\-point gap between these two endpoints indicates that the dual\-model GSM8K setting has limited room for intermediate scaffolding to help: the residual that anα\\alpha\-suffix variant could close is small, and the inference\-cost penalty of any non\-zero scaffold rapidly outweighs the available accuracy gain\.

On Countdown \(Figure[5\(a\)](https://arxiv.org/html/2607.22629#S5.F5.sf1)\), scaffolding produces a progressive and substantial recovery\. From the fully masked baseline of34\.1%34\.1\\%, accuracy rises to73\.5%73\.5\\%atα=0\.3\\alpha=0\.3,84\.2%84\.2\\%atα=0\.5\\alpha=0\.5,78\.0%78\.0\\%atα=0\.7\\alpha=0\.7, and81\.1%81\.1\\%atα=1\.0\\alpha=1\.0\. The peak of the sweep isα=0\.5\\alpha=0\.5, which slightly exceeds the non\-masked baseline \(84\.2%84\.2\\%vs\.81\.1%81\.1\\%\), and the response length at this operating point \(2,8072\{,\}807tokens\) is below the non\-masked variant’s3,0533\{,\}053\. The dual\-model Countdown frontier is therefore not just recovered but improved by the suffix scaffold:α=0\.5\\alpha=0\.5is pareto\-optimal than the non\-masked variant on both accuracy and inference cost\.

![Refer to caption](https://arxiv.org/html/2607.22629v1/figures/pareto_in_d_alltokens_dual_model_setting.png)\(a\)In\-distribution
![Refer to caption](https://arxiv.org/html/2607.22629v1/figures/pareto_ood1_targetshift_alltokens_dual_model_setting.png)\(b\)OOD: target shift
![Refer to caption](https://arxiv.org/html/2607.22629v1/figures/pareto_ood_5num_alltokens_dual_model_setting.png)\(c\)OOD: 5 numbers

Figure 5:Accuracy vs\. average response length on Countdown \(dual model setting\), for suffix\-scaffold masked distillation variants withα∈\{0,0\.3,0\.5,0\.7,1\.0\}\\alpha\\in\\\{0,0\.3,0\.5,0\.7,1\.0\\\}\. Teacher: Qwen3\-4B \(thinking\); Student: Qwen2\-0\.5B\-Instruct\.
### 5\.3Generalization under distribution shift

The RQ1 and RQ2 results analyze the accuracy v/s efficiency trade\-off on in distribution tasks\. In this section, we analyze whether the various training paradigms demonstrate similar trends in OOD settings \(RQ3\)\.

Self\-distillation setting:On the math domain \(Figure[2\(b\)](https://arxiv.org/html/2607.22629#S5.F2.sf2)\), the suffix sweep on MATH\-500 broadly follows the GSM8K ID pattern\. The base student attains69\.0%69\.0\\%on MATH\-500 and the fully masked variant \(α=0\\alpha=0\) reaches only66\.0%66\.0\\%slightly below base but every scaffolded variant lifts accuracy above the base model:75\.6%75\.6\\%atα=0\.3\\alpha=0\.3,76\.6%76\.6\\%atα=0\.5\\alpha=0\.5,77\.0%77\.0\\%atα=0\.7\\alpha=0\.7, and76\.6%76\.6\\%atα=1\.0\\alpha=1\.0\(non\-masked\)\. Theα=0\.3\\alpha=0\.3operating point recovers essentially all of the non\-masked accuracy at a∼\\sim1\.7×1\.7\\timesshorter response length \(2,5422\{,\}542vs\.4,2844\{,\}284tokens\)\. On AIME\-25 \(Figure[2\(c\)](https://arxiv.org/html/2607.22629#S5.F2.sf3)\) the base student reaches6\.7%6\.7\\%, the fully masked variant10\.0%10\.0\\%, and the scaffold variants range from16\.7%16\.7\\%to26\.7%26\.7\\%in a non\-monotonic pattern\. Interestingly, we see that theα=0\.7\\alpha=0\.7variant performs better than the teacher model while being more efficient during inference\.

On Countdown \(Figures[4\(b\)](https://arxiv.org/html/2607.22629#S5.F4.sf2)and[4\(c\)](https://arxiv.org/html/2607.22629#S5.F4.sf3)\), the suffix scaffold transfers cleanly under target\-range shift but produces a more interesting pattern under search\-depth shift\. Under target\-range shift, the base student reaches12\.5%12\.5\\%and the fully masked variant18\.3%18\.3\\%; the scaffold variants close most of the gap to non\-masked, with83\.6%83\.6\\%atα=0\.3\\alpha=0\.3,82\.0%82\.0\\%atα=0\.5\\alpha=0\.5,85\.1%85\.1\\%atα=0\.7\\alpha=0\.7, and83\.5%83\.5\\%atα=1\.0\\alpha=1\.0\. Theα=0\.3\\alpha=0\.3variant essentially matches the non\-masked baseline at this shift\. Under search\-depth shift, accuracy is much lower across the board \(1\.1%1\.1\\%base,3\.4%3\.4\\%fully masked\), but the scaffold variants actually*outperform*the non\-masked baseline:29\.9%29\.9\\%atα=0\.3\\alpha=0\.3,32\.8%32\.8\\%atα=0\.5\\alpha=0\.5,27\.2%27\.2\\%atα=0\.7\\alpha=0\.7, and25\.7%25\.7\\%atα=1\.0\\alpha=1\.0\. The peakα=0\.5\\alpha=0\.5variant is about77points above non\-masked, suggesting that under search\-shift OOD, the scaffold\-trained student generalizes better than the teacher model\.

Dual\-model setting:In math domain \(Figure[3\(b\)](https://arxiv.org/html/2607.22629#S5.F3.sf2)\), none of the distillation variants meaningfully outperform the base student on MATH\-500\. The base student attains6\.8%6\.8\\%accuracy and all distilled variants sit in a narrow6\.0%6\.0\\%–9\.4%9\.4\\%band \(9\.4%9\.4\\%fully masked,6\.0%6\.0\\%atα=0\.3\\alpha=0\.3,6\.4%6\.4\\%atα=0\.5\\alpha=0\.5,8\.4%8\.4\\%atα=0\.7\\alpha=0\.7,7\.2%7\.2\\%atα=1\.0\\alpha=1\.0\)\. The fully masked variant is the highest of the distilled variants but only marginally above base\. On AIME\-25 \(Figure[3\(c\)](https://arxiv.org/html/2607.22629#S5.F3.sf3)\), every variant including the base student reaches0%0\\%, These results indicate that on harder math benchmarks, the capacity gap between the Qwen3\-1\.7B teacher and the Qwen2\-0\.5B\-Instruct student is large enough that neither full internalization nor suffix scaffolding recovers OOD performance: the student converges to roughly the base\-student floor in every case\.

On Countdown \(Figures[5\(b\)](https://arxiv.org/html/2607.22629#S5.F5.sf2)and[5\(c\)](https://arxiv.org/html/2607.22629#S5.F5.sf3)\), the suffix scaffold transfers under target\-range shift but fails under search\-depth shift\. Under target\-range shift, the base student reaches0\.0%0\.0\\%, the fully masked variant15\.0%15\.0\\%, and the scaffolded variants65\.3%65\.3\\%\(α=0\.3\\alpha=0\.3\),76\.4%76\.4\\%\(α=0\.5\\alpha=0\.5\),73\.0%73\.0\\%\(α=0\.7\\alpha=0\.7\), and76\.8%76\.8\\%\(α=1\.0\\alpha=1\.0\)\. Theα=0\.5\\alpha=0\.5variant essentially matches the non\-masked baseline, mirroring the ID pareto\-optimal behaviour we observed in RQ2\. Under search\-depth shift, however, accuracy collapses across the board:0\.0%0\.0\\%base,0\.0%0\.0\\%fully masked,0\.2%0\.2\\%atα=0\.3\\alpha=0\.3,4\.0%4\.0\\%atα=0\.5\\alpha=0\.5,3\.8%3\.8\\%atα=0\.7\\alpha=0\.7, and9\.2%9\.2\\%atα=1\.0\\alpha=1\.0\. Unlike the self\-distillation Countdown case, the non\-masked variant is the best of the distilled variants under search\-depth shift, but even its accuracy only reaches to9\.2%9\.2\\%\.

### 5\.4Ablation study: Providing different types of scaffolds and training objectives

We investigate whether different scaffolds help the student to fit over the teacher’s distribution and whether, instead of knowledge distillation with reverse KL, supervised fine\-tuning \(SFT\) can achieve similar results\.

Masked Distillation using teaching tokens:Instead of just producing the solution tokens, we prompt the teacher model to produceteaching\-tokenswhich can be used to train a student model along with the solution tokens\. The length of these teaching tokens can be controlled and thus, we use that as lever to control the amount of scaffolding provided to the student\. We provide detailed description of teaching tokens and how they are generated in Appendix[B\.1](https://arxiv.org/html/2607.22629#A2.SS1)\. We perform this ablation in thedual\-modelsetting in the countdown domain\. Our results, refer Table[1\(a\)](https://arxiv.org/html/2607.22629#S5.T1.st1), demonstrate that the scaffolding provided in the form of teacher tokens improves task performance as compared to full masked\-distillation\. However, theα\\alpha\-suffix masked distillationvariants had higher task performance compared to the models trained on teaching tokens\. Thus, while scaffolds generally help, not all are equally effective\.

Finetuning on teacher traces:Our results show that teacher\-token scaffolding does not help under SFT \(refer Table[1\(b\)](https://arxiv.org/html/2607.22629#S5.T1.st2)\. The fully masked SFT student reaches38\.4%38\.4\\%on the in\-distribution split; the100100and10001000\-TT variants reach38\.5%38\.5\\%and30\.8%30\.8\\%respectively\. This is in sharp contrast to the reverse\-KL setting, where the same scaffolding sweep nearly doubles in\-distribution accuracy \(from34\.08%34\.08\\%to67\.25%67\.25\\%/65\.04%65\.04\\%\)\. SFT does not learn to exploit the100100or10001000\-TT scaffold, while reverse\-KL does\.

Table 1:Test accuracy on Countdown\.One in\-distribution \(ID\) split—33–44two\-digit input numbers with a two\-digit target—and three out\-of\-distribution \(OOD\) splits that each perturb a single axis: target\-range shift \(33–44nums, one\-digit target\) and search\-depth shift \(55and66input numbers\)\. Values in parentheses under each header give the sizeNNof that split\.\(a\)Reverse\-KL distillation\.
\(b\)Supervised fine\-tuning \(SFT\)\.

## 6Conclusion

In this work, we study whether the computation expressed in the intermediate tokens can be internalized into the parameters of a language model\. We proposedmasked distillation, a knowledge\-distillation framework in which a student model is trained to predict only the solution tokens conditioned on the question, while the teacher provides feedback on the student’s responses after being conditioned on the question as well as its CoT trace\. By treating intermediate tokens as a scaffold that reasoning models use to fit over the solution tokens, we also studiedα\\alpha\-suffix masked distillationvariants where student internalizes the\(1−α\)\(1\-\\alpha\)*prefix*of teacher’s intermediate trace and is trained to produce the remainingα\\alpha\-fraction*suffix*of the trace at inference, together with solution tokens\. We conduct extensive experiments across two domains, Math and Countdown under both*self\-distillation*and*dual\-model*settings\. Sweeping the suffix\-scaffold parameterα∈\{0,0\.3,0\.5,0\.7,1\.0\}\\alpha\\in\\\{0,0\.3,0\.5,0\.7,1\.0\\\}, our results show that the ability to fully internalize the information present in intermediate tokens is task\-dependent: it succeeds on GSM8K, where base student has prior exposure to the domain, but fails on Countdown\. However, our results demonstrate that theα\\alpha\-suffix scaffold closes this gap at modest cost; for example, applyingα=0\.3\\alpha=0\.3to self\-distillation on Countdown raises accuracy from41\.7%41\.7\\%\(fully masked\) to86\.2%86\.2\\%, essentially matching the teacher \(87\.3%87\.3\\%\) at a∼\\sim1\.3×1\.3\\timesreduction in inference tokens relative to the non\-masked variant\. In out of distribution tasks, models trained under various scaffold paradigms transfer cleanly under target\-range shift, and in the self\-distillation setting they generalize better than the non\-masked variant under search\-depth shift\. These findings establish suffix scaffolding as a controlled axis along which the accuracy vs\. inference\-cost trade\-off can be tuned, with the optimal operating point determined by the task and the teacher\-student capacity gap\.

#### Future Work

All variants in our experiments use the same amount of training data and a similar number of optimization steps, which is a conservative setting for the masked variants given their sparser per\-question supervision\. Matching the effective supervision across variants is left for future work\. In addition, we want to examine whether the amortization tradeoff between the additional training cost via Masked Distillation and the inference\-time savings is favorable\. Finally, the masking objective is agnostic to the source of the intermediate tokens\. The same recipe can be applied when the intermediate signal is provided through environment feedback rather than a teacher model, as in SDPO\[[12](https://arxiv.org/html/2607.22629#bib.bib12)\]\. We plan to investigate whether the source of the intermediate signal \(teacher versus environment\) affects the scaffolding–generalization tradeoff observed in our experiments\.

## Acknowledgments

This research is supported in part by grants from ONR \(N00014\-25\-1\-2301 and N00014\-23\-1\-2409\), DARPA \(HR00112520016\), DoD RAI \(via CMU subcontract 25\-00306\-SUB\-000\), and a generous gift from Qualcomm\. We also thank Cloudexe for the GPU Catalyst Fellowship\.

## References

- Agarwal et al\. \[2024\]Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos Garea, Matthieu Geist, and Olivier Bachem\.On\-policy distillation of language models: Learning from self\-generated mistakes\.In*International Conference on Learning Representations*, volume 2024, pages 21246–21263, 2024\.
- Aggarwal and Welleck \[2025\]Pranjal Aggarwal and Sean Welleck\.L1: Controlling how long a reasoning model thinks with reinforcement learning\.*arXiv preprint arXiv:2503\.04697*, 2025\.
- Arora and Zanette \[2025\]Daman Arora and Andrea Zanette\.Training language models to reason efficiently, 2025\.*URL https://arxiv\. org/abs/2502\.04463*, 2025\.
- Bhambri et al\. \[2025\]Siddhant Bhambri, Upasana Biswas, and Subbarao Kambhampati\.Interpretable traces, unexpected outcomes: Investigating the disconnect in trace\-based knowledge distillation\.*arXiv preprint arXiv:2505\.13792*, 2025\.
- Cobbe et al\. \[2021\]Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman\.Training verifiers to solve math word problems\.*arXiv preprint arXiv:2110\.14168*, 2021\.
- \[6\]Countdown\.Countdown \(game show\)\.[https://en\.wikipedia\.org/wiki/Countdown\_\(game\_show\)](https://en.wikipedia.org/wiki/Countdown_(game_show))\.\[Accessed 13\-05\-2025\]\.
- Deng et al\. \[2023\]Yuntian Deng, Kiran Prasad, Roland Fernandez, Paul Smolensky, Vishrav Chaudhary, and Stuart Shieber\.Implicit chain of thought reasoning via knowledge distillation\.*arXiv preprint arXiv:2311\.01460*, 2023\.
- Deng et al\. \[2024\]Yuntian Deng, Yejin Choi, and Stuart Shieber\.From explicit cot to implicit cot: Learning to internalize cot step by step\.*arXiv preprint arXiv:2405\.14838*, 2024\.
- Gu et al\. \[2024\]Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang\.Minillm: Knowledge distillation of large language models\.In*International Conference on Learning Representations*, volume 2024, pages 32694–32717, 2024\.
- Hinton et al\. \[2015\]Geoffrey Hinton, Oriol Vinyals, and Jeff Dean\.Distilling the knowledge in a neural network\.*arXiv preprint arXiv:1503\.02531*, 2015\.
- Hu et al\. \[2022\]Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen\-Zhu, Yuanzhi Li, Shean Wang, Liang Wang, Weizhu Chen, et al\.Lora: Low\-rank adaptation of large language models\.*Iclr*, 1\(2\):3, 2022\.
- Hübotter et al\. \[2026\]Jonas Hübotter, Frederike Lübeck, Lejs Behric, Anton Baumann, Marco Bagatella, Daniel Marta, Ido Hakimi, Idan Shenfeld, Thomas Kleine Buening, Carlos Guestrin, et al\.Reinforcement learning via self\-distillation\.*arXiv preprint arXiv:2601\.20802*, 2026\.
- Kambhampati et al\. \[2025\]Subbarao Kambhampati, Kaya Stechly, Karthik Valmeekam, Lucas Paul Saldyt, Siddhant Bhambri, Vardhan Palod, Atharva Gundawar, Soumya Rani Samineni, Durgesh Kalwar, and Upasana Biswas\.Stop anthropomorphizing intermediate tokens as reasoning/thinking traces\!In*NeurIPS 2025 Workshop on Bridging Language, Agent, and World Models for Reasoning and Planning*, 2025\.
- Kim and Rush \[2016\]Yoon Kim and Alexander M Rush\.Sequence\-level knowledge distillation\.In*Proceedings of the 2016 conference on empirical methods in natural language processing*, pages 1317–1327, 2016\.
- Lightman et al\. \[2024\]Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe\.Let’s verify step by step\.In*International Conference on Learning Representations*, volume 2024, pages 39578–39601, 2024\.
- Luo et al\. \[2025\]Haotian Luo, Li Shen, Haiying He, Yibo Wang, Shiwei Liu, Wei Li, Naiqiang Tan, Xiaochun Cao, and Dacheng Tao\.O1\-pruner: Length\-harmonizing fine\-tuning for o1\-like reasoning pruning\.*arXiv preprint arXiv:2501\.12570*, 2025\.
- Ma et al\. \[2025\]Xinyin Ma, Guangnian Wan, Runpeng Yu, Gongfan Fang, and Xinchao Wang\.Cot\-valve: Length\-compressible chain\-of\-thought tuning\.In*Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\)*, pages 6025–6035, 2025\.
- on Hugging Face \[2025\]Contributors on Hugging Face\.American invitational mathematics examination 2025 \(aime25\) dataset, 2025\.URL[https://huggingface\.co/datasets/math\-ai/aime25](https://huggingface.co/datasets/math-ai/aime25)\.
- Palod et al\. \[2025\]Vardhan Palod, Karthik Valmeekam, Kaya Stechly, and Subbarao Kambhampati\.Performative thinking? the brittle correlation between cot length and problem complexity\.*arXiv preprint arXiv:2509\.07339*, 2025\.
- Sanh et al\. \[2019\]Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf\.Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter\.*arXiv preprint arXiv:1910\.01108*, 2019\.
- Shenfeld et al\. \[2026\]Idan Shenfeld, Mehul Damani, Jonas Hübotter, and Pulkit Agrawal\.Self\-distillation enables continual learning\.*arXiv preprint arXiv:2601\.19897*, 2026\.
- Sheng et al\. \[2024\]Guangming Sheng, Chi Zhang, Zilingfeng Ye, Xibin Wu, Wang Zhang, Ru Zhang, Yanghua Peng, Haibin Lin, and Chuan Wu\.Hybridflow: A flexible and efficient rlhf framework\.*arXiv preprint arXiv: 2409\.19256*, 2024\.
- Shrivastava et al\. \[2025\]Vaishnavi Shrivastava, Ahmed Awadallah, Vidhisha Balachandran, Shivam Garg, Harkirat Behl, and Dimitris Papailiopoulos\.Sample more to think less: Group filtered policy optimization for concise reasoning\.*arXiv preprint arXiv:2508\.09726*, 2025\.
- Snell et al\. \[2022\]Charlie Snell, Dan Klein, and Ruiqi Zhong\.Learning by distilling context\.*arXiv preprint arXiv:2209\.15189*, 2022\.
- Team et al\. \[2025\]Kimi Team, Angang Du, Bofei Gao, Bowei Xing, Changjiu Jiang, Cheng Chen, Cheng Li, Chenjun Xiao, Chenzhuang Du, Chonghua Liao, et al\.Kimi k1\. 5: Scaling reinforcement learning with llms\.*arXiv preprint arXiv:2501\.12599*, 2025\.
- Valmeekam et al\. \[2025\]Karthik Valmeekam, Kaya Stechly, Vardhan Palod, Atharva Gundawar, and Subbarao Kambhampati\.Beyond semantics: The unreasonable effectiveness of reasonless intermediate tokens\.*arXiv preprint arXiv:2505\.13775*, 2025\.
- Yang et al\. \[2024\]An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong, Haoran Wei, Huan Lin, Jialong Tang, Jialin Wang, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Ma, Jianxin Yang, Jin Xu, Jingren Zhou, Jinze Bai, Jinzheng He, Junyang Lin, Kai Dang, Keming Lu, Keqin Chen, Kexin Yang, Mei Li, Mingfeng Xue, Na Ni, Pei Zhang, Peng Wang, Ru Peng, Rui Men, Ruize Gao, Runji Lin, Shijie Wang, Shuai Bai, Sinan Tan, Tianhang Zhu, Tianhao Li, Tianyu Liu, Wenbin Ge, Xiaodong Deng, Xiaohuan Zhou, Xingzhang Ren, Xinyu Zhang, Xipin Wei, Xuancheng Ren, Xuejing Liu, Yang Fan, Yang Yao, Yichang Zhang, Yu Wan, Yunfei Chu, Yuqiong Liu, Zeyu Cui, Zhenru Zhang, Zhifang Guo, and Zhihao Fan\.Qwen2 technical report, 2024\.URL[https://arxiv\.org/abs/2407\.10671](https://arxiv.org/abs/2407.10671)\.
- Yang et al\. \[2025\]An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al\.Qwen3 technical report\.*arXiv preprint arXiv:2505\.09388*, 2025\.
- Yang et al\. \[2022\]Mengjiao Sherry Yang, Dale Schuurmans, Pieter Abbeel, and Ofir Nachum\.Chain of thought imitation with procedure cloning\.*Advances in Neural Information Processing Systems*, 35:36366–36381, 2022\.

## Appendix AHyperparameter details

In our masked\-distillation experiments, the student uses a maximum prompt length of2,0482\{,\}048tokens and a maximum response length of8,1928\{,\}192tokens\. The teacher uses a longer prompt length of10,00010\{,\}000tokens, since its prompt is augmented with the intermediate tokens \(ITs\) collected in Phase 1\. Each training step uses a question batch size of1616, a mini\-batch size of11, and88rollouts per question\. Rollouts at training time are sampled with the vLLM inference engine at temperature1\.01\.0\. Validation uses one rollout per question with temperature0\.60\.6and top\-p=0\.95p=0\.95\. We optimize with AdamW at a constant learning rate of1×10−51\\times 10^\{\-5\}\.

## Appendix BWhat are teaching tokens?

In*fully masked*distillation, Phase 1 prompts the teacher with the raw input question only, eliciting a response of the form<think\>ITs</think\><answer\>STs</answer\>, from which we extract the intermediate tokens\. With the fully\-masked objective \(Eq\.[4](https://arxiv.org/html/2607.22629#S3.E4)\), the student therefore learns to generate solution tokens only\. For the teaching\-token variants, in Phase 1 we additionally pass an explicit instruction prompt to the teacher, asking it to summarize its own intermediate tokens into a concise scaffold ofk∈\{100,1000\}k\\in\\\{100,1000\\\}tokens \(the*teaching tokens*, TTs\) that the student can use as a hint while producing the solution\. The teacher’s response then takes the form<think\>ITs</think\><teaching\>TTs </teaching\><answer\>STs</answer\>\(the instruction prompt and statistics on the generated teaching tokens are reported in Appendix \.\) During Phase 2 the teacher remains conditioned on the question, the instruction prompt, and the ITs, while the student now learns to emit the teaching tokens followed by the solution tokens\.

### B\.1Prompt templates and teaching\-token statistics

In Phase 1 \(Section[3](https://arxiv.org/html/2607.22629#S3)\), we collect teacher intermediate tokens by prompting the teacher with one of two templates, depending on the distillation variant\. The student receives only the question; the teacher receives the question alone for fully masked distillation \(k=0k=0\), and the question augmented with an*instruction prompt*that elicits teaching tokens for thek∈\{100,1,000\}k\\in\\\{100,1\{,\}000\\\}variants\. Below we show a representative Countdown example\.

Prompt template \(Countdown\)Using the numbers \[10, 8, 27, 9\], create an equation that equals 14\. You can use basic arithmetic operations \(\+, \-, \*, /\) one or multiple times but each number can only be used once\. Return the final equation in<answer\> </answer\>tags, for example<answer\> \(1 \+ 2\) / 3 </answer\>\.You are a teacher\. First, solve the problem\. Then, write a teaching explanation \(approximately\{100,1,000\}\\\{100,\\,1\{,\}000\\\}words\) that captures the key insights and reasoning steps a student would need to solve this problem on their own\. Do not show your full working—distill it into useful teaching points\. Finally, give your final answer in<answer\> </answer\>tags\.Format your response as:<teaching\> \[Your teaching explanation here\] </teaching\> <answer\> \[result\] </answer\>

The unhighlighted text is the question; it is given to the student at training and inference, and to the teacher under fully masked distillation \(k=0k=0\)\. The highlighted block is the instruction prompt; it is appended to the question only for the teaching\-token variants \(k=100k=100andk=1,000k=1\{,\}000\), with the target word count set to100100or1,0001\{,\}000respectively\. Under this prompt, the teacher emits a response of the form<think\>…\\ldotsITs…\\ldots</think\><teaching\>…\\ldotsTTs…\\ldots</teaching\><answer\>…\\ldotsSTs…\\ldots</answer\>; under the question\-only prompt, it emits<think\>…\\ldotsITs…\\ldots</think\><answer\>…\\ldotsSTs…\\ldots</answer\>\.

### Generated teaching\-token length statistics

Table[2](https://arxiv.org/html/2607.22629#A2.T2)reports word\-count statistics for the teaching tokens \(TTs\) generated by the teacher under each instruction\-prompt variant on Countdown\. Both variants undershoot the target word count—the teacher tends to write tighter teaching blocks than asked—but the two distributions remain well separated, and the modal ranges align with the intent of each variant\.

Table 2:Teaching\-token length, in words, generated by the teacher under each instruction\-prompt variant\(Countdown\)\. The “target” column gives the word count requested in the instruction prompt\.

## Appendix CSFT variant

We additionally study a supervised fine\-tuning \(SFT\) variant of masked distillation, in which the student is trained on teacher responses with the standard cross\-entropy loss:

ℒSFT=−𝔼x∼𝒟,y∼πT\(⋅∣x,ITs\)​\[∑t=1\|y\|mt​log⁡πθS​\(yt∣x,y<t\)\],\\mathcal\{L\}\_\{\\mathrm\{SFT\}\}=\-\\,\\mathbb\{E\}\_\{\\,x\\,\\sim\\,\\mathcal\{D\},\\;y\\,\\sim\\,\\pi^\{T\}\(\\,\\cdot\\mid x,\\,\\mathrm\{ITs\}\)\}\\\!\\left\[\\sum\_\{t=1\}^\{\|y\|\}m\_\{t\}\\,\\log\\pi^\{S\}\_\{\\theta\}\(y\_\{t\}\\mid x,y\_\{<t\}\)\\right\],\(7\)wheremt∈\{0,1\}m\_\{t\}\\in\\\{0,1\\\}is a per\-token mask that selects which positions in the teacher response contribute to the loss\. The mask realises the same scaffolding spectrum as the reverse\-KL objective:mt=1m\_\{t\}=1on solution tokens only for the fully masked variant;mt=1m\_\{t\}=1on the \(teaching⊕\\oplussolution\) span for the masked\-kkvariants; andmt=1m\_\{t\}=1on the full \(intermediate⊕\\oplussolution\) span for non\-masked SFT\. We use exactly the teacher rollouts collected in Phase 1, so the only differences between the SFT and reverse\-KL variants are the loss form \(forward vs\. reverse KL on a single sample\) and the source of the gradient \(teacher rollouts vs\. student\-sampled responses\)\.

Hyperparameter details:For the SFT variant on Countdown, we fine\-tune the student, Qwen2\-0\.5B\-Instruct, on solution traces produced by the teacher, Qwen3\-4B operating in reasoning mode\. Teacher traces were generated by sampling at temperature 0\.6 with a maximum generation length of 8192 tokens, and the student was trained in non\-thinking mode\. We use parameter\-efficient fine\-tuning via LoRA\[[11](https://arxiv.org/html/2607.22629#bib.bib11)\], with rank r=64 and scaling factorα=128\\alpha=128\. For each variant, we determine the learning rate, and batch size after a hyperparameter sweep\.

Per\-variant configurations

- •Non\-masked SFT \(thinking\-mode student\):learning rate1×10−41\\times 10^\{\-4\}, batch size44,
- •Masked\-SFT \(1k TTs\):learning rate1×10−41\\times 10^\{\-4\}, batch size44,
- •Masked\-SFT \(100 TTs\):learning rate5×10−45\\times 10^\{\-4\}, batch size44,
- •Masked\-SFT :learning rate1×10−41\\times 10^\{\-4\}, batch size88\.