Latent-MOPD: Latent Multi-Teacher On-Policy Distillation

AI summary

Overview

Research area: Machine learning / large language model training, specifically on-policy distillation and multi-teacher knowledge transfer.

Technical level: Advanced. The paper assumes familiarity with RL fine-tuning of LLMs (policy-gradient objectives, on-policy rollouts), knowledge distillation, hidden-state representation matching, and centered kernel alignment diagnostics.

Scope: The paper introduces Latent-MOPD, a method that trains one student LLM from several frozen domain specialists by supervising it on both the teachers' next-token distributions and their internal hidden states, and reports results across nine same-family and six cross-family benchmarks.

What This Paper Is About

Existing multi-teacher on-policy distillation (MOPD) trains a student on its own generated responses, routing each response to a domain specialist and learning only from what that specialist predicts at the output layer. The authors argue that a teacher's internal hidden states carry information about how it computes those predictions, and that no prior work has used those representations when distilling from multiple specialists at once.

The goal is a single deployable student that combines the complementary strengths of several existing domain specialists (math, code, logic) without training any new teachers, by aligning to each routed teacher's late-layer hidden states as well as its token probabilities.

Key Contributions

1. A representation-level multi-teacher OPD method for LLMs. To the authors' knowledge, Latent-MOPD is the first multi-teacher on-policy distillation method that integrates specialists through both their predictions and the hidden states used to compute them, with all teachers frozen and no additional teacher training.

2. Same-family gains beyond individual teachers. A 1.5B student outperforms token-only, representation-only, and uniform-averaging baselines on all nine benchmarks spanning math, code, and logic, and surpasses the per-benchmark best teacher on five of nine at the same model size. Crossfade improves eight of nine benchmarks over constant joint weighting with other settings fixed.

3. Transfer across teacher panels and initializations. With separately developed 7B Qwen-based teachers, Latent-MOPD outperforms both single-channel baselines on all six benchmarks. Starting from a parameter-merged initialization, 62 distillation steps improve five of six benchmarks.

4. Diagnostics for organizing multi-teacher representation supervision. The paper shows, using native CKA before training and a controlled ablation, that layer selection should depend on the teacher–student relationship, and that grouping routed targets into domain-pure updates stabilizes all-layer representation-only training while interleaved updates collapse.

Main Findings

• Same-family benchmark sweep (Table 1). With three RL-tuned 1.5B same-family teachers and the default last-three-layer setting, Latent-MOPD (last 3) reaches a Norm score of 1.05, above the best-teacher envelope of 1, compared with 0.90 for MOPD-style (token-only), 0.94 for Latent-MOPD (OPRD-style) and Latent-MOPD (constant), 0.99 for Latent-MOPD (last 1), 0.79 for uniform averaging, and 0.64–0.73 for representation-only baselines.

• Specific same-family improvements over token-only MOPD-style. Minerva rises from 33.5 to 36.3, LiveCodeBench-easy from 68.3 to 73.1, and MuSR from 50.7 to 52.4, using the same prompts and number of updates.

• Beating the individual specialists. The default last-three-layer student surpasses the per-benchmark best teacher on five of nine benchmarks: BBH, MuSR, MBPP, MBPP+, and Minerva. Those gains span logic, code, and math in one student with the same parameter count as each teacher.

