Adaptive Mutual Distillation for Balanced Multi-Task Post-Training of Large Language Models

Baohang Li; Xiaocheng Feng; Yichong Huang; Chengpeng Fu; Wenshuai Huo; Zekun Zhou; Zekun Yuan; Tingjia Zhang

cs.CLOct 2, 2026

Adaptive Mutual Distillation for Balanced Multi-Task Post-Training of Large Language Models

Authors: Baohang Li, Xiaocheng Feng, Yichong Huang, Chengpeng Fu, Wenshuai Huo, Zekun Zhou, Zekun Yuan, Tingjia Zhang,

Organizations: Harbin Institute of Technology · Peng Cheng Laboratory

Abstract

Multi-task post-training of large language models (LLMs) aims to improve performance across tasks with unequal amounts of training data. Existing methods focus primarily on balancing task contributions during single-model training. Different task-balancing strategies can produce models with complementary strengths, creating opportunities for mutual distillation. However, the usefulness of cross-model supervision can vary across tasks, transfer directions, and stages of training. We propose Adaptive Mutual Distillation (AMD), a collaborative post-training framework that jointly trains two models with different task-balancing strategies. AMD evaluates candidate adjustments to distillation weights through short training probes shared across tasks, then uses task-wise validation scores to select an adjustment for each task and transfer direction. Across six benchmarks and three LLM backbones, both AMD models achieve higher average benchmark scores than supervised fine-tuning (SFT) baselines trained with the same sampling strategies. They also outperform the task-balancing methods evaluated in our experiments. Merging the two trained models can further improve their average benchmark score while yielding a single model for inference. The merged models outperform multi-task SFT by an average of 2.91 points across the three backbones.

Figures & tables

Figure 1: Overview of AMD . Two models with different task-balancing strategies undergo independent SFT warmup, followed by mutual distillation with task- and direction-specific weights adjusted using validation feedback. The trained models can optionally be merged into a single model for inference after collaborative training concludes.

Figure 2: Adaptive distillation-weight update. Three probes globally decrease, retain, or increase the weights. Validation scores determine the weight adjustment for each task and receiving model.

SFT Baselines
Method	Sampling	Math		Medicine		Law		Avg.	$\Delta$ Avg.
Method	Sampling	GSM8K	Math-500	MedQA	MedBullets	LawBench	CMMLU-Law	Avg.	$\Delta$ Avg.
Multi-task SFT	$\tau=1$	57.47	42.40	27.34	40.10	29.53	48.99	40.97	–
Smoothed SFT	$\tau>1$	57.47	34.20	30.71	44.16	28.48	50.03	40.84	–
Prior Methods
PCGrad	$\tau=1$	56.10	34.40	29.38	47.40	32.86	48.64	41.46	+0.49

Table 1: Main results on Qwen3-0.6B-Base. Avg. is the unweighted mean across six benchmarks.

\Delta

Avg. is measured against Smoothed SFT for

\tau>1

rows and Multi-task SFT otherwise, including

+

Merged. Bold indicates column-wise best results.

SFT Baselines
Method	Sampling	Math		Medicine		Law		Avg.	$\Delta$ Avg.
Method	Sampling	GSM8K	Math-500	MedQA	MedBullets	LawBench	CMMLU-Law	Avg.	$\Delta$ Avg.
Multi-task SFT	$\tau=1$	58.45	27.20	49.25	51.14	28.46	50.49	44.16	–
Smoothed SFT	$\tau>1$	50.42	26.00	50.67	55.19	35.84	47.99	44.35	–
Prior Methods
PCGrad	$\tau=1$	62.32	25.80	50.12	55.36	31.93	49.92	45.91	+1.75

Table 2: Main results on Llama-3.1-8B.

SFT Baselines
Method	Sampling	Math		Medicine		Law		Avg.	$\Delta$ Avg.
Method	Sampling	GSM8K	Math-500	MedQA	MedBullets	LawBench	CMMLU-Law	Avg.	$\Delta$ Avg.
Multi-task SFT	$\tau=1$	19.94	7.00	30.40	29.55	16.55	33.53	22.83	–
Smoothed SFT	$\tau>1$	20.39	6.80	28.20	32.79	16.70	32.94	22.97	–
Prior Methods
PCGrad	$\tau=1$	15.62	5.80	27.65	35.88	22.15	32.20	23.22	+0.39

