Introduction

Multi-Teacher On-Policy Distillation separates two difficult engineering jobs in post-training: developing strong domain specialists and integrating their capabilities into a single general-purpose model. Rather than copying static teacher completions, the student practices on its own generated trajectories while the relevant domain specialist supplies dense next-token supervision along those visited states. That setup raises immediate questions about probability distributions, state-visitation shift, gradient estimators, teacher compatibility, and serving compute.

This 26-chapter technical book works through those questions from first principles, combining full derivations, benchmark tables, worked numerical traces, implementation pseudocode, an evidence audit, a complete notation reference, and worked exercises. Each chapter is self-contained enough to read on its own while building a continuous mathematical and systems narrative from token probabilities to frontier-scale deployment.

Two papers, two meanings of MOPD

Two distinct papers share the acronym MOPD: Wenhan Ma and coauthors' Multi-Teacher On-Policy Distillation (arXiv:2606.30406v1) and Weichen Yu and coauthors' Multi-Rollout On-Policy Distillation (arXiv:2605.12652v2). Throughout this book, we refer to them as MT-MOPD and MR-MOPD.

Both methods start from on-policy distillation—sampling rollouts from the student and evaluating a teacher's next-token distributions along the resulting prefixes—but they modify orthogonal parts of the supervision pipeline. MT-MOPD changes which teacher supplies the target by routing each prompt to its domain specialist. MR-MOPD changes what the teacher sees by placing verified passing and failing peer rollouts on the same prompt into the teacher's context as privileged information. This book centers on MT-MOPD while giving MR-MOPD a complete parallel treatment, covering its objective, peer-selection combinatorics, empirical results, and manuscript inconsistencies.

Intended audience and prerequisites

This book is written for engineers and researchers who want to analyze, audit, or implement on-policy distillation pipelines. Familiarity with standard deep learning helps, though the opening chapters build up tokenizers, autoregressive policies, supervised fine-tuning (SFT), reinforcement learning with verifiable rewards, Group Relative Policy Optimization (GRPO), KL divergences, and teacher forcing from the ground up.

Several recurring analytical distinctions run through the book:

• The policy-gradient update used in sampled-token distillation differentiates the student's action distribution at a fixed prefix, which drops the harder trajectory-level path derivative through future visited states.
• The corrected top-$k$ objective restores the exact optimum on the selected vocabulary coordinates by folding unselected tail mass into a remainder bucket, without recovering information about relative preferences inside that tail.
• MT-MOPD's headline normalized score measures how much of the specialist-over-SFT headroom the integrated student captures, rather than raw accuracy gain.
• MR-MOPD's manuscript contains a discrepant GPQA Chemistry baseline between tables, edge cases where prompt-level ranking metrics become undefined on all-pass or all-fail rollout groups, and unspecified fallback rules when peer pools are empty or exceed context limits. We document these points with exact source locators in the evidence and audit chapters.

All numerical examples use natural logarithms (divergences in nats) and decimal storage units (KB/MB/GB) unless binary units (MiB) are explicitly noted.

Contents

Reading paths

• Complete sequential reading (Chapters 1–26): Chapters 1–5 build the probabilistic and RL foundations, Chapters 6–14 cover MT-MOPD, Chapters 15–19 cover MR-MOPD, and Chapters 20–26 synthesize both through worked traces, system blueprints, evaluation design, theoretical bounds, and reference material.
• Mathematical foundations (Chapters 2, 4–5, 7–9, 20, 23, 26): Focuses on KL directionality, fixed-context versus full trajectory-distribution gradients, importance-weighted policy gradients, and the corrected top-$k$ objective.
• Systems and training pipelines (Chapters 3, 6, 8–12, 15–17, 21–22): Covers specialist RL decoupling, asynchronous teacher prefill services, loss masking, rollout-group combinatorics, and end-to-end cluster throughput.
• Empirical results and source audit (Chapters 1, 13–14, 18–19, 22–24): Examines the Qwen3-32B and MiMo-V2-Flash integration benchmarks, MR-MOPD's pass@$k$ and signal-quality evaluations, and the manuscript audit.
• Fast conceptual overview (Chapters 1, 4, 6, 15, 20, 24): Traces what each algorithm modifies in the distillation loop and walks through concrete token-level updates before diving into full derivations.

Begin with Chapter 1.

1. Two Papers, One Acronym, Two Problems

A production language model has to handle very different tasks inside a single set of weights: proving mathematical claims, satisfying brittle formatting constraints, patching multi-file software repositories, and invoking external tools with well-typed arguments. Training a model to excel at any one of those domains is already demanding, but combining them into a single checkpoint introduces a distinct engineering challenge: capability integration.

Multi-Teacher On-Policy Distillation addresses capability integration by decoupling skill acquisition from skill consolidation. Domain teams first train dedicated specialists using domain-specific reinforcement learning recipes. A single generalist student then generates its own trajectories across all domains while the corresponding specialist provides dense next-token probability targets along the student's visited states. Because the student generates the prefixes on which it is supervised, it learns how to recover and complete reasoning paths it actually takes at inference time.

The naming collision

Two contemporary papers share the acronym MOPD while tackling different parts of post-training:

We use the labels MT-MOPD and MR-MOPD throughout this book to keep the two experimental programs distinct (MT paper, MR paper); whenever a chapter refers simply to "multi-teacher distillation," it means the first paper.

Both methods build on on-policy distillation, where the training loop samples a trajectory from the current student policy and queries a teacher for next-token distributions at every prefix along that trajectory. Where they diverge is how the teacher target is constructed: MT-MOPD changes which teacher scores the trajectory by routing each prompt to a domain specialist, whereas MR-MOPD changes what context the teacher sees by conditioning the teacher on verified peer rollouts from the same prompt.

Why specialist improvements do not simply add together

Suppose we start from a shared supervised fine-tuned checkpoint and train three separate copies: one on competition mathematics with symbolic answer verification, a second on instruction-following prompts with programmatic constraint checkers, and a third on repository repair inside containerized test harnesses. Each specialist improves substantially on its target benchmark.

Combining those three checkpoints into one deployable model exposes immediate trade-offs among standard post-training strategies:

• Parameter merging: Averaging the weights of the three specialists requires no extra training, but neural network parameters do not partition cleanly by domain. Because the same attention heads and MLP projections participate across tasks, linear interpolation in weight space often degrades nonlinear circuit behavior that each specialist relied on.
• Joint multi-domain reinforcement learning (Mix-RL): Training a single model on all three reward functions simultaneously forces heterogeneous workloads into one optimization loop. A learning rate, rollout length, or KL penalty tuned for 32-thousand-token math proofs can destabilize short formatting prompts, while container-bound software engineering rollouts that take minutes to execute stall synchronous batch collection for fast text-only tasks.
• Sequential reinforcement learning (Cascade RL): Training on math first, then instruction following, and finally software engineering avoids multi-environment synchronization, but later stages overwrite earlier gains unless carefully regularized. Moreover, final performance becomes sensitive to the arbitrary ordering of domains.

Integrating behavior on student-visited states

To teach a student effectively, a specialist does not need to generate the demonstration from scratch. Suppose the student is halfway through a flawed mathematical derivation. When the math teacher evaluates that exact student prefix under teacher forcing, it assigns low probability to continuations that compound the error and higher probability to tokens that pivot toward a valid proof.

By querying the teacher on prefixes the student actually produced, the training loop delivers dense token-level corrections precisely along the student's current state visitation distribution. As the student's parameters update, its rollout distribution shifts, and the teacher automatically supervises the new frontier of states the student reaches.

In interactive agent tasks such as software engineering, a trajectory interleaves model-generated tool calls with environment-generated outputs (compiler errors, test logs, file contents). During distillation, the SWE specialist reads the entire interaction prefix produced by the student's live sandbox session and provides target distributions only over the action tokens controlled by the student. The environment still dictates what each command returns, while the teacher shapes how the student reacts to those observations.

MT-MOPD scales this mechanism across multiple domains by routing each training prompt to its designated domain specialist. At inference time, the resulting student runs as a single standalone model with no routing overhead or specialist ensemble required.

Conditioning the teacher on peer attempts

