P$^2$-MOPD: Accelerating Multi-Teacher On-Policy Distillation via Parameter Priors

AI summary

Overview

Research area: Machine learning — large language model post-training, specifically multi-teacher on-policy distillation (MOPD) and model merging (task arithmetic).

Technical level: Intermediate. The paper assumes familiarity with policy-gradient-style distillation objectives, KL-based rewards, and weight-space model merging, but its central idea — making merge coefficients trainable alongside the student — is conceptually simple.

Scope: The paper proposes a single reparameterization of the distillation student that turns frozen expert weights into learnable "parameter priors," and evaluates it against MOPD and merging baselines across four domains at two model scales (Qwen3-1.7B and Qwen3-4B).

What This Paper Is About

Building a single generalist model from several domain-specialist experts is usually done either by multi-teacher on-policy distillation (expensive, teacher-supervised training) or by task arithmetic (free, but untrained and unreliable — a uniform merge recovers only 42.98% and 32.41% of the teacher gap at the 4B and 1.7B scales). The paper asks whether the partial success of merging can be turned into a training signal: if a linear combination of expert weights already carries useful capability, why not treat that combination as a trainable model and optimize its mixing coefficients together with the student's own weights under the same distillation loss?

Key Contributions

1. A differentiable "parameter prior" formulation. P²-MOPD writes each student weight matrix as its own trainable full weight plus a learnable linear combination of frozen expert task vectors, W_m = B_m + Σ_k α_{k,m} τ_{k,m}, and trains both B and the coefficients α under the unmodified MOPD objective. Because B stays unconstrained, the model class is unchanged — only the optimization path is biased toward expert directions.
2. A two-channel update with one backward pass. The same pre-update gradient feeds a student-weight channel (AdamW) and a coefficient channel (Adam on inner products ⟨G_m, τ_{k,m}⟩), with separate clipping and no weight decay on the weight channel, so the coefficient step is exactly equivalent to running the optimizer on B and recomposing.
3. Evidence that the gain comes from reachability, not extra capacity. Projections onto task vectors show student weight-only updates move only +0.03 along the lagging expert's direction while P²-MOPD moves +0.40 (4B) and +0.42 (1.7B), and controlled interventions at 1.7B show random directions in place of task vectors do not beat MOPD (TA init) at any checkpoint.
4. Faster and more balanced recovery at two scales, surpassing the final average recovery of MOPD (TA init) in half the training steps and raising worst-domain recovery from 47.2 to 89.4 (4B) and from 32.0 to 88.0 (1.7B).

Main Findings

• Faster recovery: P²-MOPD leads MOPD (TA init) — which differs only by freezing α at its task-arithmetic values — by +6.2 (4B) and +21.5 (1.7B) average-recovery points at step 25, and stays ahead in all 24 seed-paired comparisons (two scales, three seeds, four checkpoints). After half the training steps it reaches 92.3 (4B) and 90.5 (1.7B), already above the final 83.1 and 80.9 of MOPD (TA init) in every seed pair.
• Higher final average recovery: 94.85 ± 2.06 at 4B and 94.74 ± 0.91 at 1.7B, versus 83.13 ± 9.03 and 80.93 ± 3.02 for MOPD (TA init), 88.02 ± 1.33 and 67.81 ± 9.18 for IM-MOPD, and 84.20 and 79.79 for Open-MOPD. Mean raw score is higher by +2.66 (4B) and +2.17 (1.7B) points (paired, three seeds).
• Dramatically better balance: worst-domain recovery rises to 89.42 ± 2.54 (4B) and 88.01 ± 1.27 (1.7B), against 47.22 and 32.03 for MOPD (TA init); across-domain standard deviation of recovery falls from 21.8 to 4.6 (4B) and from 28.4 to 4.8 (1.7B).
• The gain is concentrated in the lagging domain: MedQA improves by +29.0 ± 8.3 (4B) and +48.2 ± 14.4 (1.7B) recovery points on average over all seed-paired comparisons, several times the difference in any other domain. CaseHOLD gains +3.8 and +4.5; FinQA is slightly lower in most pairs (−1.9 at 4B, +0.8 at 1.7B); IFBench differs in both directions with large spread (+4.8 ± 23.2, +7.6 ± 12.9).
• Distillation can erode what merging achieved: at 1.7B, MOPD (TA init) ends below the MedQA recovery of the task-arithmetic start (44%) in all three seeds, and plain MOPD from the base model collapses on MedQA (5.21 raw score, −15.74 average recovery).
• Coefficients and weights divide the work: the coefficient channel accounts for 81% (4B) and 84% (1.7B) of the displacement energy, while the student weight channel moves almost entirely outside the task-vector span (only 1.7% and 4.0% of its displacement lies inside). Total in-span energy share is 5.2%/10.7% for MOPD (TA init) versus 85.4%/90.0% for P²-MOPD.
• Cross-domain conflict becomes a few interpretable scalars: the medical gradient lies 98% on α_med and has cosine −0.78/−0.63 with the finance/IF gradients at 4B (−0.85/−0.81 at 1.7B); α_med rises from 0.25 to 0.63 (4B) and 0.65 (1.7B). The losing domains pay a small price (the slight FinQA deficit).
• Interventions isolate the mechanism (1.7B, Table 3): random directions give MedQA recovery 1.4 and average 70.2; learning only the α_med row gives 94.1 MedQA and 95.0 average (close to the full 87.2/95.8); freezing α_med returns MedQA to 24.2 versus 22.1 for MOPD (TA init). Training on medical prompts alone does not rescue MOPD (TA init), whose MedQA recovery falls to −6.9.
• Modest cost, large savings: each step is 1.24× (4B) and 1.09× (1.7B) the wall-clock of MOPD (TA init) — 225.6 versus 180.6 seconds, and 167.4 versus 153.9 seconds — but the time to exceed the baseline's final average recovery is at most 0.31× and 0.55×, since checkpoints are 25 steps apart. The actor update lengthens from 14.7 to 73.8 s (4B) and from 11.1 to 37.1 s (1.7B), but rollout generation dominates each step.

