Self-Distillation Enables Continual Learning — part 3/4

Cleaned-text code points [30962, 46658). Bibliography omitted here; source.md includes it.

Figure 5 (left) illustrates a clear trend. At small scales, such as the 3B variant, the model’s in-context learning is too weak to provide meaningful teacher guidance, and performance lags behind standard SFT. However, as we increase the size, the gains from our method grow consistently. The 7B model achieves a four-point improvement over SFT, and the 14B model widens the margin to seven points. This monotonic improvement suggests that the effectiveness of our algorithm is tightly coupled to the model’s ability to perform in-context reasoning, and, as larger models continue to exhibit stronger ICL behavior, our approach is likely to become even more advantageous at larger scales.

### 4.5 Training reasoning models without reasoning data

A major practical challenge in post-training reasoning models is the construction of high-quality supervision. Supervised fine-tuning for reasoning models typically requires access to intermediate reasoning traces, which are expensive to collect from human annotators and often unavailable from closed-source models that expose only final answers. As a result, many real-world datasets provide only the final answer, without the full chain of thought.

Naively applying SFT on such data can be harmful for reasoning-capable models. Because SFT directly matches the target responses, when demonstrations contain only short answers or abbreviated reasoning, this suppresses the model’s existing long chain-of-thought behavior. For example, consider a model that reliably produces long chains of thought, but is post-trained using demonstrations that provide only concise solutions. Direct SFT penalizes long reasoning traces in favor of short outputs that better match the supervision, causing a collapse in reasoning depth. We hypothesize that our on-policy self-distillation framework avoids this failure mode. Because the student policy is trained to match the outputs of a demonstration-conditioned teacher derived from the same model, the supervision preserves the model’s internal reasoning style even when the external data contains only final answers. In effect, the teacher induces a reasoning-consistent target distribution, allowing the model to adapt without collapsing its reasoning behavior.

We evaluate this hypothesis using Olmo-3-7B-Think ( Olmo et al., 2025 ) and fine-tune it on the medical task described earlier. Importantly, this dataset contains no explicit chain-of-thought annotations. We compare standard SFT against our method, measuring both task accuracy and the average number of generated tokens, which serves as a proxy for retained reasoning depth.

#### Results.

Table 2 presents the results. Standard SFT substantially degrades performance, reducing accuracy from 31.2% to 23.5% and sharply shortening responses, indicating a collapse in reasoning behavior. In contrast, our method significantly improves accuracy, reaching 43.7%. These results demonstrate that our approach enables effective task adaptation for reasoning models even in the absence of explicit reasoning data.

### 4.6 What drives the improvement in performance?

Our method combines two ingredients: a demonstration-conditioned teacher policy and an on-policy distillation objective. In Subsection 3.2 , we validated that the conditioned model constitutes a high-quality teacher—producing correct outputs while remaining close to the base policy. A natural question then arises: if such a teacher already exists, is on-policy learning necessary, or would standard distillation suffice?

To isolate the source of the performance gains, we compare our full algorithm against two alternative ways of leveraging the same teacher: (1) SFT from the teacher, where the student is trained offline to imitate samples generated by the teacher ( Yang et al., 2024 ; Cao et al., 2025 ) . (2) Offline distillation from the teacher, where the student minimizes a KL loss on a fixed dataset of teacher-generated outputs ( Padmanabhan et al., 2023 ; Mitra and Ulukus, 2025 ) . We use the Tool Use task for the comparison.

#### Results.

As shown in Figure 6 , neither form of offline distillation matches the performance of our on-policy approach. While distillation from the teacher improves over standard SFT, it consistently underperforms our method. This gap indicates that the benefits of SDFT cannot be attributed solely to the quality of the teacher and further highlights the importance of on-policy learning.

## 5 Discussion and Limitations

#### Relationship to on-policy RL.

SDFT is not an alternative to on-policy RL, but addresses a complementary learning regime. The two methods operate under different supervision assumptions: SDFT applies when learning from expert demonstrations without access to an explicit reward function, whereas on-policy RL assumes a reward signal and optimizes expected return through exploration.

Importantly, the two approaches can be naturally combined. As shown in Figure 5 , SDFT consistently improves pass@ k k across a wide range of k k , indicating that it increases the diversity and quality of high-probability generations. This suggests that SDFT can serve as an effective initialization for subsequent RL fine-tuning, providing a stronger starting policy.

If one nevertheless compares the training dynamics, a practical advantage of SDFT is its efficiency. Our method requires only a single on-policy generation per prompt, whereas many on-policy RL methods, such as GRPO, rely on group-based sampling to estimate relative advantages, substantially increasing generation cost. Moreover, our supervision is provided at the token level (or even the logit level), providing denser credit assignment than the trajectory-level advantage used in GRPO.

