Self-distillation factors that improve language model adaptation and memory

Disentangling Self-Distillation: Measuring and Modeling Acquisition and Retention

Artificial IntelligenceMachine Learning

Summary

When a language model learns from its own previous answers to get better, it can do this in different ways. The paper explores three key choices in this 'self-teaching' process and shows how each affects how well the model learns new tasks and remembers what it already knew. They find that using the teacher's version to guide more helps with learning hard tasks, adjusting how closely the teacher follows the student can improve learning up to a point, and how differences are measured can affect forgetting differently depending on the model. They tested many setups and built a simplified model that explains these effects.

What this means in practice

  • For machine learning engineers: Improve training routines for adapting large language models to new tasks by selecting optimal self-distillation settings to balance learning and retention.
  • For ai service operators: Enhance ongoing model updates by configuring teacher-student coupling to maintain prior knowledge while incorporating new data effectively.

Authors

Luis Zuin, Alexis Huet, Dario Rossi, Zied Ben Houidi

Abstract

Self-distillation with privileged context adapts a language model from demonstrations by letting the model, once conditioned on a reference response, teach its context-free copy token by token. Our taxonomy reveals existing methods differ along three entangled axes: (i) the rollout source (student or teacher), (ii) the teacher coupling (frozen, or an exponential moving average of the student at some coupling rate) and (iii) the KL direction (reverse or forward), yet these axes are usually studied in fixed combinations and have led to conflicting conclusions. We formalize a unifying framework to encompass all self-distillation methods vs classic supervised fine-tuning: we train every combination of the three axes, on Qwen2.5-7B and Ministral-3-3B across ordinary and contradictory tasks, totaling 1,200 adaptation runs, to systematically investigate the impact of the above axes. We propose a controlled model of the same objective to explain the resulting acquisition-retention trade-offs. We find that (i) the rollout source matters mostly where the task contradicts the pretrained behavior: there teacher rollouts raise acquisition well above what student rollouts achieve, with almost no change in retention; (ii) the teacher coupling changes acquisition most, on every task: acquisition rises with the coupling rate, then falls past a task-specific rate; (iii) switching the KL direction costs retention in one model but not the other so which axis to tune first depends on the model. The controlled model reproduces the three trends.