PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2609.34605v1cs.LG

TL;DR

Multi-teacher on-policy distillation must integrate capabilities through shared parameters even when improving one domain can suppress another. PMOPD protects task-update subspaces while organizing teacher updates, and improves every evaluated capability over MOPD. Across three tasks, average scores rise by 2.54 points on Qwen2.5-7B and 2.09 points on Llama-3.1-8B.

  • Problem

    MOPD aggregates multiple capabilities in shared parameters, where incompatible objectives can make improving one capability suppress another.

    Source: §1 · p. 2

  • Method

    PMOPD builds subspace memories from cumulative task displacements and projects gradients and optimizer updates away from components that interfere with protected task directions.

    Source: §1 · p. 2, §8 · p. 9

  • Results

    PMOPD improves every evaluated capability over MOPD, raising average scores across Code, Reason, and Math by 2.54 points on Qwen2.5-7B and 2.09 points on Llama-3.1-8B.

    Source: §8 · p. 10

  • Takeaways & Limitations

    Parameter-update subspace protection is an effective and transferable strategy for integrating specialized teachers into a single balanced model.

    Source: §8 · p. 10

Abstract

from arXiv · show

Multi-teacher on-policy distillation (MOPD) has emerged as a popular post-training paradigm for integrating specialized capabilities in frontier language models. Existing OPD research has primarily focused on optimizing single-task distillation through objective design, distillation scope, and teacher signal construction, whereas MOPD must aggregate multiple capabilities in shared parameters and address the resulting capability seesaw, in which improving one domain suppresses capabilities acquired from another. Inspired by the distinctive update geometry of OPD, we find that parameter updates from different tasks rapidly concentrate in their respective low-dimensional subspaces during MOPD, providing a direct geometric basis for identifying and controlling cross-task interference. We therefore propose PMOPD (Projection-based Multi-Teacher On-Policy Distillation), which constructs subspace memories from the cumulative parameter displacements of different tasks and projects both gradients and optimizer updates to remove components that interfere with protected task directions. We further develop a lightweight conflict probe to characterize task interactions and guide task ordering, together with a cycling strategy that balances subspace estimation and timely task revisitation. Experiments on representative Code, Reason, and Math tasks show that PMOPD improves every evaluated capability over MOPD, raising the average score across the three tasks by 2.54 points on Qwen2.5-7B and 2.09 points on Llama-3.1-8B. These consistent gains establish geometry-aware optimization as an effective and transferable approach to balanced multi-teacher distillation.

1 INTRODUCTION

MOPD must integrate specialized capabilities through shared parameters despite potentially conflicting task objectives. PMOPD uses task-update geometry to protect acquired capabilities, organize training, and improve results across the evaluated tasks.

1 INTRODUCTION

  • MOPD's shared parameters create a capability seesaw: improving one domain can suppress capabilities acquired from another, and standard task mixing does not protect useful earlier changes.The stated challenge is to preserve acquired capabilities while retaining enough freedom to learn new ones.Source: §1 · p. 2
  • Cumulative OPD displacements form early-stabilizing, within-task-consistent low-dimensional subspaces with low cross-task overlap, offering a compact target for interference control.This approach avoids treating each noisy mini-batch gradient independently and does not globally freeze the parameter space.Source: §1 · p. 2
  • PMOPD builds per-matrix memories from completed-task displacements and projects both gradients and preconditioned optimizer updates to remove components interfering with protected directions.Projecting optimizer updates matters because adaptive preconditioning can rotate a projected gradient back toward protected subspaces.Source: §1 · p. 2
  • A bidirectional conflict probe ranks tasks, while cycling balances timely teacher revisitation against reliable block-displacement estimates; in this setting, Code→Reason→Math and four cycles perform best.The probe measures directional pairwise interference before training.Source: §1 · p. 3
  • Across Qwen2.5-7B and Llama-3.1-8B, PMOPD raises average scores across three evaluated tasks over MOPD by 2.54 and 2.09 points, respectively, improving every evaluated task in both model families.Source: §1 · p. 3

