USA: Update-aware SAM for Cross-domain On-Policy Disitllation of Language Agents

2026-09-28

AI Summary

This study addresses the issue of negative transfer and performance degradation during model merging in multi-domain knowledge distillation, caused by coupled cross-domain parameter updates. To mitigate this, we propose USA, a method that for the first time quantifies cross-domain update coupling as a perturbation radius. Built upon an update-aware SAM optimizer, USA adaptively determines the perturbation radius by measuring parameter update magnitudes during a warm-up phase. It specifically suppresses merging interference along high-displacement coordinates and reduces curvature in critical dimensions to alleviate merging conflicts. By integrating online policy distillation with flat minima optimization, our approach comprehensively outperforms single-domain baselines across mathematics, science, and code tasks, achieving average improvements exceeding four percentage points while effectively eliminating negative transfer among conflicting domains.

Abstract

On-policy distillation instils multi-turn agentic reasoning through dense token-level supervision on the student's own trajectories, but a single domain saturates early, so further supervision has to be drawn from other domains. Multi-domain data mixing is the most direct way of incorporating them, at the cost of conflicts between their data distributions and of retraining the entire model whenever one domain is revised. Model merging avoids both by distilling every domain independently and fusing the resulting task vectors afterwards. We find instead that the benefit polarizes across domain pairs: on those exhibiting negative transfer, every merging operator we evaluate falls below the single-domain reference. We attribute this to cross-domain update coupling, where a substantial fraction of coordinates is updated comparably by both domains and a merge can therefore displace them by as much as their own updates. To overcome this limitation, we propose USA, which converts per-parameter update magnitudes measured during a brief warm-up into per-coordinate perturbation radii, reducing curvature precisely on the coordinates that carry most of the merging displacement. Experiments across mathematics, science and code at two student scales show USA strongest in all six transfer directions, ahead of the single-domain reference by more than four points on average, and reverse the negative transfer of the conflicting pairs.

Problem

• On-policy distillation
• Cross-domain transfer
• Model merging
• Negative transfer
• Language agents

Innovation

• On-policy Distillation
• Model Merging
• Sharpness-Aware Minimization
• Negative Transfer
• Update Coupling