Table 3: Main results on Llama-3.2-1B.

Method	Sampling	Avg.
SFT baselines	$(1,3)$	22.83 / 22.97
Same-temp MD	$(1,1)$	22.50 / 22.45
Same-temp MD	$(3,3)$	23.18 / 22.22
Fixed-Bi	$(1,3)$	24.07 / 23.92
Adaptive	$(1,3)$	25.03 / 24.54

Table 4: Component comparisons on Llama-3.2-1B. Scores are six-benchmark macro-averages. Sampling lists the two temperatures; paired scores follow the same model order. Task-Tied MD selects one action per direction, shared across tasks.

Figure 3: Task- and direction-specific distillation-weight trajectories on Llama-3.2-1B. The first point in each trajectory shows the weight after the initial adaptive update at step 3000.

Method	Sampling	Macro Avg. by seed			Mean $\pm$ Std.
Method	Sampling	42	43	44	Mean $\pm$ Std.
Fixed-Bi	$\tau=1$	24.07	23.79	23.83	23.90 $\pm$ 0.15
Fixed-Bi	$\tau=3$	23.92	24.14	24.44	24.17 $\pm$ 0.26
Fixed-Uni	$\tau=1$	24.38	24.52	24.16	24.35 $\pm$ 0.18
Fixed-Uni	$\tau=3$	24.51	24.40	24.45	24.45 $\pm$ 0.06
Adaptive	$\tau=1$	25.03	24.32	24.44	24.60 $\pm$ 0.38

Table 5: Training-seed variability of the three AMD regimes on Llama-3.2-1B. Scores are unweighted six-benchmark macro-averages, summarized across three runs as mean

\pm

sample standard deviation. The

+

Merged row reports the merged Adaptive model from each run.

Training	Empirical sampling	Smoothed sampling	Merged	$\Delta$ Merge
SFT baselines	22.83	22.97	24.08	+1.11
AMD	25.03	24.54	25.53	+0.50

Table 6: Model merging on Llama-3.2-1B. Scores are six-benchmark macro-averages;

\Delta

Merge is the gain over the better model before merging.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Category	Dataset	Cap
Math	NuminaMath-1.5 ( AI-MO, 2025 )	100k
Medicine	MedQA ( Jin et al., 2021 )	9,178
Law	DISC-Law ( Yue et al., 2023 )	50k
General	InfinityInstruct ( Li et al., 2025a )	10k

Appendix

Table 7: Training mixture for the main experiments. Each row lists one task category, its source dataset, and the configured sample cap.

SFT Baselines
Method	Sampling	Math		Medicine		Law		Avg.	$\Delta$ Avg.
Method	Sampling	GSM8K	Math-500	MedQA	MedBullets	LawBench	CMMLU-Law	Avg.	$\Delta$ Avg.
Smoothed SFT	$\tau>1$	20.39	6.80	28.20	32.79	16.70	32.94	22.97	–
Prior Methods
PCGrad	$\tau>1$	21.30	7.00	26.47	34.90	20.78	33.14	23.93	+0.96
MFTCoder	$\tau>1$	15.16	4.80	28.36	33.93	21.51	31.13	22.48	-0.49

Appendix

Table 8: Sampling-matched comparison on Llama-3.2-1B under temperature-smoothed sampling. Each reported model or branch uses the same

\tau>1

sampling distribution and six-benchmark evaluation protocol. All rows report raw scores;

\Delta

Avg. gives the macro-average gain relative to the Smoothed SFT reference. Bold marks the highest value in each column.

Panel of Figure 4 — Figure 4: Sensitivity to fixed distillation weights on Llama-3.2-1B. Scores are six-benchmark macro-averages. Dashed and solid horizontal lines indicate the SFT baseline and Adaptive, respectively, for each sampling branch.

