Nash-MOPD improves multi-domain training with adaptive task scheduling
Nash-MOPD: Optimizing Task Trade-offs in Multi-Teacher On-Policy Distillation
Machine LearningArtificial Intelligence
Summary
Training one AI to learn from many expert teachers can be tricky because the AI has limited time and resources to practice each subject. Previous methods either split training time evenly or focus on subjects that show quick improvement, which can leave harder subjects behind. The authors created Nash-MOPD, a smarter way to decide which subject to train on next, balancing attention so slower topics get enough practice. Their approach leads to better overall learning without needing extra teacher data.
What this means in practice
- •For machine learning engineers: Boost training efficiency of multi-domain AI models by dynamically allocating training steps based on progress across tasks.
- •For software developers building ai assistants: Improve a single AI assistant’s performance by better combining expert knowledge from various specialized models into one.
Authors
Siyuan Gan, Bo Yang, Linjian Meng, Boyan Wang, Xiaolong Qi, Jing Huo, Yang Gao
Abstract
Multi-teacher on-policy distillation (MOPD) aggregates the capabilities of multiple domain-expert teachers into a single student model. However, all domains in MOPD share the student's finite budget of optimization steps and rollouts. Fixed uses constant domain sampling ratios based on corpus sizes; by contrast, D$^3$-MOPD adjusts the ratios using the teacher-student reverse-KL divergence and its recent decrease. However, favoring large recent decreases can direct more training toward domains that already improve quickly, while domains that improve more slowly may receive too few examples. To address inflexible allocation and the risk of giving insufficient training to domains that improve more slowly, we propose Nash-MOPD, an adaptive domain scheduler based on a local minimax game using existing reverse-KL statistics. The equilibrium maximizes the smallest predicted proportional reduction under the local model and gives more prompts to domains that need more examples for comparable improvement. Consequently, Nash-MOPD responds to changing learning progress and supports domains that improve more slowly without extra teacher queries or rollouts. At step 100, Nash-MOPD exceeds both baselines' final mathematics and code checkpoint scores on DeepSeek-R1-Distill-Qwen-1.5B and reaches Fixed's on Qwen3-8B-Base. Furthermore, it leads eight of nine reported metrics and is best or tied best on all six mathematics and code accuracy metrics, respectively.