• Cross-family transfer (Table 2). Using Skywork-OR1-Math-7B, AceReason-Nemotron-7B, and DeepSeek-R1-Distill-Qwen-7B (each roughly 2.3x the student's hidden width) distilled into the same 1.5B student, the default last-layer Latent-MOPD with a linear map outperforms token-only and representation-only baselines on all six benchmarks (Norm 0.26, versus 0.16 for MOPD-style and −0.05 for Rep-only (last 1)). Relative to MOPD-style, AIME24 improves from 36.5 to 39.4 and MBPP from 55.4 to 57.9; GYM rises from 29.2 to 34.1 and BBH from 45.8 to 49.5. A two-layer MLP also beats both single-channel baselines on all six (Norm 0.27), with the two projector configurations splitting benchmark wins three to three.

• Layer choice depends on the teacher panel. Same-family, three-layer supervision raises Norm from 0.99 to 1.05 and improves eight of nine benchmarks; representation-only controls also favor three layers over last-layer and all-layer alternatives on seven of nine benchmarks each (Norm 0.73 vs 0.64). Cross-family with linear projectors, last-layer supervision reaches Norm 0.26 versus 0.25 for three layers, splitting wins three to three while retaining the aggregate gain with fewer targets.

• Batch organization determines whether all-layer representation-only training survives. Holding teachers, loss, coefficient, and budget fixed, domain-pure blocks of 32 complete all 62 steps and reach 91.7 on in-loop MATH-500 validation. Interleaving the same rows collapses MATH-500 to 22.8 at step 20 and 9.8 at step 30, and it remains far below the base at step 62 (43.2 vs 84.1).

• Cross-family representation-only supervision is unstable. In the cross-family setting, representation-only supervision falls below the student's initial performance in aggregate (Norm −0.05), motivating a schedule rather than sustained latent supervision.

• Training dynamics. On MATH-500 with 8 samples per problem, the base score is 84.1; the same-family run reaches 90.3 at step 20 and 92.0 at step 62, while cross-family reaches 86.2 and 87.7 respectively. Both remain above the base from step 20 through the later token-only phase.

• Merge initialization. From a uniform parameter merge of the same-family teachers (Norm 0.87), Latent-MOPD improves five of six benchmarks: GYM from 46.3 to 52.2, BBH from 60.6 to 63.4, and AIME24 from 48.1 to 53.8. Norm rises from 0.87 to 1.03, above token-only (0.97) and representation-only (0.91) training from the same merge. The student exceeds the best same-family teacher on BBH, MBPP, and MBPP+.

• General-ability retention. On three suites absent from the training pool (ARC-Challenge, WinoGrande, HellaSwag), the same-family student's scores remain within 1.1 points of the base under 0-shot likelihood scoring; ARC-Challenge and WinoGrande improve, while HellaSwag changes from 44.7 to 44.2.

• Representation loss declines and stays low. Normalized to each domain's first logged value, the math, code, and logic losses fall to 0.30, 0.38, and 0.39 at update 10 per domain, and reach 0.13, 0.19, and 0.15 at their final recorded updates, with β excluded from the logged loss.

• Matching the routed specialist after training. On teacher-disagreement positions, the math row's maximum similarity shifts from the code teacher under MOPD-style to the math teacher under Latent-MOPD, while code and logic retain their domain-matched maxima. The math matching-teacher margin increases by 0.0648 on selected positions and 0.0327 over all positions, with positive paired 95% instance-bootstrap intervals in both cases.

• Storage cost of layer selection. Selecting 3 of 28 layers requires 89.3% less selected-state storage at fixed token count and precision: for 16,384 positions and BF16 width-1,536 states, the student and teacher tensor copies total 0.28 GiB for three layers versus 2.63 GiB for all 28.

Methodology in Plain English

The setup starts from a base student model, DeepSeek-R1-Distill-Qwen-1.5B, and three existing frozen specialists, one for math, one for code, and one for logic. Each training prompt carries a domain label, and the student generates its own responses. Because the student is scored on its own rollouts, this is on-policy distillation rather than training on a fixed dataset.

For every response, the domain's teacher supplies two kinds of supervision. The first is the token channel: the teacher's next-token distribution scores the student's top-16 candidate tokens using teacher-to-student log-probability ratios, with those ratios converted into fixed advantages that are scaled by the domain's masked standard deviation. The second is the representation channel: the student's hidden state at selected layers is compared to the same teacher's hidden state, after L2 normalization, using a squared distance.

Three design choices organize that representation supervision. Layer selection follows the teacher–student relationship: within a shared model lineage, the authors align the final three layers; across families, where correspondence at intermediate depths is weak, they align only the final state, which has the same functional role in both models (it directly feeds the next-token head). Width bridging handles the fact that same-family teachers share the student's hidden width (identity map) while cross-family teachers are wider, requiring a trainable linear map shared across routed teachers and layers, initialized by a ridge fit to domain-routed student–teacher state pairs. Batch organization means each optimizer update contains prompts from one domain only, and the three domains rotate in a math–code–logic order without shuffling.

The other key ingredient is the crossfade. Rather than keeping both channels weighted equally throughout training, each teacher's representation weight β(u) decays from 1 to 0 over its own window of completed updates, while the token weight α(u) rises from 0 to 1 over the same window. The window is 10 per-teacher steps same-family and 7 cross-family, spanning 30 and 21 of the 62 global optimizer steps under the three-teacher rotation. Each teacher sees the same sequence of weights, and the representation term is switched off entirely once its window closes.

Training uses 62 optimizer steps, 32 prompts per step with 4 responses each, a constant learning rate of 1e-5, and AdamW. Prompts come from DAPO-Math-17k, OpenCodeReasoning, and Reasoning Gym, pooled to 17,856 prompts (5,952 per domain). The 62 updates process 1,984 prompt presentations and generate 7,936 responses with no verifier rewards. Evaluation covers the teachers, the base student, a token-only MOPD-style baseline, representation-only baselines (OPRD-style over all 28 layers, and last-1 and last-3 variants), uniform averaging, and Latent-MOPD variants. A separate experiment initializes the student from the uniform weight average of the same-family teachers and then distills for 62 steps.

Why This Matters

Impact on research. The paper extends multi-teacher distillation from the output level to the representation level, and provides evidence that how representation supervision is scheduled and batched matters as much as how it is selected. The domain-pure versus interleaved contrast is a concrete result about optimizer-step composition in multi-teacher training, and the crossfade provides a transferable recipe for phasing out a latent objective. Reusing public specialist checkpoints without retraining them is also a practical alternative to training one monolithic model.

Potential applications:

• Building a single production model that covers math, coding, and logic by reusing separate specialist checkpoints rather than maintaining several deployed models.
• Reducing serving cost by consolidating three domain experts into a student with the same parameter count as one of them, which the same-family results support.
• Recycling previously trained or publicly released RL-tuned checkpoints into new models without access to their original training data or pipelines.
• Improving distillation transfer when teachers and students come from different model scales or post-training lineages, as tested in the cross-family setting with 7B teachers and a 1.5B

Authors’ abstract

On-policy distillation (OPD) trains a student on the responses it generates. Existing LLM multi-teacher OPD transfers what specialists predict through their output distributions. We introduce Latent-MOPD, to our knowledge the first representation-level multi-teacher OPD method for LLMs. It integrates existing specialists through both their predictions and the hidden states used to compute them, without additional teacher training. To coordinate representation supervision from multiple specialists, we select late-layer targets according to the teacher-student relationship, bridge unequal hidden widths with a shared projection, and group updates by domain. Each teacher's supervision gradually shifts from hidden states to token predictions, with both channels using the same routed specialist. In our main same-family setting, Latent-MOPD outperforms the token-only, representation-only and uniform-averaging baselines on all nine benchmarks across math, code and logic. With the same parameter count as each teacher, the student also surpasses the per-benchmark best teacher on a majority of these benchmarks. With larger, separately developed cross-family teachers, Latent-MOPD outperforms both single-channel baselines on all benchmarks. A same-family all-layer representation-only control remains stable with domain-pure updates but collapses when teacher domains are interleaved within an update. Our results show that a single student can integrate capabilities from several specialists through both their output distributions and internal representations.