PMOPD: Task Ordering, Cycling, and Parameter-Update Subspace Protection in Multi-Teacher On-Policy Distillation
Youzhi Liu, Ruobing Zheng, Boyuan Tong, Tianqi Li, Pingqi Li, Hanbo Bi, Yi Yuan, Jingdong Chen
cs.LG
2026-09-28
PMOPD projects gradients and Adafactor updates off protected task subspaces; vs MOPD, 3-task avg rises 2.54 on Qwen2.5-7B and 2.09 on Llama-3.1-8B, every task up.
Post-training for frontier models often trains separate math, reasoning, and code teachers, then distills them into one student. Multi-teacher on-policy distillation (MOPD) skips inference-time routing and multi-model serving: teachers give token-level targets on trajectories the student actually samples.
Most OPD papers tune a single task through the loss, the distillation scope, and the teacher signal. Several teachers writing into the same weights need not agree. Gains in one domain suppress another. Mixing tasks inside a mini-batch only shortens the gap between domains. It does not name conflicting update directions, and it does not protect useful displacements an earlier teacher already wrote into the weights.
The object that matters is the cumulative parameter displacement ΔW over a full task block. A single mini-batch gradient is too noisy to serve as a memory. For every 2-D weight matrix, a truncated SVD keeps the top-16 right singular vectors. Protection multiplies on the right of the gradient and the optimizer update, so the method stores right singular vectors. By 20% of training those subspaces already match the final ones at average similarity 0.62: the dominant directions lock early. Pairwise similarities among Math, Code, and Reason top-16 subspaces sit between 0.133 and 0.151. The small overlap is where a later task steps on an earlier skill.
Figure 1 plots preferred next steps on a shared PCA plane. From the same checkpoint the three arrows point far apart. The measured MOPD path wobbles on the joint loss surface. After stripping components along protected directions, PMOPD walks more steadily into a shared low-loss region.
PMOPD runs reverse-KL OPD in ordered task blocks. Inside each cycle:
The memory stores the displacement that actually happened after clipping and adaptive steps, unlike PCGrad's surgery on instantaneous gradients. Unlike GPM, the basis does not come from activations.
Task order comes from a cheap conflict probe. Build memory from task a, measure the fraction of task-b's gradient that would be removed, average both directions, then average each task against the others. Lower conflict trains first. Here that ranking is Code 22.34%, Reason 28.50%, Math 29.38%, which is Code→Reason→Math. A 60-example probe recovers the same order in all ten folds.
Cycle count trades two failures under a fixed prompt budget. Long blocks let the trajectory drift between revisits. Short blocks make ΔW estimates noisy. Four cycles hit the peak same-task subspace similarity across revisits (mean 0.525) and the peak evaluation average: 66.97 versus 65.81 with one cycle and 65.37 with ten.
Two bases, Qwen2.5-7B and Llama-3.1-8B. Teachers start from the same checkpoint and specialize with in-domain RL, so differences come from data, not unrelated pretraining. Each task has 600 OPD prompts; under four cycles each block has 150. Batch size is 16. Tests are the orca-math split of NuminaMath-CoT (450), CommonsenseQA (1221), and APPS (300). PMOPD is averaged over three seeds. Every OPD run uses reverse KL and the same prompt budget.
| Method | Qwen avg | Llama avg |
| Base student | 50.55 | 32.41 |
| Parameter merge | 58.56 | 35.29 |
| MOPD (within-batch mix) | 64.43 | 38.95 |
| BB-MOPD (between-batch) | 63.89 | 36.29 |
| Open-MOPD | 63.54 | 35.14 |
| PMOPD | 66.97 | 41.04 |
Against MOPD on Qwen, Math moves 65.56→67.78 (+2.22), Reason 71.74→73.46 (+1.72), Code 56.00→59.67 (+3.67). On Llama, Math 18.22→21.33 (+3.11), Reason 68.63→69.45 (+0.82), Code 30.00→32.33 (+2.33). Every task rises.
On Qwen, PMOPD also beats each specialist on that specialist's own task: Math 67.78 vs 66.89, Reason 73.46 vs 70.35, Code 59.67 vs 58.00. On Llama it beats the Math and Reason teachers but still trails the Code teacher at 42.00.
Under the same order and four-cycle budget, turning projection off drops the Qwen average to 63.85. Gradient projection alone reaches 66.22. Adding optimizer-update projection reaches 66.97, 3.12 above the no-projection run. The six one-cycle orders span 2.97 points; the probe's C→R→M still wins at 65.81. The strongest mixing variant, within-batch ordered, scores 64.83. Full PMOPD stays 2.14 above that.
Teams already running MOPD mostly mix or interleave batches. This paper adds a geometric constraint that is cheap to implement: task updates live in nearly separate subspaces, and the overlap is small enough to block with a top-16 memory per matrix. The second projection exists because Adafactor can undo a gradient-only correction.
The gains are incremental. Qwen's average goes 64.43→66.97, Llama's 38.95→41.04. Llama Math is still 21.33, and Reason only beats MOPD by 0.82. The setting that fits is same-base teachers, a handful of tasks, and a willingness to train in blocks. It does not replace teacher quality, and it does not answer whether production serving should still route.
There is no Limitations section. The study covers three domains, 7B and 8B bases, and 600 OPD prompts per task. Teacher differences come from in-domain RL on one checkpoint, not from experts with distant pretraining histories. K=16 is inherited from prior OPD geometry work and is not swept here.
The probe is meant to skip factorial search when the task set grows. This paper has exactly three tasks, so exhaustive search is six runs. Whether the probe still picks the best order on larger sets is untested. Wall-clock cost of randomized SVD and the extra projection is not reported. Llama's absolute scores are low; the Math +3.11 sits in the narrow band 18.22→21.33, so transfer to stronger bases is an open bet. The main table does not include PCGrad or GPM as direct baselines, only a mechanistic contrast in related work.