Methodology in Plain English

The setup follows the IM-MOPD protocol. A general SFT reference model serves as the base, and four domain teachers (medical, law, finance, instruction following) are fine-tuned from it, so each teacher's "task vector" — its weights minus the base's — is well defined. A training batch holds 128 prompts, 32 per domain, and each prompt is scored by the teacher of its own domain on responses the student itself generates. Runs last 100 steps, with checkpoints at 25, 50, 75 and 100 scored by the official IM-MOPD harness using three samples per item at temperature 1.0 and evaluation seeds 42, 43 and 44. Performance is reported as recovery, the fraction of the base-to-teacher gap that is closed, plus worst-domain recovery and the area under the learning curve.

The methodological trick is a change of variables. Instead of treating the merged model as a starting point, the authors define every attention and MLP projection as a trainable weight plus a weighted sum of the frozen experts' task vectors, with one scalar per expert per module. Everything else — embeddings, normalization weights, the tied output head — is trained exactly as in MOPD. Initialization sets the trainable part to the base weights and the coefficients to the uniform merge values (0.25 each), so the model starts precisely at the task-arithmetic merge. During training, the gradient with respect to a coefficient is just the inner product of the ordinary loss gradient with that expert's task vector, which measures whether nudging toward that expert would reduce the loss on the current batch. Both channels are updated from the same gradient computation, so there is no second backward pass; the coefficients can alternatively be tied across a layer or the whole model, though the main experiments use per-module coefficients. The authors also show that, under plain SGD assumptions, the scheme is equivalent to preconditioning the gradient with a positive-semidefinite, rank-at-most-K term that amplifies movement along expert directions.

To diagnose why this helps, they project each method's final displacement from the merge onto the four task vectors, and they run a set of one-factor-at-a-time interventions at 1.7B: replacing task vectors with random directions, training on medical prompts only, learning only the medical coefficient row, and freezing that row.

Why This Matters

Impact on research. The paper reframes model merging from an operation that outputs a model into a parameterization that shapes optimization, and shows the two views can be combined inside a single distillation objective. It also offers a concrete diagnosis of why MOPD recovers domains unevenly: student weight-only updates simply do not travel far along the task

Authors’ abstract

Multi-teacher on-policy distillation consolidates expert capabilities by supervising student-generated trajectories, whereas task arithmetic offers a training-free alternative through linear parameter combinations. We view the linear combination of student and expert parameters as a single differentiable model and propose P$^2$-MOPD, which jointly optimizes the mixing coefficients and the student weights under the same on-policy distillation objective. Unlike task arithmetic, which stays within the span of the frozen student and expert parameters for any choice of merging coefficients, P$^2$-MOPD keeps the student trainable, so the experts act as parameter priors without confining it to this span. Across four domains at 1.7B and 4B scales, P$^2$-MOPD achieves faster and more balanced recovery than MOPD from the same initialization, surpassing its final average recovery in half the training steps. At the end of training, our 1.7B model reaches 94.7% average recovery, versus 80.9% for MOPD.