Teacher-Guided Distillation for a Compact Agent Model
A compact agent must learn broad capabilities and recover from its own imperfect generation paths; Muse Glimmer uses teacher-provided token distributions first during pre-training and then on student-generated rollouts before reinforcement learning.
What it is
- Two places for teacher guidance A larger model supplies token-level teaching signals during the compact model’s pre-training, and later evaluates the compact model’s own generated sequences.
- Learn on your own paths The compact model generates rollouts; the teacher gives feedback along those same paths, so training includes the kinds of imperfect continuations the student will produce at use time.
- A staged handoff The recipe moves from supervised examples to teacher feedback on student rollouts, then to reinforcement learning.
Why it matters
Training only on a teacher’s polished trajectories can leave a student unprepared for its own mistakes. Student-generated rollouts let the teacher provide feedback on the student’s visited paths before reinforcement learning.
Before getting into the training recipe, here is the basic idea: Muse Glimmer learns from a larger model in two different settings. First, the teacher supplies token-level guidance during pre-training; later, the student generates its own sequences and the teacher gives feedback on those sequences. That second step matters because an agent has to continue from what it actually said, not from an idealized teacher trajectory.
-
Transfer: A compact model needs broad language and reasoning ability without learning only from hard yes-or-no labels. The canonical example here is next-token prediction: the teacher can provide a distribution over likely next tokens, not just one answer token. Glimmer’s pre-training uses logit distillation from Muse Spark, retaining its top 256 logits. In this setup, the first 95% uses KL divergence, and the final 5% switches to cross-entropy on ground-truth targets. Put plainly, the training starts by matching the teacher’s retained preferences, then finishes by anchoring the student to the actual target text. The stated reason for that final switch is bias from truncating the teacher distribution. What exactly is lost when only the largest logits are kept?
-
On the student’s own trail: A teacher-generated example can be clean, but the student may take a different turn when it is running. The canonical example is a multi-step tool-use exchange: if the student’s response changes the next prompt or tool result, the rest of its sequence follows from that changed history. During on-policy distillation, Glimmer samples rollouts and Muse Spark evaluates the student’s generated distributions step by step on those visited prefixes. Think of it as a coach responding to the route you actually took, rather than handing you notes for someone else’s route; the analogy stops short because the teacher supplies token distributions, not conversational advice. The stated goal is to reduce train–inference distribution shift and help the student recover when its path drifts. How should that fit with ordinary supervised examples and reinforcement learning?
-
A staged handoff: The post-training sequence is supervised fine-tuning, on-policy distillation, then multi-objective reinforcement learning. In practical terms, SFT teaches the expected response format and instruction-following behavior; the middle stage brings teacher feedback to student-generated paths; RL then optimizes against reward models, verifiers, and task-success signals. The reported comparison favors including the on-policy stage over moving directly from SFT to RL, but the passage gives no numeric effect size. My read is that the stages address different failure modes rather than replace one another. Which agent behaviors are supplied by the surrounding data and training recipe is a separate question.
Read the deep dive
Why People Wanted It
A compact agent has to do more than imitate a polished answer. It must reason, use tools, and keep going when its own next-token choices lead somewhere imperfect. The real question is how to transfer a larger model’s behavior without training only on the larger model’s preferred paths—and how to prepare the student for the later reinforcement-learning stage.
The Core Move
Muse Glimmer uses Muse Spark’s token distributions to teach the student during pre-training, then asks the teacher to evaluate trajectories the student generates during post-training. In the reported pre-training setup, the teacher’s top 256 logits are retained; KL divergence is used for the first 95% of pre-training, followed by ground-truth cross-entropy alone for the final 5%.
Why the Final Cross-Entropy Phase Addresses Truncation Bias
Here is the part that gets technical: the transition is specifically motivated by bias from retaining only the top- teacher logits, rather than a general claim that one loss is always better. For most of pre-training, the objective is KL divergence over the co-distilled teacher signal. For the final portion, the objective switches to cross-entropy against ground-truth targets, with no continuing blend of soft and hard targets described for that phase. The quoted method gives the schedule and target change, but no explicit loss equation, so there is no formula to reproduce here.
Why Student-Generated Prefixes Change the Teacher’s Job
Now the version for readers who want more of the machinery: during on-policy distillation, the student samples autoregressive rollouts, and the teacher evaluates the student’s generated distributions step by step on those visited prefixes. The distinction is the prefix distribution: the teacher is not only supervising trajectories it generated itself; it supplies per-token distributional targets along paths the student actually took. That is the mechanism behind the stated aim of addressing train–inference distribution shift and helping the student recover when its generation path drifts.
Why Distillation Sits Between SFT and RL
Here is the training-sequence detail: post-training is described as supervised fine-tuning, on-policy distillation, and multi-objective reinforcement learning. SFT aligns the base checkpoint to instruction-following, formatting, and reasoning-style targets; the on-policy stage then trains against teacher feedback on student rollouts; RL optimizes toward reward models, verifiers, and task-success signals. The reported comparison says the intermediate on-policy stage improves downstream RL outcomes relative to moving directly from SFT to RL. The source does not give a numeric effect size here.
Minimal Causal Code
This loop is schematic: it exposes the student-rollout and teacher-feedback dependency, not the exact production loss implementation. The paper describes the stages but does not specify a complete executable loss formula.
for batch in prompts:
student_rollout = student.sample(batch)
teacher_targets = teacher.token_distributions(
prefixes=student_rollout.prefixes
)
loss = distillation_loss(
student.distributions(student_rollout.prefixes),
teacher_targets,
)
loss.backward()
optimizer.step()
optimizer.zero_grad()
What You Gain, and What You Lose
The approach gives the compact student teacher-provided distributional supervision during pre-training and later feedback on its own visited prefixes. The reported intermediate stage improves downstream RL outcomes compared with going directly from SFT to RL, and the stated rationale is to reduce train–inference distribution shift and teach recovery from drift.
There are tradeoffs. Keeping only the teacher’s top 256 logits truncates its distribution, which is why the pre-training schedule ends with ground-truth cross-entropy. On-policy feedback also requires the student to generate rollouts and the teacher to evaluate them step by step. Finally, the method is a sequence of stages, not one interchangeable objective: broad pre-training distillation, SFT, student-rollout distillation, and RL each have distinct targets. How much of the result comes from the intermediate stage versus the broader training recipe is not quantified in these passages.
Teach the student from the teacher’s distributions, then make the teacher meet the student on the paths the student actually takes.