Self-Distillation Fine-Tuning for Continual Learning

ML Training & Fine-Tuning advanced How AI Works 32 min read

Overview

Fine-tuning a language model on a new skill is easy to start and hard to keep stable. Standard supervised fine-tuning (SFT) trains the model to imitate a fixed dataset of expert demonstrations. That can improve the target task, but it also changes the model under an offline distribution that may not match how the model behaves at inference time. The result is a familiar continual-learning failure mode: the model learns the new behavior while forgetting earlier capabilities, becoming brittle outside the demonstration format, or collapsing reasoning traces when the training examples contain only short final answers.

Self-Distillation Fine-Tuning (SDFT) is a post-training method designed for the regime where demonstrations exist but explicit reward functions do not. Instead of treating the demonstration as a static target sequence, SDFT uses the model’s own in-context learning ability to create a demonstration-conditioned teacher. The student samples responses from its current policy, while a frozen or moving-average teacher sees the same prompt plus a relevant expert demonstration. Training then distills the teacher distribution into the student on the student’s own trajectories.

The key shift is from offline imitation to on-policy distillation. The model is still learning from demonstrations, but the correction signal is computed where the student actually goes, not only where the dataset says an expert went. In the SDFT paper, this lets a single model acquire science Q&A, tool-use, medical reasoning, and new factual knowledge while preserving general benchmark performance better than standard SFT.

SDFT belongs in the fine-tuning strategy toolbox alongside SFT, preference optimization, and reinforcement learning. It is not a drop-in replacement for every adapter run: it requires online generation during training, teacher scoring, and a base model strong enough to use demonstrations in context. But when a team needs continual skill acquisition from demonstrations and cannot write a reliable reward function, SDFT offers a practical middle ground between plain SFT and full reinforcement learning.

Core Concepts

Off-Policy Imitation vs On-Policy Learning

SFT is off-policy because the model is optimized on trajectories from the dataset, not trajectories generated by the model being trained. This matters because errors compound. A student may produce a partial answer that differs from the expert prefix, then receive no useful supervision for recovering from that state because the dataset never visits it. The larger the mismatch between training states and inference states, the more fragile the learned behavior becomes.

On-policy learning trains under the model’s own induced distribution. Reinforcement learning is the common version: sample completions, score them with a reward, and update the policy. SDFT keeps the on-policy sampling idea but replaces the external reward with a teacher distribution derived from the same model conditioned on demonstrations.

Demonstration-Conditioned Teacher

The teacher is not a separate larger model. It is the base or exponential-moving-average model evaluated with extra context: the task prompt plus an expert demonstration. If the model has strong in-context learning, that demonstration shifts the output distribution toward task-appropriate reasoning while remaining close to the original model’s style. This closeness is important because large deviations from the pretrained distribution are associated with forgetting and brittle specialization.

The teacher acts like a local, prompt-specific policy improvement operator. It does not merely copy the demonstration; it interprets the demonstration as evidence about the task and produces a response distribution for the student’s sampled trajectory.

Reverse-KL Distillation on Student Rollouts

Training minimizes divergence between the student’s token distribution and the demonstration-conditioned teacher distribution on responses sampled by the student. The SDFT paper frames this as reverse-KL distillation, with gradients estimated over the student’s current trajectories. Because the student supplies the rollouts, the update is on-policy. Because the teacher supplies token-level probabilities, the learning signal is denser than a single scalar trajectory reward.

Key Concept: Demonstrations Become Policy Improvement Signals

SDFT turns each demonstration into a contextual teacher rather than a string to imitate. The student asks, “given the state I actually reached, how would my demonstration-aware version distribute probability over next tokens?” That makes the demonstration useful even when no reward function exists and even when the demonstration omits the full reasoning trace.

How It Works

Construct a Teacher Prompt Per Training Example

For each query, SDFT pairs the prompt with an expert demonstration. The demonstration can be an answer, a solved example, a tool-use trace, or factual evidence transformed into question-answer form. The teacher prompt must make the demonstration informative without forcing verbatim copying. The paper uses a simple pattern that asks the model to answer a new question after seeing an example response.

# Conceptual teacher prompt shape. The delimiter names are placeholders,
# not a required API.
teacher_context = """
Question:
{query}

This is an example for a response to the question:
{demonstration}

Now answer with a response of your own, including the thinking process:
"""

student_context = """
Question:
{query}
"""

The student sees only the query, matching the deployment setting. The teacher sees the query plus the demonstration. During training, the two policies are evaluated over the same student-generated tokens so their probability distributions can be compared token by token.

Generate Student Rollouts

The training loop samples one or more completions from the student policy. This is the expensive part relative to SFT: the model must generate online rather than simply teacher-force over a dataset. Those rollouts represent the states where the deployed model is likely to operate, including its partial mistakes and stylistic tendencies.