#	Schedule	Macro Avg.
1	$\displaystyle\mu=1$	23.96
2	$\displaystyle\mu=\sqrt{\frac{T_{\max}-0.8\,t}{T_{\max}}}$	24.89
3	$\displaystyle\mu=\sqrt{\max\left(0,\frac{T_{\max}-1.2\,t}{T_{\max}}\right)}$	23.64
4	$\displaystyle\mu=\sqrt{\frac{T_{\max}-t}{T_{\max}}}$	25.03

Appendix

Table 9: Sensitivity of the Llama-3.2-1B empirical-sampling model to the step-size schedule

\mu

. Scores are six-benchmark macro-averages; bold marks the best score.

Method	Macro Avg.
Single-model baselines, six epochs
Multi-task SFT	23.05
Smoothed SFT	22.94
PCGrad	23.50
FAMO	24.16
CoBa	23.35

Appendix

Table 10: Longer-budget stress test on Llama-3.2-1B. Single-model baselines train for six epochs, twice the three-epoch main protocol used by AMD. For each AMD regime, we report the higher-scoring model from the collaborative run. Scores are six-benchmark macro-averages.

LLM Post-Training

Large language model post-training adapts pretrained models into more useful, reliable, and specialized assistants through supervised and preference-based learning.

See all 246 papers

Thinking at the Right Size: Amortized Distillation Across Post-Trained LLMs
Aug 24, 2026Yan Zhou, Sara Kangaslahti, Jonathan Geuter +4Language Model Distillation Knowledge Distillation
Harvard University · Kempner Institute · IST Austria
Post-Hoc Merging is Not Enough: Many-Shot Model Merging with Loss-Gap Balancing
Jun 15, 2026Kyungjin Im, Miru Kim, Chanin Eom +1Multi-Task Learning Language Model Merging
Soongsil University, Republic of Korea · Department of Electrical and Computer Engineering, Sungkyunkwan University, Republic of Korea
Decomposing the Basic Abilities of Large Language Models: Mitigating Cross-Task Interference in Multi-Task Instruct-Tuning
May 7, 2026Bing Wang, Ximing Li, Changchun Li +3LLM Fine-Tuning Gradient Interference
College of Computer Science and Technology, Jilin University · Key Laboratory of Symbolic Computation and Knowledge Engineering, Ministry of Education, Jilin University · RIKEN Center for Advanced Intelligence Project +1

Aug 24, 2026J/K move · Enter open · S save

Thinking at the Right Size: Amortized Distillation Across Post-Trained LLMs

Yan Zhou, Sara Kangaslahti, Jonathan Geuter, Nihal V. Nayak+3

Harvard University · Kempner Institute · IST Austria

Language Model Distillation Knowledge Distillation Language Model Post-Training

Open paper AskPDF arXiv

Practical deployment of large language models (LLMs) requires families of post-trained variants---instruction-tuned, reasoning-tuned, and chat-style models---each at multiple sizes to meet diverse latency and memory budgets. Producing each (variant, size) pair independently is prohibitive, so model families typically span only a handful of coarse-grained sizes per post-trained variant. Boomerang distillation (Kangaslahti et al., 2026) reduces this cost along the size axis for base models. Through model size interpolation, it constructs models of intermediate sizes from a single teacher-student pair without additional training. However, it still treats each post-trained variant as a separate object of optimization. We introduce ADAPT---Amortized Distillation Across Post-Trained LLMs---a framework for amortizing distillation across both axes of a model family: size and post-training variant, producing

L \times K

models for

L

interpolated sizes across

K

post-trained variants with a single distillation run. ADAPT combines two components. First, a two-phase distillation procedure constructs post-trained students through pre-training alignment and supervised fine-tuning distillation, enabling smooth size--performance interpolation on generation and reasoning tasks. Second, weight-delta initialization approximates this construction across post-trained variants by transferring the distillation-induced weight change from the base model to students initialized from different post-trained variants. The resulting continuum of interpolated models also enables adaptive model-size selection at inference time, improving the compute--accuracy trade-off for long-form reasoning tasks.

Abstract

LLM Post-Training

Explore similar work

Thinking at the Right Size: Amortized Distillation Across Post-Trained LLMs

Post-Hoc Merging is Not Enough: Many-Shot Model Merging with Loss-Gap Balancing

Decomposing the Basic Abilities of Large Language Models: Mitigating Cross-Task Interference in Multi-Task Instruct-Tuning