Even a strong teacher can struggle to provide sharp token-level credit assignment when it only sees the problem statement and the student's partial attempt. Consider a coding prompt where the student generates eight rollouts: two pass the unit tests, while six fail because they confuse 0-indexed and 1-indexed array bounds. A binary test verifier flags which rollouts passed and failed, but that scalar outcome does not tell the student which token in a 500-token function introduced the off-by-one bug.

MR-MOPD converts those group rollouts into privileged context for the teacher. Before scoring a target rollout $y^{(i)}$, the teacher's prompt is augmented with verified passing and failing peer rollouts from the same group. Seeing a passing peer shows the teacher a concrete working implementation for that exact problem, while seeing a failing peer highlights common pitfalls.

Crucially, this setup works even in a self-distillation regime where the teacher and student start from the exact same model weights. Because the teacher conditions on verified peer solutions that are withheld from the student, it solves an easier conditional prediction problem and translates that privileged hindsight into dense token-level probability targets for the student.

Sources: Problem formulation and method taxonomies follow MT-MOPD §§1, 3, 5 and MR-MOPD §§1, 3.

Previous: Book introduction · Contents · Next: 2. Language Models as Probability Distributions and Policies

2. Language Models as Probability Distributions and Policies

In reinforcement learning and on-policy distillation, an autoregressive language model acts as a stochastic policy over discrete token sequences: at each step, it conditions on the prompt and preceding history to output a categorical distribution over the vocabulary. This chapter establishes the probabilistic notation used throughout the book and highlights implementation invariants—vocabulary alignment, log-normalizer handling, causal index shifts, and action masking—that determine whether a distillation pipeline optimizes the intended objective.

Tokens, logits, and probabilities

A tokenizer maps raw strings into sequences of integer token IDs drawn from a finite vocabulary $V$. Given a context $c$, the language model's forward pass produces a vector of unnormalized real-valued logits $z(c) \in \mathbb{R}^{|V|}$, which the softmax function converts into a categorical probability distribution:

Notice the role of the partition function $\sum_{u\in V}\exp z_u(c)$ in the denominator. If a teacher service returns only the top $k$ raw logits without returning either the full log-normalizer $\log \sum_{u\in V}\exp z_u(c)$ or pre-normalized log-probabilities, the student cannot recover the teacher's true vocabulary probabilities.

Every token probability $p(v\mid c)$ is strictly conditional on the exact prefix $c$. Furthermore, computing a token-level divergence between a student $p$ and a teacher $q$ requires their tokenizers and vocabulary tables to match identically. If token ID 123 maps to different byte sequences in the two models, comparing $p(123\mid c)$ against $q(123\mid c)$ is meaningless. In both MT-MOPD and MR-MOPD, teachers and students share the same base model family, tokenizer, and chat template formatting.

Autoregressive trajectories

Let $x$ denote the input prompt and

denote a generated response of length $T$, where $y_{<t} = (y_1, \ldots, y_{t-1})$ is the prefix before step $t$ (with $y_{<1}$ empty). By the chain rule of probability, a student policy parameterized by weights $\theta$ assigns sequence probability

which becomes a sum in log space:

Throughout the book, we abbreviate the context at step $t$ as $c_t=(x,y_{<t})$, the student's next-token distribution $\pi_\theta(\cdot\mid c_t)$ as $p_t$, and the teacher's target distribution as $q_t$. All logarithms are natural ($\ln$), so entropies and KL divergences are measured in nats.

For single-turn text generation, every token $y_t$ is an action chosen by the policy. For multi-turn tool-use or software-engineering agents, a trajectory interleaves model-generated action tokens with environment-generated observation tokens (such as bash stdout or compiler diagnostics). Only the model's own decisions are parameterized by $\theta$; environment observations must be conditioned on as context but masked out of the policy gradient.

Sampling and the visited-state distribution

A policy's underlying probability distribution $\pi_\theta(\cdot\mid c)$ and the decoding rule used to generate rollouts are distinct mechanisms. Greedy decoding deterministically picks $\arg\max_v \pi_\theta(v\mid c)$, while temperature sampling rescales logits by $\tau > 0$ before applying softmax:

Nucleus (top-$p$) or top-$k$ sampling further truncates the support of the sampling distribution by zeroing out low-probability tails and renormalizing. In policy-gradient training, importance ratios and score-function estimators assume that actions are drawn from a known behavior policy. If rollout workers apply temperature scaling or top-$p$/top-$k$ truncation, the behavior probability logged for each sampled token must reflect the actual sampling distribution used during generation.

Keep the two meanings of $k$ separate: top-$k$ decoding truncates which tokens a rollout worker is allowed to sample during generation, whereas top-$k$ distillation restricts how many teacher token probabilities are transmitted to the loss function during training. Using $k=64$ in top-$k$ distillation does not imply truncating student generation to 64 tokens.

Teacher forcing as a parallel scoring pass

When a teacher evaluates a student rollout $y$, it operates in teacher-forcing (or prefill) mode: at step $t$, the teacher reads the exact student-generated prefix $(x, y_{<t})$ as fixed input history and outputs its own next-token distribution $q_t(\cdot\mid x, y_{<t})$. It never overwrites $y_{<t}$ with the tokens it would have generated itself.

If the student writes 2 + 2 = 5, the teacher's distribution at the final position is conditioned on the prefix 2 + 2 =, allowing it to assign near-zero probability to the student's sampled token 5 and high probability to 4. Because all tokens of the student rollout $y_1, \ldots, y_T$ are already materialized before the teacher is called, a causal Transformer teacher can evaluate $q_1, \ldots, q_T$ across all $T$ positions in a single parallel prefill forward pass. Causal attention masks still prevent position $t$ from attending to tokens at $\ge t$, making teacher scoring vastly faster than autoregressive generation of a $T$-token response.

In code, causal language model outputs are shifted by one index relative to their inputs: the hidden state produced after reading input token $t-1$ predicts target token $t$. Off-by-one indexing bugs here are notorious because the tensor shapes still match and the loss still decreases, even though the student is being trained to predict the token it just emitted.

Full distributions versus single sampled tokens

To see how supervision granularity affects credit assignment, consider a three-action toy vocabulary at a single prefix $c_t$ where the student is editing a function:

If the student samples continuation A and we only evaluate an end-of-sequence unit test, a binary reward tells us merely that the full program failed somewhere. If we use sampled-token distillation, we evaluate the teacher's probability on the single realized action A, observing $q_t(\text{A}) = 0.10 < p_t(\text{A}) = 0.50$; this single-sample signal immediately pushes down the student's probability of A at position $t$, while leaving the relative preference between B and C to be discovered on future visits. If we use distribution-level (or top-$k$) distillation, the update at step $t$ simultaneously suppresses A, boosts B from $0.30$ toward $0.70$, and leaves C unchanged on that single forward pass.

Note that a high local teacher probability $q_t(v)$ means the teacher prefers token $v$ given prefix $c_t$, not that every continuation starting with $v$ is guaranteed to pass downstream verification.

Action masks and length normalization

A serialized training sequence typically concatenates system instructions, user prompts, assistant reasoning tokens, tool calls, tool execution outputs, and right-padding. Let $m_t \in \{0, 1\}$ be a binary loss mask that equals $1$ on assistant-generated decision tokens and $0$ on prompts, environment observations, and padding. The sequence-normalized loss for a trajectory $y$ is

provided $\sum_{t=1}^T m_t > 0$. Normalizing per trajectory gives short and long responses equal weight in the batch gradient, whereas summing $\sum_t m_t \ell_t$ across the batch weights trajectories proportionally to their response length. Masking out environment observations ($m_t = 0$) is mandatory in agentic rollouts so the policy gradient never attempts to train the student to predict deterministic or stochastic environment outputs that the model does not control.

Student, teacher, and verifier roles

Three distinct components interact across the MOPD papers:

• The student $\pi_\theta(\cdot\mid c_t)$ is the parametric policy being updated and ultimately deployed.
• The teacher supplies a target distribution $q_t$ over the vocabulary at each student prefix $c_t$. In MT-MOPD, $q_t = \pi_{\phi_d}(\cdot\mid c_t)$ comes from the specialist for domain $d$. In MR-MOPD, $q_t = q(\cdot\mid c_t, C_i)$ comes from a model conditioned on a set of verified peer attempts $C_i$.
• The verifier evaluates completed trajectories to return a task outcome (such as unit test pass/fail or exact-match math grading). During specialist RL and MR-MOPD rollout partitioning, the verifier grounds correctness; during distillation updates, the teacher translates its specialized weights or verifier-labeled context into dense per-token probabilities.

Sources: Trajectory and context notation follows MR-MOPD §2 and MT-MOPD §3.1 (Eq. 1).

Previous: 1. Two Papers, One Acronym, Two Problems · Contents · Next: 3. Post-Training, Rewards, and Specialist Development

3. Post-Training, Rewards, and Specialist Development

Multi-Teacher MOPD operates as a multi-stage post-training pipeline. Rather than assuming pre-existing external teachers, it develops its own domain specialists through reinforcement learning starting from a shared supervised fine-tuned (SFT) checkpoint. Understanding how those specialists acquire their capabilities—and why training them separately works better than training them jointly—requires examining the mechanics of SFT, verifiable-reward reinforcement learning, and multi-task gradient interference.

Supervised fine-tuning establishes the shared base

Pretraining equips a language model with broad world knowledge and linguistic fluency, while supervised fine-tuning (SFT) shapes it to follow user instructions, adopt chat templates, format tool calls, and emit structured chain-of-thought reasoning. Given a curated demonstration $y^*$ for prompt $x$, SFT minimizes the token-level cross-entropy loss:

Every token in the demonstration provides a direct gradient signal, making SFT fast and stable. However, the model is only supervised on prefixes $y_{<t}^*$ drawn from the static demonstration corpus. Once the model generates its own tokens at inference time and drifts slightly off the demonstration manifold, SFT provides no explicit training on how to recover from its own errors.

In MT-MOPD, a broad multi-domain SFT stage produces the common ancestor checkpoint $\theta_{\mathrm{SFT}}$ that initializes both every domain specialist teacher and the final unified student.

Reinforcement learning and temporal credit assignment

Reinforcement learning closes the loop between generation and evaluation by sampling rollouts $y \sim \pi_\theta(\cdot\mid x)$ from the current policy and scoring them with a task reward $R(x, y)$:

Applying the log-derivative trick yields the score-function (REINFORCE) policy gradient:

In practice, subtracting a context-dependent baseline $b(x)$ from $R(x, y)$ reduces gradient variance without biasing the expectation, and modern policy-gradient optimizers add importance-ratio clipping and reference-policy KL penalties.

Even with variance reduction, outcome-based RL faces a severe temporal credit assignment bottleneck: a single scalar reward $R(x, y)$ at the end of a $10{,}000$-token rollout weights the log-probability gradient of every token in the sequence equally. If a software patch fails its test suite because of a single mistyped variable name on line 400, the negative advantage penalizes all 399 preceding lines of valid setup and reasoning alongside the actual bug.

Verifiable rewards and their scope

Post-training pipelines rely heavily on programmatic, verifiable rewards across four core domains:

• Mathematics: Symbolic equivalence checkers parse the model's boxed final answer and compare it against a ground-truth expression.
• Code generation and SWE agents: Sandboxed execution environments compile the model's patch and run unit or regression test suites.
• Instruction following: Rule-based verifiers check explicit structural constraints (word counts, JSON schemas, forbidden keywords, bullet formatting).
• Agentic tool use: Mock or live environment harnesses verify whether the final environment state matches the target configuration.

While programmatic verifiers eliminate reward-model hacking on surface style, they still measure only what they test. A finite unit test suite can pass code that fails on edge cases; a final-answer math grader rewards a lucky guess reached via a flawed proof; a constraint checker verifies formatting without judging factual accuracy. Because MT-MOPD distills the behavior of RL-trained specialists, any blind spots in a domain's verifier propagate through the specialist into the student. In MR-MOPD, the verifier plays an even more direct role during distillation by partitioning peer rollouts into passing and failing examples for the teacher's context.

Group Relative Policy Optimization (GRPO)

To avoid training a separate value-network critic of equal size to the policy, Group Relative Policy Optimization (GRPO) samples a group of $N$ rollouts $\{y^{(1)}, \ldots, y^{(N)}\}$ for each prompt $x$, scores each with the verifier to obtain rewards $\{R_1, \ldots, R_N\}$, and computes a standardized within-group advantage:

where $\bar R$ and $\operatorname{std}(R_1, \ldots, R_N)$ are the sample mean and standard deviation across the $N$ rollouts for prompt $x$. Rollouts that outperform their peers on the same prompt receive positive advantage ($A_i > 0$), while worse rollouts receive negative advantage ($A_i < 0$).

When a prompt is so easy that all $N$ rollouts pass ($R_1 = \cdots = R_N = 1$) or so hard that all $N$ fail ($R_1 = \cdots = R_N = 0$), the group standard deviation vanishes and $A_i = 0$ across the entire group, contributing zero gradient. Dynamic Sampling filters out these zero-variance groups on the fly and continues sampling until the batch is filled with prompts that exhibit reward variance.

Note how the two stages of MT-MOPD differ here: the domain specialists are trained using GRPO with $N=8$ rollouts per prompt and Dynamic Sampling enabled, whereas the subsequent MT-MOPD distillation stage generates just one rollout per prompt ($N=1$) with no verifier calls or Dynamic Sampling required.

Why domains benefit from decoupled RL recipes

Different capabilities place conflicting operational demands on an RL training loop:

• Rollout length and latency: Competition math problems require long single-turn reasoning chains, instruction following uses short responses, and interactive SWE tasks require multi-turn container execution that can take several minutes per rollout.
• Reward sparsity and variance: A batch of formatting prompts may yield mixed rewards immediately, while hard coding or math benchmarks may produce mostly all-zero groups that require heavy Dynamic Sampling filtration.
• Hyperparameter sensitivity: Clipping thresholds, entropy bonuses, and sampling temperatures that encourage exploration in math reasoning can cause an instruction-following model to violate strict output schemas.

Decoupling specialist development lets each domain team run its own RL loop, environment cluster, rollout horizon, and stopping criterion independently without blocking or destabilizing other domains.

First-order mechanics of multi-domain interference

When a single parameter vector $\theta$ is updated on two domain losses $\mathcal L_a(\theta)$ and $\mathcal L_b(\theta)$ with gradients $g_a = \nabla_\theta \mathcal L_a(\theta)$ and $g_b = \nabla_\theta \mathcal L_b(\theta)$, a gradient step $\Delta \theta = -\eta g_a$ on domain $a$ changes domain $b$'s loss to first order by

Whenever the inner product $g_b^\top g_a < 0$, the two gradients point in conflicting half-spaces, and stepping to improve domain $a$ directly degrades domain $b$. In joint multi-task RL, this conflict is exacerbated by high-variance sequence-level policy gradients and shifting prompt distributions under Dynamic Sampling.

Replacing sparse sequence-level rewards with dense token-level distillation targets dramatically lowers gradient variance, though it does not eliminate shared-parameter capacity limits—the unified student in MT-MOPD still optimizes a single set of weights across a mixture of domain prompts, a trade-off we analyze in Chapter 10.

The four integration baselines in MT-MOPD

MT-MOPD evaluates its specialist-then-distill pipeline against four standard multi-domain post-training alternatives:

1. Mix-RL: Pools prompts and verifiers from all domains into a single joint GRPO training run on one model.
2. Cascade RL: Trains a single model sequentially through one domain RL stage after another (whose final balance depends heavily on stage ordering and stopping points).
3. Off-Policy Finetune: Trains independent domain specialists via RL, uses them to generate static completions on the training prompts, and fine-tunes the student via standard off-policy cross-entropy on those teacher trajectories.
4. Param-Merge: Trains independent domain specialists via RL and merges their parameter tensors directly (e.g., via weight averaging) without further behavioral training.

Sources: Specialist RL configurations and integration baselines follow MT-MOPD §§3–4 and Appendix A.

Previous: 2. Language Models as Probability Distributions and Policies · Contents · Next: 4. Why the Student Must Practice Its Own Paths

