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

Aug 24, 2026

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

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.

Yan Zhou, Sara Kangaslahti, Jonathan Geuter +4

Harvard University · Kempner Institute · IST Austria

Language Model Distillation Knowledge Distillation Language Model Post-Training

Ask PDF

Jun 15, 2026

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

Model merging has become a practical post-training strategy for building a single multi-task large language model (LLM) by combining multiple task-specialized models. However, most existing approaches rely on post-hoc merging, in which task-specific models are merged only once after training. This one-shot aggregation often suffers from task interference, leading to information erasure across individual tasks. In this work, we show that replacing post-hoc merging with an iterative many-shot merging protocol is effective in improving multi-task performance. Building on this insight, we propose METIS, Mitigating Erasure from Task Interference for Stable many-shot merging. METIS is a loss-aware many-shot merging method that addresses information erasure in post-hoc merging through task-wise loss-gap weighting and consensus-based masking. Notably, METIS exhibits significant performance improvement on the worst-performing task, effectively mitigating information erasure. (Project page: https://imkyungjin.github.io/METIS/)

Kyungjin Im, Miru Kim, Chanin Eom +1

Soongsil University, Republic of Korea · Department of Electrical and Computer Engineering, Sungkyunkwan University, Republic of Korea

Multi-Task Learning Language Model Merging Model Merging Task Interference

Ask PDF

May 7, 2026

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

Recently, the prominent performance of large language models (LLMs) has been largely driven by multi-task instruct-tuning. Unfortunately, this training paradigm suffers from a key issue, named cross-task interference, due to conflicting gradients over shared parameters among different tasks. Some previous methods mitigate this issue by isolating task-specific parameters, e.g., task-specific neuron selection and mixture-of-experts. In this paper, we empirically reveal that the cross-task interference still exists for the existing solutions because of many parameters also shared by different tasks, and accordingly, we propose a novel solution, namely Basic Abilities Decomposition for multi-task Instruct-Tuning (BADIT). Specifically, we empirically find that certain parameters are consistently co-activated, and that co-activated parameters naturally organize into base groups. This motivates us to analogize that LLMs encode several orthogonal basic abilities, and that any task can be represented as a linear combination of these abilities. Accordingly, we propose BADIT that decomposes LLM parameters into orthogonal high-singular-value LoRA experts representing basic abilities, and dynamically enforces their orthogonality during training via spherical clustering of rank-1 components. We conduct extensive experiments on the SuperNI benchmark with 6 LLMs, and empirical results demonstrate that BADIT can outperform SOTA methods and mitigate the degree of cross-task interference.

Bing Wang, Ximing Li, Changchun Li +3

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

LLM Fine-Tuning Gradient Interference Instruction Following Instruction Tuning Parameter-Efficient Fine-Tuning

Ask PDF

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