# Pseudocode for one SDFT-style step.
for batch in train_loader:
    queries = batch["query"]
    demos = batch["demonstration"]

    # 1. Sample from the current student policy.
    sampled_tokens = student.generate(
        prompts=queries,
        max_new_tokens=max_tokens,
        temperature=sample_temperature,
    )

    # 2. Score the sampled tokens under the student context.
    student_logprobs = student.logprobs(
        prompts=queries,
        completions=sampled_tokens,
    )

    # 3. Score the same sampled tokens under a demonstration-aware teacher.
    teacher_prompts = build_teacher_prompts(queries, demos)
    with no_grad():
        teacher_logprobs = teacher.logprobs(
            prompts=teacher_prompts,
            completions=sampled_tokens,
        )

    # 4. Update the student toward the teacher distribution on its own rollouts.
    loss = reverse_kl(student_logprobs, teacher_logprobs)
    loss.backward()
    optimizer.step()
    update_ema_teacher(student, teacher)

Use a Stable Teacher Parameterization

The teacher can be parameterized in several ways. A fixed base teacher is simple, but it may lag as the student improves. The SDFT paper’s default uses an exponential moving average (EMA) of the student weights for the teacher while still conditioning the teacher on demonstrations. EMA gives the teacher enough stability to avoid chasing every noisy student update while keeping it near the policy being optimized.

Evaluate Both New Skill and Retained Capabilities

SDFT is a continual-learning method, so target-task accuracy alone is insufficient. The paper evaluates new-task accuracy, previous-capability retention across general benchmarks, out-of-distribution reasoning for knowledge acquisition, and sequential learning across multiple skills. A production implementation should copy that evaluation shape: test whether the model learned the new capability and whether it preserved unrelated behavior that users already rely on.

Practical Applications

Continual Skill Acquisition

A model that first learns tool-use routing, then medical reasoning, then a specialized scientific task should not require a separate adapter for every skill if the goal is one general assistant. SDFT is useful when a team wants a single model to accumulate skills sequentially. The student can be trained on each new demonstration set while retention suites monitor previously learned domains and broad general reasoning.

Knowledge Injection Without Reward Engineering

Some updates are about factual integration rather than a reusable skill. A company might need a model to internalize new procedures, product details, or incident-history facts that are not in pretraining. RAG can expose facts at inference time, but it does not update model parameters. Plain SFT may memorize question-answer pairs without using the facts compositionally. SDFT targets this gap by using demonstrations to shape an on-policy learning signal and testing indirect questions that require reasoning over the injected knowledge.

Reasoning Adaptation from Answer-Only Data

Many useful datasets contain final answers but not full reasoning traces. Direct SFT on short answers can teach a reasoning model to shorten or suppress its internal problem-solving behavior. The SDFT paper reports an answer-only medical reasoning setting where standard SFT reduced performance and shortened responses, while SDFT improved accuracy and preserved longer reasoning behavior. The mechanism is that the demonstration-conditioned teacher reconstructs a model-native reasoning style instead of forcing the student to imitate abbreviated targets.

Initialization Before Reinforcement Learning

SDFT can also serve as a pre-RL stage. If SDFT moves the model toward better high-probability generations using dense token-level feedback, a later RL phase may start from a stronger policy and need less exploration. This is especially attractive when reward functions are available only for a subset of tasks, while demonstrations cover a broader domain.

Implementation Patterns

Batch Around Shared Teacher and Student Scoring

An SDFT trainer has three expensive operations: student generation, student log-probability scoring, and teacher log-probability scoring. The implementation should batch by prompt length and expected completion length so the same sampled tokens can be scored efficiently under both contexts. Store tokenized student prompts, teacher prompts, and completion masks explicitly; otherwise it is easy to compare log probabilities over misaligned token positions.

def masked_reverse_kl(student_logp, teacher_logp, completion_mask):
    # Both tensors are [batch, tokens, vocab] log-probabilities for the
    # same completion positions. The mask excludes prompt and padding tokens.
    student_p = student_logp.exp()
    token_kl = student_p * (student_logp - teacher_logp)
    token_kl = token_kl.sum(dim=-1)
    return (token_kl * completion_mask).sum() / completion_mask.sum()

Keep Demonstration Retrieval Separate from Training

SDFT assumes a relevant demonstration for each query. In a small curated dataset, the pair may be stored directly. In a larger corpus, demonstration selection becomes a retrieval problem: choose an example that expresses the desired task behavior without leaking the answer to the specific query. For tool use, this might mean same API family but different user request. For knowledge acquisition, it might mean a fact-grounded example from the same source corpus but not the exact held-out question.

Use EMA for Teacher Stability

When the teacher tracks the student too closely, noisy student updates can destabilize the target. When the teacher is fixed forever, it may become stale after several skills. EMA is a pragmatic compromise. The teacher weights are updated slowly from the student after optimizer steps, while gradients never flow through teacher scoring.

def update_ema_teacher(student, teacher, decay=0.995):
    for student_param, teacher_param in zip(student.parameters(), teacher.parameters()):
        teacher_param.data.mul_(decay)
        teacher_param.data.add_(student_param.data, alpha=1.0 - decay)

Design Retention Gates Before Training

SDFT’s value is retention, so a training pipeline should define retention gates before the first run. Include a fixed set of general benchmarks, representative production prompts, previously learned tasks, and refusal/safety probes if those behaviors matter. Checkpoint selection should optimize a Pareto frontier: target skill gain, previous-capability retention, and compute cost.

Performance Considerations

Common Pitfalls

Further Reading