4. Why the Student Must Practice Its Own Paths

Standard knowledge distillation works by imitation: you collect responses from a teacher model and train the student to reproduce them. On-policy distillation flips the order of operations. The student generates a response first, and the teacher scores the next-token distributions along the path the student actually took.

In an autoregressive language model, a state is a prompt concatenated with the partial response generated so far. Evaluating the teacher on states visited by the student makes the supervision responsive to the student's current habits, mistakes, and recovery paths.

A branching example

Suppose a coding problem requires composing two lookup tables. A strong teacher keeps the two mappings separate from line one, whereas the student conflates them under a misleading variable name. By the fourth line of code, the student has entered a prefix that would never appear in a corpus of teacher demonstrations.

Offline imitation keeps drilling the student on the teacher's clean opening lines. That may eventually fix the student's initial variable declaration, but it provides no direct supervision on how to continue, recover, or catch the bug once the bad prefix is already on screen. On-policy supervision hands that exact fourth-line student prefix to the teacher and asks how the teacher distributes probability mass over the next token. The benefit comes from matching the sampled histories to the student's current policy, regardless of whether the student and teacher share the same parameter count or architecture.

State distributions formalize the mismatch

Let $d_{\pi,t}(c)$ denote the distribution of contexts reached at generation step $t$ when policy $\pi$ rolls out responses. A corpus of teacher trajectories samples contexts from $d_{q,t}$, while deploying the student exposes it to contexts drawn from $d_{\pi_\theta,t}$.

An offline distillation objective evaluates the divergence over teacher contexts:

An on-policy local objective shifts the outer expectation to the student's state distribution, and typically reverses the divergence direction as well:

Notice that the context distribution ($d_q$ vs. $d_{\pi_\theta}$) and the divergence direction (forward vs. reverse KL) are independent design axes. You can compute forward KL on student-generated prefixes or reverse KL on a static offline dataset, so calling a loss "on-policy" specifies where the prefixes come from, not which divergence measures disagreement at each token.

Sampling from $d_{\pi_\theta}$ ensures that the teacher evaluates states the student is actually likely to reach—including flawed intermediate states where the student most needs corrective steering.

How errors compound across long horizons

In a simplified $T$-step generation process where the student has a constant per-step probability $\epsilon$ of stepping off the teacher's preferred path, independent steps would leave a $(1-\epsilon)^T$ chance of staying on track. Even at $\epsilon=0.01$ over $T=100$ tokens, $(1-0.01)^{100} \approx 0.366$.

Token transitions in real language models are heavily correlated and many deviations are benign paraphrases, yet the basic compounding dynamic remains: tiny local discrepancies rapidly shift the distribution of long prefixes. In multi-turn agent workflows, early actions alter the external environment itself. Editing the wrong line of code triggers a different compiler traceback, which changes every subsequent token in the context window. At that point, the student isn't just continuing a slightly unusual phrasing; it is operating in a different environment state altogether.

What "no exposure bias" covers—and what it leaves out

The MT-MOPD paper states that training on student rollouts eliminates exposure bias because the integration targets are evaluated on prefixes generated by the current student rather than on teacher-generated prefixes.

Training on $d_{\pi_\theta}$ closes the gap between the prefixes seen during training and those generated at inference under the training distribution, though several operational mismatches still remain in practice:

• Deployment traffic rarely matches the prompt distribution of the training corpus.
• Production serving may use different decoding temperatures, system prompts, tool schemas, or maximum context lengths.
• In asynchronous distributed training, worker rollouts can lag behind the optimizer by several steps, making the sampled prefixes slightly off-policy relative to the latest weights $\theta$.

On-policy sampling tracks a moving student distribution during training rather than guaranteeing an exact match to production inference.

The five steps of an on-policy distillation update

A single training step follows a straightforward loop:

1. Sample a task prompt (and its domain label, if routing across specialists).
2. Roll out a complete response from the student under its current policy $\pi_\theta$.
3. Pass the prompt and student-generated response to the teacher for causal scoring.
4. Compute the token-level divergence between the student and teacher distributions at all valid response positions.
5. Backpropagate the loss to update $\theta$, then draw fresh student rollouts.

Crucially, the teacher never has to generate text token by token. Because the full student sequence is already known, a teacher prefill service can score every prefix of the trajectory in a single parallel forward pass.

The two MOPD variants hook into this loop at different points: MT-MOPD routes step 3 to a domain-specific teacher based on the prompt's task label, whereas MR-MOPD inserts a verification and context-assembly step between steps 2 and 3 so the teacher conditions on peer rollouts alongside the prompt.

Why failed rollouts still carry rich training signal

When every rollout in a batch fails a binary unit test or math verifier, outcome-based RL methods like GRPO get zero relative advantage across the group and waste the batch. A teacher model scoring those same failed trajectories token by token still differentiates good local continuations from bad ones—assigning higher probability to a missing boundary check, a better tool call, or a cleaner algebraic step.

Local supervision does have a structural limit: if a student prefix has already painted itself into a dead end by token 50, token-level teacher targets at token 51 can encourage the model to backtrack or hedge, but they cannot erase the first 50 tokens of that trajectory. Instead, the gradients applied at the earlier positions shift $\pi_\theta$ so that subsequent rollouts are less likely to enter that dead end in the first place.

Rollout freshness and data reuse

A strictly on-policy loop uses each rollout batch for a single gradient step and discards it. Many large-scale RL pipelines instead take multiple epochs or minibatch updates over each collected rollout buffer to amortize generation cost. The first gradient step is genuinely on-policy; subsequent steps are off-policy updates over slightly stale trajectories, typically stabilized with importance-sampling ratios and PPO-style clipping.

MT-MOPD specifies single-use rollouts for its specialist RL stage, though its appendix leaves open the exact number of gradient updates per rollout batch during the distillation integration stage (a detail we revisit in Chapter 22).

Why this remains distillation rather than environment RL

Even when implemented inside a policy-gradient trainer, the training signal in MOPD comes from a teacher's predictive distribution $\log q(y_t \mid x, y_{<t})$ rather than a ground-truth environment reward. That teacher may itself be the product of extensive RL, yet its token probabilities are still model beliefs: if a teacher assigns high probability to a plausible-sounding fallacy, the student is rewarded for matching it.

On-policy distillation borrows the rollout collectors and policy-gradient estimators of reinforcement learning to optimize distributional agreement with a teacher. Verifying that this agreement translates into actual task performance still requires independent outcome evaluation, and ensuring the supervision is trustworthy makes teacher selection and routing central to the pipeline.

Sources: On-policy distillation definitions follow MR-MOPD §2 and MT-MOPD §3.

Previous: 3. Post-Training, Rewards, and Specialist Development · Contents · Next: 5. What KL and Jensen–Shannon Actually Teach

5. What KL and Jensen–Shannon Actually Teach

A distillation loss measures the discrepancy between two probability distributions over the vocabulary at a given context. Which divergence you choose determines which disagreements get penalized most heavily, how the estimator interacts with sampling variance, and how the loss behaves when one model assigns near-zero probability to a token.

Throughout this chapter, $p$ denotes the student's next-token distribution and $q$ denotes the teacher's next-token distribution at a fixed context $c$, defined over a common vocabulary $V$ with natural logarithms.

Forward KL

The forward (teacher-to-student) Kullback–Leibler divergence weights log-ratios by the teacher's probabilities:

Because the teacher distribution $q$ does not depend on the student parameters $\theta$, minimizing $D_{\mathrm{KL}}(q\|p)$ with respect to $\theta$ is identical to minimizing the teacher-weighted cross-entropy:

Whenever the teacher places noticeable mass on a token $v$ ($q_v > 0$) and the student assigns it nearly zero probability ($p_v \to 0$), the $-\log p_v$ term blows up. Forward KL therefore forces the student to spread probability mass across every token the teacher considers plausible, even tokens the student would rarely sample on its own.

In variational inference and distillation, this property is called mode covering (or zero-avoiding). When the student is capacity-limited and cannot match a multimodal teacher distribution $q$ exactly, forward KL pushes $p$ to bridge across the teacher's modes rather than dropping one. If the student's output layer can represent any categorical distribution over $V$ and optimization is unconstrained, the unique global minimum of both forward and reverse KL is simply $p = q$.

