Building Multi-Task Agentic LLMs via Two-Phase 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.

• The paper demonstrates a two-phase distillation method that overcomes mode-covering limitations to recover RL-expert performance.
• The off-policy phase creates a strong initialization, while on-policy refinement steers the model toward optimal task-specific behaviors.
• Empirical results on conversational and game-based tasks validate the method’s efficacy across different scales and model capacities.

Building Multi-Task Agentic LLMs via Two-Phase Distillation

Introduction

The consolidation of reinforcement learning (RL) agents into multi-task agentic LLMs remains a significant open problem in compositional generalization and scalable agent design. The dominant approaches—multi-task RL, parameter-space model merging, and knowledge distillation—each present notable challenges for practical deployment. This work rigorously analyzes the bottlenecks of applying off-policy (supervised) and on-policy (RL-style) distillation for multi-task agentic LLMs, and proposes a two-phase distillation framework—off-policy followed by on-policy—that mitigates mode-covering drawbacks and leverages mode-seeking behavior, achieving RL-expert-level performance across diverse domains.

Limitations of Single-Phase Distillation Approaches

Off-Policy Distillation and Mode-Covering Pathologies

Off-policy distillation, typically implemented as SFT on rollouts from RL experts, is effective in single-task transfer, largely recovering RL performance when student and teacher share the same architecture. However, in the multi-task case, simple aggregation of expert trajectories exposes a capacity-mismatch-induced pathology: the forward KL objective is fundamentally mode-covering. When presented with a union of multiple expert distributions, a capacity-limited student model optimizes to average across modes, exacerbating behavioral interference and reducing per-task performance. Empirical failures are most pronounced in pass^4 metrics on more complex tasks (e.g., telecom)—the student's outputs become inconsistent and less reliable due to this averaging.

Figure 1: Learning curves comparing SFT and off-policy distillation (single-task setting) on airline and telecom, where off-policy distillation is essentially SFT on teacher-generated data.

In single-task settings, off-policy distillation from task-aligned RL teachers even outperforms SFT on rollouts from stronger, but more divergent, upstream models under a limited compute budget, confirming mode-averaging as the dominant factor over exposure bias.

On-Policy Distillation and Initialization Sensitivities

On-policy distillation's reverse KL objective is mode-seeking; learning occurs only in regions actively sampled by the student. While potentially alleviating averaging, direct application is unstable unless the student model is already tuned to output valid task-distributions—otherwise, the gradient signal provides no constructive correction. The RL literature's intuition for policy bootstrap applies: SFT initialization often yields poor exploration for multi-task settings, leading to sub-optimal convergence and underperformance relative to single-task experts.

The Two-Phase Distillation Method

This work proposes an explicit two-phase protocol:

1. Off-Policy Distillation: Initialize the multi-task student with SFT on a union of RL-expert trajectories (forward KL).
2. On-Policy Distillation: Refine this initialization with on-policy RL-style imitation (reverse KL) from the suite of experts, with rewards for agreement at the token level.

The theoretical motivation is substantiated with a toy example: in limited-capacity regimes, off-policy distillation drives the model to “average” solution modes, yielding partial task success ($<$1 reward even if each expert solves its domain), but a subsequent mode-seeking phase steers the solution toward high-probability support aligned with the tasks, recovering full reward.

Figure 2: Training curves for the toy example. For 'Off.+On.', reward is plotted against on-policy distillation steps. Off-policy alone saturates at reward 0.5 and on-policy alone fails at 0, while the two-phase approach achieves full reward 1.

Empirical Results

Evaluation is conducted on two $\tau^2$-bench conversational tasks (airline, telecom) and two GEM text-based games (sudoku, mastermind), with Qwen3-8B and Qwen3-30B as architectures. Metrics include pass@1 and pass^4, measuring per-trial and aggregated task success. Results demonstrate:

• In single-task RL expert distillation, performance is nearly perfectly recovered by off-policy distillation alone.
• In the multi-task setting, off-policy distillation consistently underperforms RL baselines, with larger drops on tasks requiring capacity (e.g., telecom pass^4 drop up to 10.6% for 30B model).
• Two-phase Off.+On. distillation recovers (and sometimes exceeds) RL-expert performance across all tasks, eliminating the off-policy averaging-induced degradation.
• On-policy distillation without a strong initialization (i.e., off-policy pretraining) fails—especially at lower scales—with substantial instability and persistent performance gaps observed.

Additional ablations on data mixture ratios, per-prompt rollout diversity, token filtering, and model merging validate that neither naive dataset rebalancing nor parameter/weight averaging addresses the multi-task averaging pathology, nor does multi-task RL (which can exasperate task interference (Wu et al., 22 Oct 2025, Ramesh et al., 5 Feb 2026)).

Analysis and Discussion

The root cause of degradation in multi-task off-policy distillation is conclusively traced to a model capacity mismatch with the union of behavioral modes present in expert rollouts, enforced by forward KL optimization. As model scaling increases, the number of high-quality modes per expert also grows, maintaining the capacity gap at practical model sizes—a phenomenon observed at both 8B and 30B scale.

Mode-seeking (reverse KL) property in the two-phase approach robustly selects feasible mode subsets consistent with the student’s optimal behaviors, eliminating unwanted averaging and yielding high per-task reliability. Token-level filtering distillation data to match the chosen modes also corroborates the mode-count hypothesis; performance recovers nearly to baseline levels when non-conforming samples are pruned.

Model merging methods (parameter averaging, task arithmetic, DARE) perform inconsistently and are highly sensitive to model and task characteristics, sometimes dominated by the largest task-dependent post-SFT parameter shifts and generally showing substantial degradation outside of a few constrained domains (e.g., math/coding).

Implications and Future Directions

The proposed two-phase approach suggests a robust recipe for scaling multi-task agentic LLMs with RL-derived expertise. Theoretical and empirical findings point to persistent limitations in single-phase off-policy or on-policy distillation as the number of tasks or behavioral complexity grows. This has implications for modular RL-based agentic systems, foundation model consolidation, and efficient deployment on resource-constrained hardware. As model and task scale increases, algorithmic improvements for capacity allocation (dynamic mode selection, adaptive distillation mixture scheduling) are promising directions. Further work is warranted to extend this method to very large model regimes and tasks requiring intricate reasoning and long-horizon planning.

Conclusion

A systematic two-phase distillation approach—off-policy followed by on-policy—effectively overcomes the mode-covering limitations of single-phase multi-task distillation, yielding agentic LLMs with per-task performance matching single-task RL experts. The analysis clarifies that neither naive off-policy nor on-policy distillation alone, nor parameter merging strategies, are sufficient for scalable multi-task learning in agentic LLMs under realistic compute or capacity constraints. These findings inform practical recipes and design choices for multi-expert model consolidation and future research in efficient, scalable agentic architectures.