#### Computational Costs.

A practical consideration when adopting SDFT is that, unlike standard SFT, it requires generating on-policy rollouts during training. In our experiments, this results in approximately 2.5× the computational cost in FLOPs and roughly 4× the wall-clock training time compared to SFT. However, this additional cost must be viewed in context: many existing continual learning approaches, such as Re-invoke, require sequential training stages—first performing SFT, then conducting additional on-policy training to restore degraded capabilities. When accounting for this multi-stage process, SDFT can actually reduce total training time while simultaneously achieving better performance on the new tasks.

#### Learned Artifacts.

A subtle failure mode of our approach is that the student can inherit spurious linguistic patterns from the teacher. Because the teacher is conditioned on demonstrations or text passages, it may produce responses prefaced with phrases like ”Based on the text…” or ”Following the example…” The student, although receiving no such context, sometimes nevertheless reproduces these markers, having learned them as part of the teacher’s output distribution. Empirically, we find that masking the loss over the first few tokens during training effectively suppresses these artifacts without harming downstream accuracy. While this workaround is effective in practice, it is fundamentally a heuristic fix. A more principled solution remains an open problem.

#### Requirements for Model Capability.

The effectiveness of SDFT depends critically on the base model’s in-context learning capabilities. As demonstrated in our scaling experiments (Figure 4a), smaller models with weak ICL abilities fail to provide meaningful teacher signals. Additionally, the fundamental design of our algorithm imposes constraints on the types of adaptations it can support. SDFT excels at acquiring new skills or knowledge while preserving existing capabilities, but struggles when the desired behavior requires a fundamental shift in the model’s generation patterns. For instance, we found that transforming a non-reasoning model into one that produces explicit chain-of-thought traces proved difficult. Future work should explore modified objectives or prompting techniques to enable more aggressive behavioral changes when necessary.

#### Future Work.

Several promising directions remain for extending SDFT. First, our method can be naturally combined with on-policy RL, either sequentially by using SDFT as an initialization before reward-based fine-tuning, or simultaneously by blending demonstration-based and reward-based learning signals. Second, although on-policy learning substantially reduces catastrophic forgetting compared to off-policy methods, some degradation of prior capabilities remains. Developing complementary techniques that further minimize forgetting represents an important avenue for future research. Finally, while we focus on expert demonstrations, extending SDFT to learn from non-expert or noisy demonstrations, or from unstructured data such as user conversations, would broaden its applicability to real-world continual learning settings where high-quality supervision is scarce.

## Author Contributions

Idan Shenfeld Lead the project and contributed to all aspects of experiments and writing.

Mehul Damani supported this project, made significant contributions to the design of experiments in discussions and gave valuable advice on presentation and writing.

Jonas Hübotter supported this project, made significant contributions to the core algorithmic design, and experimental setting. Gave valuable advice on writing.

Pulkit Agrawal co-developed the project direction, advised IS, and played a significant role in paper writing.

## Acknowledgment

We want to express our gratitude to Nitish Dashora, Jyo Pari, Isha Puri, Zhang-Wei Hong, and members of the Improbable AI lab for the helpful discussion on the paper. We also thank Nofit Segal, Dan Haramati, Yael Vinker, and Assaf Ben-Kish for their support. We are grateful to MIT Supercloud and the Lincoln Laboratory Supercomputing Center for providing HPC resources. The research was supported in part by Hyundai Motor Company, Google, Amazon and MIT-IBM Watson AI Lab. The research was sponsored by the Army Research Office and was accomplished under Grant Number W911NF-21-1-0328 and an ARO MURI under Grant Number W911NF-23-1-0277. The research was also sponsored by the Office of Naval Research and was accomplished under Grant Number N00014-22-1-2740. Research was also sponsored by the Department of the Air Force Artificial Intelligence Accelerator and was accomplished under Cooperative Agreement Number FA8750-19-2-1000. The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the Department of the Air Force or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein. The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the Army Research Office, Naval Research Office, Air Force, or the U.S. Government. JH was supported by the Swiss National Science Foundation under NCCR Automation, grant agreement 51NF40 180545.

## Appendix A Additional Ablations

### A.1 Estimating the KL gradient

A central component of our objective is the gradient of the KL divergence between the current policy π θ ​ ( y | x ) \pi_{\theta}(y|x) and the teacher policy π ⁡ ( y | x , c ) \pi(y|x,c) . For sequence models, the KL divergence is defined at the sequence level as: KL ( π θ | | π ) = 𝔼 y ∼ π θ [ log π θ ​ ( y | x ) π ⁡ ( y | x , c ) ] \mathrm{KL}(\pi_{\theta}||\pi)=\mathbb{E}_{y\sim\pi_{\theta}}\left[\log\frac{\pi_{\theta}(y|x)}{\pi(y|x,c)}\right] where y = ( y 1 , … , y T ) y=(y_{1},\dots,y_{T}) is generated autoregressively. Differentiating this quantity is non-trivial because π θ \pi_{\theta} appears both in the sampling distribution and inside the logarithm, and different practical estimators trade off bias, variance, and computational cost ( Tang and Munos, 2025 ) .