2 RELATED WORK

The related work situates PMOPD among on-policy distillation, multi-task optimization, and continual-learning approaches. Its distinction is to constrain later optimization using realized OPD block displacements and their task-specific geometry.

2.1 KNOWLEDGE DISTILLATION AND ON-POLICY SUPERVISION

  • OPD queries teachers on student-generated trajectories, supplying token-level targets for visited prefixes and reducing mismatch between training and inference states.Prior MOPD systems use multiple teachers to integrate domain specialists; this work studies interference through shared-parameter update geometry.Source: §2.1 · p. 3

2.2 MULTI-TASK OPTIMIZATION AND CAPABILITY INTEGRATION

2.3 CONTINUAL LEARNING AND LOW-DIMENSIONAL UPDATE GEOMETRY

  • Continual-learning methods use penalties, episodic constraints, or activation-derived projections; unlike GPM, PMOPD derives per-matrix bases from realized OPD displacements and rebuilds memory each cycle.GPM extracts protected bases from activations.Source: §2.3 · p. 3
  • Cumulative OPD updates show early subspace locking, strong same-task alignment, and low cross-task overlap, making overlapping directions a tractable target for interference control.The related analysis describes low-dimensional adaptation in intrinsic-dimension methods and LoRA (Aghajanyan et al., 2021; Hu et al., 2022).Source: §2.3 · pp. 3–4

3 PROBLEM FORMULATION AND GEOMETRIC MOTIVATION

The formulation describes a shared student trained with specialist teachers on student-generated prefixes, then characterizes task updates through cumulative parameter displacements and their subspaces. The reported geometry shows early within-task stability and separation across task domains.

3.1 MULTI-TEACHER ON-POLICY DISTILLATION

  • The task set is T = {Math, Reason, Code}, with D_t denoting the prompt distribution for task t.The formulation also names the shared student π_θ and corresponding specialist teacher π_t^T.Source: §3.1 · p. 4
  • For prompts from a task distribution, the shared student samples a response and that task's teacher is evaluated on the same student-generated prefixes.This follows the stated on-policy distillation setup (Agarwal et al., 2024; Ma et al., 2026).Source: §3.1 · p. 4

3.2 TASK-UPDATE SUBSPACES

  • For a weight matrix, the accumulated task-block update is the difference between its after- and before-block weights, and its singular value decomposition supplies the update geometry.The passage defines the displacement as ΔW^t = W_after − W_before.Source: §3.2 · p. 4
  • The method retains the top-K right-singular vectors as a compact task-update basis and uses subspace overlap to compare two blocks.The overlap formula itself is not needed to state the role described in the passage.Source: §3.2 · p. 4
  • With K = 16, subspaces extracted after 20% of training average 0.62 similarity to their corresponding final task subspaces across Math, Reason, and Code.The passage reports this as evidence that each task's dominant update subspace emerges early and remains stable.Source: §3.2 · p. 4
  • Final top-16 update-subspace similarities across Math, Code, and Reason range from 0.133 to 0.151, indicating substantially different parameter-update subspaces.The passage identifies overlapping directions as the directions in which later updates interfere with protected capabilities.Source: §3.2 · p. 4

4 METHOD

PMOPD trains a shared student in ordered task blocks, builds per-matrix memories from completed-task displacements, and projects later gradients and optimizer updates against those memories. Each cycle rebuilds memory from empty, while the method uses reverse-KL on-policy distillation and extracts task-update subspaces.

4.1 OVERVIEW

  • Each cycle trains Code, Reason, then Math; the first task starts with empty memory, and completed block displacements build protection for later tasks.The final student initializes the next cycle, where memory is rebuilt from new task blocks.Source: §4.1 · p. 4, Figure 2 · p. 5

4.2 REVERSE-KL ON-POLICY DISTILLATION

  • The student samples responses and each task’s teacher evaluates the student-generated prefixes, with reverse KL used in the main experiments.Task identity selects which specialized teacher supplies the target distribution.Source: §4.2 · p. 5

