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

arXiv: https://arxiv.org/abs/2609.34225 · HTML: https://mlsanity.com/paper/2609.34225/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Reinforcement learning for LLM agents, one of 16 directions in LLM agents and tool use.
- Builds on 65 of its 69 references that are in the map.
- First posted on arXiv 28 Sep 2026.

## Placed in

- LLM agents and tool use / [Reinforcement learning for LLM agents](https://mlsanity.com/agents/reinforcement-learning-for-llm-agents/)
- Knowledge distillation / [On-policy and self-distillation for LLMs](https://mlsanity.com/distill/on-policy-and-self-distillation-for-llms/)

## Builds on (most cited first, 30 of 65)

- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [ReAct: Synergizing Reasoning and Acting in Language Models](https://mlsanity.com/paper/2210.03629/) (2022)
- [Toolformer: Language Models Can Teach Themselves to Use Tools](https://mlsanity.com/paper/2302.04761/) (2023)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [GPQA: A Graduate-Level Google-Proof Q&A Benchmark](https://mlsanity.com/paper/2311.12022/) (2023)
- [DAPO: An Open-Source LLM Reinforcement Learning System at Scale](https://mlsanity.com/paper/2503.14476/) (2025)
- [HybridFlow: A Flexible and Efficient RLHF Framework](https://mlsanity.com/paper/2409.19256/) (2024)
- [LiveCodeBench: Holistic and Contamination Free Evaluation of Large Language Models for Code](https://mlsanity.com/paper/2403.07974/) (2024)
- [Search-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement Learning](https://mlsanity.com/paper/2503.09516/) (2025)
- [The Rise and Potential of Large Language Model Based Agents: A Survey](https://mlsanity.com/paper/2309.07864/) (2023)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time](https://mlsanity.com/paper/2203.05482/) (2022)
- [Text Embeddings by Weakly-Supervised Contrastive Pre-training](https://mlsanity.com/paper/2212.03533/) (2022)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [When Not to Trust Language Models: Investigating Effectiveness of Parametric and Non-Parametric Memories](https://mlsanity.com/paper/2212.10511/) (2022)
- [Editing Models with Task Arithmetic](https://mlsanity.com/paper/2212.04089/) (2022)
- [On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes](https://mlsanity.com/paper/2306.13649/) (2023)
- [TIES-Merging: Resolving Interference When Merging Models](https://mlsanity.com/paper/2306.01708/) (2023)
- [Language Models are Super Mario: Absorbing Abilities from Homologous Models as a Free Lunch](https://mlsanity.com/paper/2311.03099/) (2023)
- [Linear Mode Connectivity and the Lottery Ticket Hypothesis](https://mlsanity.com/paper/1912.05671/) (2019)
- [Loss Surfaces, Mode Connectivity, and Fast Ensembling of DNNs](https://mlsanity.com/paper/1802.10026/) (2018)
- [FinQA: A Dataset of Numerical Reasoning over Financial Data](https://mlsanity.com/paper/2109.00122/) (2021)
- [Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models](https://mlsanity.com/paper/2601.18734/) (2026)
- [ReTool: Reinforcement Learning for Strategic Tool Use in LLMs](https://mlsanity.com/paper/2504.11536/) (2025)
- [Reinforcement Learning via Self-Distillation](https://mlsanity.com/paper/2601.20802/) (2026)
- [Rethinking On-Policy Distillation of Large Language Models: Phenomenology, Mechanism, and Recipe](https://mlsanity.com/paper/2604.13016/) (2026)
- [SciBench: Evaluating College-Level Scientific Problem-Solving Abilities of Large Language Models](https://mlsanity.com/paper/2307.10635/) (2023)
- [Model Merging in LLMs, MLLMs, and Beyond: Methods, Theories, Applications and Opportunities](https://mlsanity.com/paper/2408.07666/) (2024)
- [AdaMerging: Adaptive Model Merging for Multi-Task Learning](https://mlsanity.com/paper/2310.02575/) (2023)

## Abstract (arXiv)

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.