We consider and ablate several commonly used KL gradient estimators.

#### Token-level (partial) estimator

. A widely used approximation decomposes the KL into token-level terms and differentiates each independently: g ^ token = ∑ t = 1 T log ⁡ π θ ​ ( y t | y < t , x ) π ⁡ ( y t | y < t , x , c ) ​ ∇ θ ​ log ​ π θ ​ ( y t | y < t , x ) \widehat{g}_{\text{token}}=\sum_{t=1}^{T}\log\frac{\pi_{\theta}(y_{t}|y_{<t},x)}{\pi(y_{t}|y_{<t},x,c)}\nabla_{\theta}\log\pi_{\theta}(y_{t}|y_{<t},x) As shown in recent analyses, this estimator corresponds to a partial derivative of the sequence-level KL: it ignores the effect of early tokens on future token distributions and is therefore biased with respect to the true gradient ( Tang and Munos, 2025 ) .

#### Full analytic per-token estimator.

An alternative is to compute the KL analytically at each timestep by marginalizing over the vocabulary: g ^ analytic = ∑ t = 1 T ∑ v ∈ 𝒱 log ⁡ π θ ​ ( v | y < t , x ) π ⁡ ( v | y < t , x , c ) ​ ∇ θ ​ log ​ π θ ​ ( v | y < t , x ) \widehat{g}_{\text{analytic}}=\sum_{t=1}^{T}\sum_{v\in\mathcal{V}}\log\frac{\pi_{\theta}(v|y_{<t},x)}{\pi(v|y_{<t},x,c)}\nabla_{\theta}\log\pi_{\theta}(v|y_{<t},x) This estimator has strictly lower variance than sample-based token estimators, but it remains biased at the sequence level, since it still does not account for how the choice of y t y_{t} influences future states y > t y_{>t} . Despite this bias, it is often computationally attractive because it leverages quantities already produced during the forward pass.

#### Rao-Blackwellized estimator.

Following recent work ( Amini et al., 2025 ) , one can further reduce variance by Rao-Blackwellizing the KL estimator, analytically integrating over next-token distributions while retaining Monte-Carlo sampling over prefixes. This yields an unbiased estimator of both the KL and its gradient with provably lower variance than standard Monte-Carlo estimators . However, this estimator is more expensive to compute.

g ^ rb = ∑ t = 1 T [ ∑ v ∈ 𝒱 log ⁡ π θ ​ ( v | y < t , x ) π ⁡ ( v | y < t , x , c ) ​ ∇ θ ​ log ​ π θ ​ ( v | y < t , x ) + k θ ​ ( y < t ) ​ ∑ i = 1 t − 1 ∇ θ ​ log ​ π θ ​ ( y i | y < i , x ) ] \widehat{g}_{\text{rb}}=\sum_{t=1}^{T}\left[\sum_{v\in\mathcal{V}}\log\frac{\pi_{\theta}(v|y_{<t},x)}{\pi(v|y_{<t},x,c)}\nabla_{\theta}\log\pi_{\theta}(v|y_{<t},x)+k_{\theta}(y_{<t})\sum_{i=1}^{t-1}\nabla_{\theta}\log\pi_{\theta}(y_{i}|y_{<i},x)\right]

Where k θ ​ ( y < t ) k_{\theta}(y_{<t}) is the stepwise KL term KL ( π θ ( ⋅ | y < t , x ) | | π ( ⋅ | y < t , x , x ) ) \mathrm{KL}(\pi_{\theta}(\cdot|y_{<t},x)||\pi(\cdot|y_{<t},x,x)) .

We empirically ablate all three estimators in our training pipeline. Despite its theoretical bias, we find that the full analytic per-token estimator consistently yields the most stable optimization and best downstream performance. In contrast, the token-level estimator exhibits higher variance and weaker KL control, while the Rao–Blackwellized estimator did not provide measurable gains in our setting relative to its additional complexity.

We also experimented with drawing multiple trajectories per prompt to reduce variance in the gradient estimator, which is theoretically beneficial for Monte-Carlo estimates. In practice, however, increasing the number of samples per prompt produced negligible improvements while substantially increasing compute. As a result, we adopt a single-trajectory-per-prompt setup combined with the analytic per-token KL estimator in all main experiments.

### A.2 The Importance of Demonstration-Conditioned Context