4.3 EXTRACTING TASK-UPDATE SUBSPACES

  • For each trainable matrix, PMOPD extracts dominant input-side directions from the cumulative displacement across a task block using randomized truncated SVD.The displacement includes effects of the task loss, gradient clipping, and the optimizer; right singular vectors are used because protection acts on the right side of updates.Source: §4.3 · pp. 5–6

4.4 ORTHOGONAL TASK-SUBSPACE MEMORY

  • After a remembered task block, its new directions are concatenated with the existing memory and orthogonalized; a redundancy threshold removes numerically redundant directions.The resulting orthonormal basis defines the orthogonal projector onto the protected input-side subspace.Source: §4.4 · p. 6

4.5 GRADIENT AND OPTIMIZER-UPDATE PROJECTION

  • PMOPD projects gradients before global norm clipping and also projects the preconditioned Adafactor update, which can otherwise rotate a projected gradient toward protected directions.The projected update does not directly change the transformation along inputs in the protected subspace, at the level of that linear map.Source: §4.5 · p. 6

5 EXPERIMENTAL SETUP

The experiments integrate Math, Reason, and Code teachers specialized from shared base checkpoints on two model families. They compare PMOPD with shared-model and multi-teacher baselines using task-specific scores and their unweighted average.

5.1 MODELS, TASKS, AND TEACHERS

  • Qwen2.5-7B and Llama-3.1-8B teachers share their family’s base checkpoint and are specialized by domain-specific reinforcement learning.The Qwen setup uses NuminaMath-CoT for Math, CommonsenseQA for Reason, and APPS for Code, with 600 OPD prompts per task.Source: §5.1 · p. 6, Table 2 · p. 6

5.2 BASELINES AND EVALUATION

  • Comparisons include the base model, specialists, parameter merging, MOPD, BB-MOPD, and Open-MOPD; OPD baselines use reverse KL and the same total task prompts.PMOPD uses Code→Reason→Math with four cycles and batch size 16; task-specific percentage scores and their unweighted average are reported, with PMOPD results averaged over three seeds.Source: §5.2 · p. 7, Table 3 · p. 7

6 MAIN RESULTS

PMOPD outperforms MOPD across all tasks in both model families, while ablations show benefits from gradient projection and further gains from projecting optimizer updates.

6.1 QWEN2.5-7B; 6.2 LLAMA-3.1-8B

  • PMOPD raises the three-task average over MOPD by 2.54 points on Qwen2.5-7B and 2.09 points on Llama-3.1-8B, improving every evaluated task.It achieves the best score in every Qwen task column and the best integrated-model average in both families.Source: §6.1 · p. 7, §6.2 · p. 7

6.3 PROJECTION ABLATION

  • Under the matched order, cycle schedule, data, and optimization budget, gradient projection raises the average from 63.85 to 66.22, and optimizer-update projection further raises it to 66.97.The complete pipeline outperforms the matched no-projection variant by 3.12 points.Source: Table 4 · p. 7, §6.3 · pp. 7–8
  • Gradient projection particularly improves Reason and Code, while optimizer-update projection yields the strongest Math and Code results.The ablation isolates gradient projection and optimizer-update projection under the same order and cycle schedule.Source: Table 4 · p. 7, §6.3 · p. 8

7 ANALYSIS

The analysis uses geometric diagnostics to guide task ordering and cycle count. Code→Reason→Math ranks highest in the order sweep, and four cycles produce the strongest average score and most consistent subspace estimates.

7 ANALYSIS

  • The analysis connects lightweight geometric diagnostics to selecting task order and cycle count before a full MOPD run.The diagnostics are motivated by the increasing cost of exhaustive search as task count or training scale grows.Source: §7 · p. 8

