Building Multi-Task Agentic LLMs via Two-Phase Distillation

2026-06-29

AI Summary

This work addresses the performance degradation commonly observed in multitask large language models when trained via mixed-task learning or single-stage distillation, often due to conflicting behavioral patterns across tasks. To mitigate this issue, the authors propose a two-stage distillation framework: first, individual single-task reinforcement learning experts are trained; then, their knowledge is transferred through a combination of off-policy distillation—ensuring broad behavioral coverage—and on-policy fine-tuning, which hones in on task-specific strategies. This approach effectively alleviates behavioral interference among tasks and successfully recovers the performance of each expert in both conversational agent and text-based game benchmarks, significantly outperforming baselines that rely solely on either off-policy or on-policy distillation.

Abstract

A key step toward artificial general intelligence is to train models that can perform multiple tasks. In this paper, we study how to build such models by first training separate RL experts for individual tasks and then consolidating them via distillation, as an alternative to directly training a single model on mixed tasks. We show that off-policy distillation degrades in multi-task settings due to the mode-covering nature of forward KL: aggregating data from multiple tasks introduces a large number of behavioral modes that can exceed the student's capacity, forcing it to average across behaviors and leading to degraded performance. In contrast, on-policy distillation is mode-seeking but requires strong initialization. Inspired by these observations, we propose a two-phase approach: off-policy distillation followed by on-policy refinement. Evaluation across conversational agents and text-based games confirms that this two-phase approach matches single-task RL expert performance for each individual task, whereas off-policy or on-policy distillation alone fails to match this performance.

Problem

• multi-task learning
• distillation
• large language models
• reinforcement learning
• artificial general intelligence

Innovation

• two-phase distillation
• multi-task agentic LLMs
• off-policy distillation
• on-policy distillation
• forward KL