Reverse KL

The reverse (student-to-teacher) KL divergence weights the log-ratio by the student's probabilities:

Here the expectation is taken under $p$. If the student puts significant probability mass on a token that the teacher considers virtually impossible ($q_v \to 0$), $\log(p_v/q_v)$ becomes huge and the student is sharply penalized. Conversely, if the student assigns near-zero probability to a token the teacher likes, the product $p_v \log(p_v/q_v)$ approaches zero, and a sampled-token estimator under $p$ will rarely even visit $v$ to observe the discrepancy.

Expanding the log-ratio makes the role of student entropy explicit:

where $H(p)=-\sum_v p_v\log p_v$ is the Shannon entropy of the student distribution. Minimizing reverse KL simultaneously pushes the student to put mass on tokens with high teacher log-likelihood $\log q_v$ while maintaining its own entropy $H(p)$ so it does not collapse onto a single argmax token faster than the teacher warrants.

Reverse KL is often called mode seeking (or zero-forcing) because an under-parameterized student—or an on-policy sampler that has already concentrated around one valid solution path—can achieve low loss by locking onto a single high-probability mode of $q$ and ignoring others.

Numerical comparison on a two-token vocabulary

Consider a binary vocabulary where the student strongly favors the first token while the teacher is indifferent between the two:

The reverse KL evaluates to

whereas the forward KL is larger because it weights the under-covered second token by $q_2 = 0.5$ instead of $p_2 = 0.1$:

If we train with a sampled-token estimator that draws $v \sim p$ and uses the log-ratio $\log(q_v/p_v)$ as a policy-gradient advantage:

• Drawing the first token (probability $0.9$) yields $\log(0.5/0.9) \approx -0.5878$, nudging $p_1$ downward.
• Drawing the second token (probability $0.1$) yields $\log(0.5/0.1) = \log 5 \approx 1.6094$, pulling $p_2$ upward.

Although the second token carries a strong positive learning signal, a single-sample estimator only sees that signal on 10% of draws, linking divergence geometry directly to Monte Carlo variance and exploration.

Support mismatches, truncation, and numerical floors

Mathematically, $D_{\mathrm{KL}}(p\|q) = \infty$ whenever $q_v = 0$ and $p_v > 0$, while $D_{\mathrm{KL}}(q\|p) = \infty$ whenever $p_v = 0$ and $q_v > 0$. An exact softmax over finite logits produces strictly positive probabilities across the entire vocabulary, but practical serving and training pipelines introduce effective zeros in three ways:

1. Floating-point underflow in fp16 or bf16 when exponentiating large negative logits.
2. Hard vocabulary masks or mismatched special-token IDs between models.
3. Top-$k$ or top-$p$ truncation in prefill services that only return the highest-scoring teacher tokens over the network.

Clamping teacher log-probabilities at a floor (or renormalizing a truncated top-$k$ payload) prevents NaN gradients, at the cost of altering the target distribution in the tails. Working directly in log-space via log_softmax and subtracting log-probabilities ($\log p_v - \log q_v$) avoids exponentiation underflow altogether.

Jensen–Shannon divergence

Defining the midpoint mixture distribution $m = \frac{1}{2}(p+q)$, the Jensen–Shannon (JS) divergence averages the KL from each distribution to $m$:

JS divergence is symmetric, bounded between $0$ and $\log 2$ for normalized distributions, and remains finite even when $p$ and $q$ have completely disjoint support. That boundedness softens gradient spikes when a student and teacher strongly disagree on tail tokens.

When differentiating $D_{\mathrm{JS}}(p_\theta, q)$ with respect to $\theta$, note that $m = \frac{1}{2}(p_\theta + q)$ also depends on $p_\theta$. Exact automatic differentiation propagates gradients through both appearances of $p_\theta$ (yielding $\frac{1}{2}\sum_v \nabla_\theta p_v \log(p_v / m_v)$), whereas detaching $m$ changes the update into a different surrogate.

Across the two MOPD papers, divergence choices vary by setting: MR-MOPD uses reverse KL for mathematical reasoning and JS divergence for its remaining benchmarks, while MT-MOPD uses reverse KL everywhere (implemented via either a sampled-token policy gradient or a tail-corrected top-$k$ sum).

Temperature scaling changes the target

Dividing logits by a temperature $\tau$ before the softmax reshapes the target distribution: $\tau > 1$ flattens the distribution and exposes dark knowledge in lower-ranked tokens, whereas $\tau < 1$ sharpens mass onto the top candidates. Because changing $\tau$ alters the target distribution $q_\tau$, three temperatures need to be kept distinct when reading or reproducing distillation experiments:

• The rollout sampling temperature used by the student to generate trajectories.
• The scoring temperature applied to teacher and student logits inside the divergence loss.
• The evaluation decoding temperature used at benchmark time.

Divergence, task accuracy, and diversity measure different things

Driving $D_{\mathrm{KL}}(p\|q)$ toward zero only certifies that the student matches the teacher's next-token probabilities on visited prefixes. It does not guarantee task correctness or output diversity:

• Correctness vs. agreement: A student can achieve zero KL against a teacher that hallucinates confidently, or have high KL against a teacher simply because it writes a valid proof using different variable names.
• Token entropy vs. semantic diversity: High average token entropy $H(p)$ can reflect superficial synonym variation (choosing between "however" and "nevertheless") along a single rigid algorithm, while a localized shift on a single branching token can switch between two completely different solution strategies with almost no change in sequence-averaged entropy.

Evaluating a distillation run therefore requires pairing distributional metrics (KL, token entropy) with ground-truth task verifiers (Pass@1, Pass@k, and semantic solution-cluster coverage).

Sources: Divergence definitions and task-specific loss choices in MR-MOPD appear in MR-MOPD §3 and Appendix; MT-MOPD's reverse-KL formulations appear in MT-MOPD §3.2.

Previous: 4. Why the Student Must Practice Its Own Paths · Contents · Next: 6. The Multi-Teacher Pipeline from Beginning to End

6. The Multi-Teacher Pipeline from Beginning to End

Multi-Teacher MOPD decouples capability acquisition from capability integration. Instead of training a single model with a fragile, multi-objective RL loop across every domain at once, domain specialists are trained independently with domain-tailored RL and then distilled into a single student model on the student's own rollouts.

Stage 1: Establish a shared SFT origin

The pipeline starts from a pretrained base model and runs a broad supervised fine-tuning (SFT) phase across all target domains to produce a shared starting checkpoint $\theta_0$. This SFT stage establishes common chat formatting, reasoning delimiters, and baseline competence across tasks.

Both the domain teachers and the eventual unified student initialize from $\theta_0$. Consequently, every specialist in the pool is a branch from the same ancestor—sharing an identical vocabulary, tokenizer, and architectural backbone—rather than an arbitrary collection of third-party checkpoints. In MT-MOPD's controlled experiments, $\theta_0$ is a custom SFT checkpoint trained on top of Qwen3-30B-A3B-Base.

Stage 2: Train domain specialists independently

Let $d\in\{1,\ldots,D\}$ index the target domains. Starting from $\theta_0$, each domain team runs an independent reinforcement learning pipeline to produce specialist weights $\phi_d$:

Each specialist run uses its own prompt dataset $\mathcal D_d$, reward verifier $R_d$, rollout horizon, and RL hyperparameters. In the paper's controlled 30B study, $D=3$ covers competition mathematics, instruction following (IF), and software engineering (SWE); the production MiMo-V2-Flash deployment scales this up to also include competitive code generation and multi-turn agentic tool use.

Because the specialist runs do not share weights during Stage 2, a reward bug or instability in the SWE agent environment never corrupts the math specialist's weights. Domain teams can iterate asynchronously, tune domain-specific rollout lengths or exploration temperatures, and swap updated teachers into the pool without forcing other teams to restart their RL runs.

Stage 3: Freeze teachers and integrate via on-policy distillation