7.1 SELECTING TASK ORDER BY AVERAGE UNDIRECTED CONFLICT

  • Among six one-cycle task orders, Code→Reason→Math achieves the highest average of 65.81, with order averages spanning 2.97 points.The analysis introduces average undirected conflict to explain the preferred order and avoid exhaustive enumeration as task sets or training scale grow.Source: §7.1 · p. 8, Table 5 · p. 8
  • The average undirected conflict scores are 22.34% for Code, 28.50% for Reason, and 29.38% for Math, yielding the same low-to-high order as the strongest exhaustive-search result.Lower-conflict tasks are placed earlier so accumulated memory is less broadly conflicting with later optimization.Source: §7.1 · p. 8
  • Using 60 examples per fold, all ten probe folds recover Code < Reason < Math, with mean estimates of 28.04%, 35.75%, and 36.97%, respectively.The compact probe selects the best observed order before full training, avoiding exhaustive permutation evaluation.Source: §7.1 · pp. 8–9, Figure 3 · p. 8

7.2 SELECTING THE NUMBER OF CYCLES BY SUBSPACE CONSISTENCY

  • With task order and total prompts fixed, the average score rises from 65.81 at one cycle to 66.97 at four, then falls to 65.37 at ten cycles.The evaluated schedules include 1, 2, 4, 5, 8, and 10 cycles.Source: §7.2 · p. 9, Figure 4 · p. 9
  • Four cycles yield the highest same-task subspace similarity, 0.525, and balance reliable subspace estimation with timely revisitation.Similarity is 0.487 for two cycles, then decreases to 0.488, 0.452, and 0.424 for 5, 8, and 10 cycles; high cycle counts shorten displacement-estimation blocks, while low counts increase trajectory drift between revisits.Source: §7.2 · p. 9

8 CONCLUSION

PMOPD uses task-update subspaces to protect against interference and improves every evaluated capability over MOPD. The results support subspace protection as an effective and transferable strategy for integrating specialized teachers into one balanced model.

8 CONCLUSION

  • PMOPD constructs memories from realized task displacements and projects gradients and adaptive-optimizer updates away from protected directions.Task updates rapidly concentrate in distinct low-dimensional subspaces, with overlapping components providing a target for interference control.Source: §8 · p. 9
  • PMOPD improves every evaluated capability over MOPD, raising the three-task average by 2.54 points on Qwen2.5-7B and 2.09 points on Llama-3.1-8B.The paper concludes that parameter-update subspace protection is effective and transferable for integrating specialized teachers into a single balanced model.Source: §8 · p. 10

AI USAGE DISCLOSURE

The authors report using generative AI tools during manuscript preparation to polish English and assist with literature retrieval and discovery.

AI USAGE DISCLOSURE

  • Generative AI tools assisted with English polishing and retrieval and discovery of relevant literature during manuscript preparation.The authors report reviewing and revising all AI-assisted outputs.Source: AI USAGE DISCLOSURE · p. 10
  • The authors state that they retained full control over the scientific content and citation selection and take responsibility for the final manuscript.Source: AI USAGE DISCLOSURE · p. 10

A.1 CONFLICT ESTIMATES AND DIRECTIONAL MEASUREMENTS

Conflict probes produce a stable task ordering despite directional asymmetry; four-cycle scheduling maximizes revisit consistency, while complete PMOPD outperforms task mixing alone.

A.1 CONFLICT ESTIMATES AND DIRECTIONAL MEASUREMENTS

  • All ten 60-example folds recover Code < Reason < Math, while directional removed-gradient ratios differ for every task pair.The two directions are measured by constructing memory from one task and measuring the removed-gradient ratio on the other.Source: Table 6 · p. 12, §A.1 · p. 12, Table 7 · p. 12

A.2 PER-TASK REVISIT CONSISTENCY

  • The four-cycle schedule achieves the highest revisit similarity for every task and the highest mean, providing task-wise support for the consistency criterion used to select the cycle count.Source: §A.2 · p. 13

A.3 MOPD MIXING VARIANTS

  • 64.83 for the strongest task-mixing variant versus 66.97 for complete PMOPD, a 2.14-point advantage on Qwen2.5-7B.The comparison is between the strongest mixing variant and the complete PMOPD configuration.Source: Table 9 · p. 13, §A.3 · p. 13
Loading 2609.34605v1…