Once the domain specialists $\phi_1, \dots, \phi_D$ are ready, their weights are frozen and deployed as inference-only scoring services. A single student model $\pi_\theta$ is initialized at $\theta_0$ and trained on a mixed stream of prompts drawn from all domains.

For each prompt $x$, the student generates an on-policy trajectory $y \sim \pi_\theta(\cdot\mid x)$. The prompt's domain tag $d(x)$ routes the completed trajectory (x, y) to the matching specialist teacher $\phi_{d(x)}$, which runs a causal forward pass over the student's tokens and returns token-level probability targets. The student then takes a gradient step to reduce the reverse KL divergence against that teacher:

                     broad SFT checkpoint (theta_0)
                       /           |            \
                  math RL        IF RL         SWE RL
                     |             |              |
                 teacher M     teacher I      teacher S
                       \           |            /
                   domain-routed prefill supervision
                                   |
             student rollouts initialized from theta_0
                                   |
                       one unified student model

The multiple teachers exist only during training; at deployment, only the single unified student checkpoint $\theta$ is served.

Prompt-level routing vs. token-level ensembling

In MT-MOPD, every training prompt $x$ carries a known domain label $d(x)$ from its source dataset, so the teacher target at step $t$ is simply the distribution of that domain's specialist:

Notice what this routing mechanism does not do:

• It does not query all $D$ teachers on every prompt and average their probability distributions.
• It does not train a learned gating network to switch teachers token by token mid-response.
• It does not require a classifier at inference time, because the distilled student internalizes all domains into one set of parameters.

Routing by prompt dataset keeps prefill compute at $1\times$ teacher forward pass per student rollout rather than $D\times$, and avoids averaging conflicting stylistic priors (such as a terse IF teacher and a verbose chain-of-thought math teacher) on the same trajectory.

Why frozen teachers still provide an adaptive curriculum

Even though the specialist parameters $\phi_d$ are frozen throughout an integration round, the supervision targets $q_{d(x)}(\cdot\mid x, y_{<t})$ change continuously as the student evolves. Early in integration, the student visits clumsy prefixes $y_{<t}$, and the teacher provides corrective next-token distributions from those states. As $\pi_\theta$ improves, its rollouts reach deeper reasoning steps or cleaner code structures, and the frozen teacher automatically scores those more advanced prefixes.

When the team wants to push capabilities further across multiple generations, they can start a new round by branching fresh RL specialists from the newly integrated student checkpoint—an iterative generational loop examined in Chapter 14.

Anatomy of a single integration step

To run one optimization step across a distributed training cluster:

1. Sample and rollout: Draw a mixed-domain batch of prompts, recording each prompt's text, domain ID $d(x)$, generated token IDs $y$, loss mask (masking out prompt tokens and injected environment outputs), and behavior-policy log-probabilities.
2. Route and score: Dispatch each trajectory (x, y) to the prefill service hosting teacher $\phi_{d(x)}$ using the exact same tokenizer and chat template, and return either the teacher's log-probability on the sampled tokens $y_t$ or its top-$k$ logits at each position.
3. Update the student: Run a forward-backward pass on the student trainer to compute current student log-probabilities and step the optimizer using either the sampled-token policy-gradient loss (Chapter 8) or the tail-corrected top-$k$ reverse-KL loss (Chapter 9).

To prevent silent supervision bugs when swapping specialists in and out of a cluster, treat each teacher as a versioned teacher bundle that locks together the weight checkpoint, tokenizer, chat template, special-token masks, and scoring temperature.

Reported hyperparameters for the 30B controlled study

In Appendix A of the MT-MOPD paper, the controlled 30B integration stage uses:

• Global batch size: 2,048 prompts per step
• Rollouts per prompt: 1
• Domain mixture (Math : IF : SWE): $0.35 : 0.35 : 0.30$
• Dynamic sampling filtering: disabled during distillation
• Advantage clipping ceiling: $5.0$ (for the policy-gradient formulation)
• Top-$k$ payload size: $k = 64$ (for the top-$k$ formulation)

The paper leaves several secondary training hyperparameters (such as the exact integration learning-rate schedule and optimizer betas) unspecified in the public appendix; Chapter 21 and Chapter 22 detail sensible defaults and verification checks when building a reproduction pipeline.

Sources: The three-stage architecture, domain routing, and prefill service design are presented in MT-MOPD §3, with integration hyperparameters in MT-MOPD Appendix A.

Previous: 5. What KL and Jensen–Shannon Actually Teach · Contents · Next: 7. The Reverse-KL Objective and Its Gradient

7. The Reverse-KL Objective and Its Gradient

Multi-Teacher MOPD minimizes the token-level reverse KL divergence between the student and the routed domain teacher along trajectories sampled from the student. Because the student parameters $\theta$ govern both the local next-token distribution at each prefix and the probability of visiting that prefix in the first place, deriving the gradient requires distinguishing the local fixed-context gradient from the full trajectory-occupancy gradient.

The token-averaged reverse-KL objective

Let $x$ be a prompt with domain label $d(x)$, and let $y=(y_1,\dots,y_T)\sim\pi_\theta(\cdot\mid x)$ be a trajectory generated by the student. Writing the conditional next-token distributions at prefix $c_t = (x, y_{<t})$ as

MT-MOPD states its training objective (Equation 1) as the expected per-token reverse KL across student rollouts:

Here the inner sum over $V$ evaluates the exact categorical reverse KL at prefix $c_t$, the sum over $t=1,\dots,T$ averages across the response length, and the outer expectation samples trajectories from the current student policy $\pi_\theta$. Although the teacher parameters $\phi_{d(x)}$ are frozen, $q_t(v)$ still varies from rollout to rollout because it conditions on the student's prefix $y_{<t}$.

Differentiating reverse KL at a fixed context

Hold a single visited context $c$ fixed and consider the local reverse KL as a function of $\theta$:

Applying the product rule with respect to $\theta$ while holding $c$ and $q(\cdot\mid c)$ constant gives

Because $p_\theta(\cdot\mid c)$ is a normalized probability distribution over $V$, the sum of its probability gradients vanishes:

Dropping the $+1$ term and applying the log-derivative identity $\nabla_\theta p_v = p_v\nabla_\theta\log p_v$ yields the score-function form of the local reverse-KL gradient:

Sampling the actual next token $y_t \sim p_\theta(\cdot\mid c_t)$ during the student rollout therefore provides a single-sample Monte Carlo estimate of $\nabla_\theta K(\theta;c_t)$.

The log-probability ratio as a token-level advantage

MT-MOPD defines the reverse-KL log-ratio as a token-level advantage signal:

With this definition, the local gradient becomes $-\mathbb E_{v\sim p}[A(v;c)\nabla_\theta \log p_v]$, matching the form of a policy-gradient update:

• When $q(v\mid c) > p(v\mid c)$, the teacher considers the sampled token more likely than the student does, so $A(v;c) > 0$ and the update increases $\log p_\theta(v\mid c)$.
• When $q(v\mid c) < p(v\mid c)$, the student over-estimates the token relative to the teacher, so $A(v;c) < 0$ and the update pushes $\log p_\theta(v\mid c)$ down.
• For example, if $p_v=0.1$ and $q_v=0.4$, then $A=\log 4\approx +1.3863$; if $p_v=0.4$ and $q_v=0.1$, then $A=-\log 4\approx -1.3863$; and when $p_v = q_v$, $A = 0$.

Notice that $A(v;c)$ is an immediate conditional log-likelihood difference, not a cumulative return-to-go over future steps. In addition, subtracting any context-only baseline $b(c)$ leaves the unclipped expectation unchanged because $\mathbb E_{v\sim p}[\nabla_\theta \log p_v] = 0$.

Why stop_gradient is required in the surrogate loss

To implement this update with standard automatic differentiation, we form the per-token surrogate loss

where $\operatorname{sg}[\cdot]$ detaches the advantage from the computation graph. Differentiating $\ell_{\mathrm{sur}}$ produces $-A(v;c)\nabla_\theta\log p_\theta(v\mid c)$, whose expectation over $v\sim p_\theta$ equals $\nabla_\theta K(\theta;c)$.

If you omit $\operatorname{sg}[\cdot]$ and allow gradients to flow through the $-\log p_\theta(v\mid c)$ term inside $A(v;c)$, differentiation introduces an extra $+\nabla_\theta\log p_\theta(v\mid c)$ term on the sampled token (which only vanishes in expectation over the full vocabulary, adding unnecessary single-sample variance and breaking importance-sampled or clipped surrogates). Also keep in mind that the scalar value of $\ell_{\mathrm{sur}}$ is not the KL divergence itself—it can be positive or negative and should not be logged as a proxy for $D_{\mathrm{KL}}(p\|q)$.

Local semi-gradients vs. full trajectory-distribution gradients

Now return to the outer expectation over student trajectories $y \sim \pi_\theta(\cdot\mid x)$ in Equation 1. Let

denote the trajectory-averaged local KL. Because both the integrand $f_\theta(y)$ and the sampling distribution $\pi_\theta(y\mid x)$ depend on $\theta$, the exact gradient of $\mathbb E_{y\sim\pi_\theta}[f_\theta(y)]$ has two parts:

1. The local step-wise term $\mathbb E_y[\nabla_\theta f_\theta(y)]$ differentiates the token-level KL at each visited prefix $c_t(y)$ while treating the trajectory prefixes themselves as fixed.
2. The state-occupancy term $\mathbb E_y[f_\theta(y)\nabla_\theta\log\pi_\theta(y\mid x)]$ accounts for how changing an early token probability $\pi_\theta(y_t\mid c_t)$ shifts the probability of reaching downstream states $c_{t'}$ ($t' > t$) where local KL costs may be higher or lower.

MT-MOPD's Equation 2—like most practical on-policy distillation algorithms—implements only the first term: it rolls out $y \sim \pi_\theta$, treats the visited prefixes $c_1,\dots,c_T$ as fixed contexts, and applies the token-wise score-function update $-A(y_t;c_t)\nabla_\theta\log p_\theta(y_t\mid c_t)$ at each position $t$ independently. In reinforcement learning terminology, this is a local (or myopic) semi-gradient estimator on the current state distribution.

Contrast with sequence-level reverse KL

To see how trajectory occupancy enters when differentiating a full sequence divergence, consider the unnormalized sequence-level reverse KL over fixed-length trajectories:

By the chain rule of KL divergence, the scalar value of this sequence KL equals the sum of expected step-wise conditional KLs: $\sum_{t=1}^T \mathbb E_{c_t}[K(\theta;c_t)]$. However, its exact policy gradient at step $t$ weights $\nabla_\theta\log p_\theta(Y_t\mid c_t)$ by the cumulative future log-ratio (reward-to-go) $\sum_{t'=t}^T \log\frac{p_\theta(Y_{t'}\mid c_{t'})}{q(Y_{t'}\mid c_{t'})}$, rather than only the immediate step-$t$ log-ratio.

A two-step branching example makes the difference concrete: suppose at $t=1$ the student chooses between branch A and branch B, and on this first token $p_1$ already matches $q_1$ closely on both branches ($A_1 \approx 0$). However, entering branch A leads to a state at $t=2$ where the student and teacher strongly disagree ($K(\theta; c_2^A) \gg 0$), whereas branch B leads to a state where they agree ($K(\theta; c_2^B) = 0$).

• The local token-wise update gives zero gradient at $t=1$ because $p_1(A) = q_1(A)$, and only updates the student's second-step behavior inside branch A whenever branch A happens to be visited.
• The full sequence-KL gradient additionally pushes down $p_1(A)$ at $t=1$ to avoid visiting the high-discrepancy state $c_2^A$ in the first place.

In practice, the myopic token-wise update is simpler, avoids the high variance of summing log-ratios over thousand-token trajectories, and prevents the student from gaming sequence KL by prematurely terminating responses or avoiding difficult reasoning branches just to dodge future teacher disagreement.

Sources: The token-averaged reverse-KL objective, log-ratio advantage, and stop-gradient surrogate correspond to MT-MOPD §3.2, Equations 1–4.

Previous: 6. The Multi-Teacher Pipeline from Beginning to End · Contents · Next: 8. Implementing the Sampled-Token Update

8. Implementing the Sampled-Token Update

The policy-gradient formulation of Multi-Teacher MOPD reduces the teacher's network payload to a single float—one log-probability—per sampled student token. Because the student rollout already specifies the exact sequence of chosen tokens $y_1, \dots, y_T$, the teacher prefill service only needs to report how likely it considers those specific tokens at their respective prefixes.

Turning this formula into a reliable training loop requires keeping three probability distributions distinct at every token position: the behavior policy that generated the rollout, the current student policy being differentiated, and the frozen domain teacher.

The three distributions at a single token position

At context $c_t = (x, y_{<t})$ and sampled token $y_t$, three models enter the calculation:

1. $p_b(y_t\mid c_t)$: the behavior student snapshot that generated the rollout.
2. $p_\theta(y_t\mid c_t)$: the current student weights on the trainer node receiving the gradient update.
3. $q_d(y_t\mid c_t)$: the routed domain teacher $\phi_{d(x)}$.

In a strictly synchronous, single-update loop, $p_b$ and the pre-update $p_\theta$ are identical up to floating-point kernel differences between the inference engine (e.g., SGLang or vLLM) and the training backend (e.g., Megatron-LM or FSDP). If rollouts are buffered across asynchronous steps or reused for multiple minibatch updates, $p_\theta$ drifts away from $p_b$. Logging and storing $\log p_b(y_t\mid c_t)$ at rollout time makes any policy lag explicit and enables importance-sampling diagnostics.

The immediate reverse-KL token advantage is

where $\log p(y_t\mid c_t)$ is evaluated under the student distribution.

Two-sided advantage clipping

When a student samples a token that lies in the far tail of the teacher's distribution (or vice versa), the raw log-ratio $A_t$ can spike to large magnitudes. To prevent single-token outliers from destabilizing the optimizer, MT-MOPD clips the advantage symmetrically:

and minimizes the trajectory-averaged surrogate loss

In Appendix A, MT-MOPD sets $A_{\max} = 5.0$, which caps the effective probability ratio $q_d/p$ at $e^5 \approx 148.4$ (and $e^{-5} \approx 0.00674$ on the negative side).

Note that advantage clipping on $\log(q_d/p)$ is distinct from PPO probability-ratio clipping on $p_\theta / p_b$:

• Advantage clipping bounds how strongly any single teacher-student disagreement can pull the gradient, introducing a controlled bias into the local reverse-KL estimator to reduce variance.
• PPO ratio clipping bounds how far the updated actor $p_\theta$ can move away from the behavior snapshot $p_b$ across multiple off-policy epochs.

When monitoring training runs, track positive and negative advantage clipping fractions separately; a surge in negative clipping (A_t < -5.0) usually signals that the student has drifted into prefixes or formatting conventions that the teacher strongly rejects.

Tracing the gradient sign across token types

Minimizing $-\bar A_t \log p_\theta(y_t\mid c_t)$ scales the score function $\nabla_\theta \log p_\theta(y_t\mid c_t)$ by $\bar A_t$. Because the vocabulary probabilities are coupled through the softmax normalizer, raising the logit of $y_t$ when $\bar A_t > 0$ slightly lowers the probabilities of competing tokens, whereas pushing down the logit of $y_t$ when $\bar A_t < 0$ redistributes mass across the remaining logits according to the network's representations.

Consider three representative token positions inside a single failed coding trajectory:

Where an outcome-based RL update would assign the exact same negative advantage to all three tokens because the unit test ultimately failed, the teacher's conditional probabilities reinforce the valid boundary check while specifically penalizing the buggy variable assignment.

Reference training step

The pseudocode below outlines a synchronous, single-update sampled-token step:

batch = sample_domain_labeled_prompts(domain_mixture)
rollouts = student_sampler.generate(batch, policy_version=current_version)
records = attach_masks_and_behavior_logprobs(rollouts)

# Route each sequence to its domain specialist for causal prefill scoring
teacher_logp = score_sampled_tokens_by_domain(records, frozen_teachers)
student_logp = actor.logprobs(records.token_ids, records.attention_mask)

# In a fresh single-step update, student_logp.detach() or behavior_logp may be used
advantage = (teacher_logp - student_logp.detach()).clamp(-5.0, 5.0)

per_token_loss = -advantage * student_logp
per_response_loss = masked_response_mean(per_token_loss, records.actor_mask)
loss = per_response_loss.mean()

optimizer.zero_grad()
loss.backward()
optimizer.step()

If your training infrastructure reuses rollouts across multiple gradient steps or tolerates asynchronous policy lag between student_sampler and actor, multiply the per-token surrogate by the importance ratio $r_t(\theta) = \exp(\log p_\theta(y_t\mid c_t) - \log p_b(y_t\mid c_t))$ (with optional PPO clipping on $r_t$), while recomputing the reverse-KL advantage against $\operatorname{sg}[\log p_\theta]$.

Always transmit normalized log-probabilities, never raw logits

A common implementation bug when adapting inference servers for distillation is returning raw pre-softmax logits $z^{\text{teacher}}_{y_t}$ instead of normalized log-probabilities $\log q_d(y_t\mid c_t)$. Because adding an arbitrary context-dependent constant $C(c_t)$ to every logit leaves the softmax distribution unchanged, raw logit differences $z^{\text{teacher}}_{y_t} - z^{\text{student}}_{y_t}$ are contaminated by the mismatch between the two models' log-partition functions $\log \sum_v \exp(z_v)$.

The teacher prefill endpoint must therefore compute the full-vocabulary log_softmax at each position $t$ before gathering the scalar entry corresponding to the student's sampled token ID $y_t$.

Loss masking and sequence length normalization

Only tokens actively sampled by the student policy should contribute to $\mathcal L_{\mathrm{PG}}$. The binary actor_mask must zero out:

• User prompt and system instruction tokens.
• Right-padding tokens in batched tensors.
• External environment observations injected mid-trajectory (such as compiler outputs, file contents, or tool responses in SWE and agentic domains).

Leaving tool outputs unmasked is particularly damaging in agentic trajectories: a 500-token stack trace returned by a bash tool is deterministic environment state, not an action chosen by $\pi_\theta$, and training the student to match teacher probabilities over stack-trace tokens wastes gradient capacity and distorts sequence weighting.

Notice also how MT-MOPD normalizes across sequence length: Equation 1 averages over the $T$ active tokens within each trajectory first ($\frac{1}{T}\sum_{t=1}^T$), and then averages across the batch. Computing a flat global mean over all active tokens in a batch weights a 16,000-token SWE trajectory 32 times more heavily than a 500-token IF response, silently skewing the effective domain mixture toward whichever domain produces the longest rollouts.

Production telemetry and sanity checks

Because the surrogate loss $\mathcal L_{\mathrm{PG}} = -\frac{1}{T}\sum_t \bar A_t \log p_\theta(y_t\mid c_t)$ has a detached advantage multiplier, its scalar value can be positive or negative and does not measure KL divergence. Instead, log the following metrics per domain at every step:

• Estimated reverse KL: the batch mean of $-\tilde A_t = \log p_\theta(y_t\mid c_t) - \log q_d(y_t\mid c_t)$ (since $\mathbb E_{y_t\sim p_\theta}[\log p_\theta(y_t\mid c_t) - \log q_d(y_t\mid c_t)] = D_{\mathrm{KL}}(p_\theta\|q_d)$).
• Advantage quantiles and clip rates: fraction of tokens hitting $+A_{\max}$ and $-A_{\max}$.
• Student token entropy and response length: broken down by domain to catch premature length collapse or verbosity drift.
• Inference-trainer logprob gap: $|\log p_\theta^{\text{init}}(y_t\mid c_t) - \log p_b(y_t\mid c_t)|$ on step 0 of a fresh batch, which catches chat-template, tokenizer, or RoPE position-ID mismatches immediately.

Sources: The sampled-token advantage, clipping threshold ($A_{\max}=5$), and policy-gradient surrogate loss follow MT-MOPD §3.2.1 and Appendix A.

Previous: 7. The Reverse-KL Objective and Its Gradient · Contents · Next: 9. Why Top-k Distillation Needs a Correction

9. Why Top-k Distillation Needs a Correction

While the sampled-token policy gradient sends just one float per token position, full-vocabulary KL distillation requires transmitting and storing ~150,000 floats per position. Multi-Teacher MOPD evaluates a middle ground: having the teacher service return its top-$k$ highest-probability tokens at each student prefix.

Truncating a reverse-KL sum to a subset of the vocabulary introduces a subtle mathematical trap. In a full-vocabulary KL sum, the constant $+1$ terms in the derivative $\frac{\partial}{\partial p_v}(p_v \log(p_v/q_v)) = \log(p_v/q_v) + 1$ cancel out because $\sum_{v \in V} p_v = 1$ is fixed. Over an unnormalized subset $S \subset V$, the partial sum $\sum_{v \in S} p_v$ is free to shrink, so a naive partial KL sum actively pushes student probability mass out of the top-$k$ set and into the unscored tail. MT-MOPD adds a linear Bregman correction term to restore the minimum to $p_v = q_v$ on every scored coordinate.

The tail-corrected top-$k$ objective

At a fixed prefix, let $S = \operatorname{TopK}_k(q)$ denote the indices of the $k$ largest probabilities under the teacher distribution $q$ (with $k = 64$ in MT-MOPD's experiments). Using the full-vocabulary normalized probabilities $p_v$ and $q_v$ (without renormalizing them over $S$), MT-MOPD defines the corrected sparse reverse-KL loss (Equation 5) as

Each summand is the scalar Bregman divergence generated by the negative entropy function $\phi(u) = u \log u - u$, which is nonnegative everywhere and vanishes if and only if $p_v = q_v$.

Why naive partial reverse KL fails

To see what goes wrong without the $-p_v + q_v$ correction, consider the naive truncated sum:

Differentiating with respect to a single selected probability $p_v$ ($v \in S$) while allowing the remaining mass $1 - \sum_{u \in S} p_u$ to flow into unscored tail tokens $V \setminus S$ yields

Setting this derivative to zero puts the stationary point at

rather than $p_v = q_v$!

For a concrete two-token example, suppose $q = (0.8, 0.2)$ and we score only the top-$1$ token $S = \{1\}$:

• At the true teacher match $p_1 = 0.8$, $L_{\mathrm{naive}} = 0.8 \log(0.8/0.8) = 0$.
• At $p_1 = 0.8 / e \approx 0.2943$ (which dumps $p_2 = 0.7057$ into the unscored second token), $L_{\mathrm{naive}} = 0.2943 \log(1/e) \approx -0.2943 < 0$.

Minimizing $L_{\mathrm{naive}}$ therefore drains student probability mass out of the teacher's top-$k$ tokens and hides it in the unscored tail so that $p_v \log(p_v/q_v)$ becomes negative. Whenever a training log shows a negative "top-$k$ KL loss," the code is computing $L_{\mathrm{naive}}$ without the normalization correction.

How the $-p_v + q_v$ correction fixes the optimum

Adding $-p_v + q_v$ to each coordinate gives the corrected summand

Now the first and second partial derivatives with respect to $p_v$ are

The $+1$ bias vanishes, $f(p_v, q_v)$ is strictly convex in $p_v > 0$, and its unique minimum occurs at $p_v = q_v$ where $f(q_v, q_v) = 0$. Writing $r = p_v / q_v > 0$ makes nonnegativity immediate:

since $r\log r - r + 1 \ge 0$ for all $r > 0$ with equality only at $r = 1$. In our two-token example with $q_1 = 0.8$, evaluating the corrected loss at the naive optimum $p_1 = 0.2943$ now yields $f(0.2943, 0.8) \approx -0.2943 - 0.2943 + 0.8 = +0.2114 > 0$, penalizing the mass evacuation and pulling $p_1$ back to $0.8$.

What the sparse loss still leaves unobserved in the tail

Restoring the coordinate-wise optimum on $S$ does not make $L_S(p,q)$ an unbiased estimator of the full-vocabulary reverse KL. Any redistribution of student probability mass among the unscored tail tokens $V \setminus S$ is invisible to $L_S$ as long as the total tail mass $1 - \sum_{v \in S} p_v$ stays fixed.