PaperlayerView exact revision on arXiv

Source-linked AI summary

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, Bing Qin

arXiv:2610.02856v1cs.CL

TL;DR

Unequal task data make multi-task performance a balancing challenge, while the value of cross-model supervision can vary by task, direction, and training stage. AMD jointly trains differently balanced models with validation-guided mutual distillation; across benchmarks and backbones, both models beat matched SFT baselines, and merging yields further gains.

  • Problem

    Multi-task training must manage trade-offs from unequal task data, while the usefulness of cross-model supervision can differ by task, transfer direction, and training stage.

    Source: §1 · p. 1

  • Method

    AMD jointly trains two differently balanced models and uses short shared training probes with task-wise validation scores to select distillation-weight adjustments for each task and direction.

    Source: §1 · p. 2

  • Results

    2.91 points: merged AMD models exceed multi-task SFT on average across three backbones, while both AMD models outperform sampling-matched SFT baselines across the evaluated benchmarks.

    Source: §1 · p. 2, §4.2 · p. 7, §6 · p. 9

  • Takeaways & Limitations

    Merging the collaboratively trained models can further improve average benchmark scores while producing a single model for inference.

    Source: §4.2 · p. 7, §6 · p. 9

    General instruction-following data are included in SFT but do not participate in mutual distillation.

    Source: §3.1 · p. 3

Abstract

from arXiv · show

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.

1 INTRODUCTION

Unequal task data create trade-offs: emphasizing underrepresented tasks can improve their performance while reducing performance on others. AMD has differently balanced models collaborate through adaptive, task- and direction-specific distillation; both models and their merged version outperform evaluated baselines.

1 INTRODUCTION

  • Unequal task quantities make multi-task performance a balancing problem: emphasizing underrepresented tasks may improve them but reduce performance on other tasks.The stated goal is to manage these trade-offs to improve overall multi-task performance.Source: §1 · p. 1
  • Existing approaches primarily adjust task contributions within a single model, while different balancing strategies can produce complementary strengths that invite collaboration.The value of mutual-distillation supervision can vary by task, transfer direction, and training stage.Source: §1 · p. 1
  • AMD jointly trains two differently balanced models: after independent SFT warmup, they continue under their respective sampling distributions while exchanging supervision.Task-wise validation feedback guides distillation weights separately for each task and transfer direction.Source: §1 · p. 2, Figure 1 · p. 2
  • Three short probes test decreasing, maintaining, or increasing all distillation weights; validation scores select an adjustment for each task and transfer direction.The probes are shared across tasks.Source: §1 · p. 2
  • Merged AMD models achieve the highest average benchmark scores among evaluated methods, exceeding multi-task SFT by 2.70–3.24 points across three backbones and six benchmarks.Both AMD models also outperform SFT baselines trained with the same sampling strategies and the evaluated task-balancing methods.Source: §1 · p. 2

2 RELATED WORK

Prior work balances task contributions through data mixtures, losses, gradients, and distillation. AMD combines differently balanced multi-task models and adapts mutual-distillation weights using task-wise validation feedback.

2 RELATED WORK

  • Task-balancing methods adjust data allocation, task losses, or gradients; examples include GradNorm, PCGrad, FAMO, CoBa, MFTCoder, MaZO, DoReMi, Mixture-of-Skills, DynamixSFT, and HBO.AMD instead brings models with different task-balancing strategies into collaborative training while each retains its own strategy.Source: §2 · p. 3
  • Mutual-distillation research includes peer-supervision aggregation, online multi-task distillation, and methods that adapt teacher contributions or task and distillation losses.The passage names OKDDip, online distillation, MMKD, and DTO-KD as examples.Source: §2 · p. 3
  • AMD distinguishes itself by jointly training two differently balanced multi-task models and adjusting distillation weights per task and transfer direction using validation scores from shared probes.The update probes test globally decreasing, retaining, or increasing weights before task-wise selection.Source: §2 · p. 3, Figure 2 · p. 3

3 METHODOLOGY

AMD uses an independent SFT warmup followed by collaborative training with adaptive mutual distillation. Shared probes test global weight adjustments, while task-wise validation scores select updates for each task and direction.

3 METHODOLOGY

  • AMD has two stages: independent SFT warmup, then collaborative training with adaptive mutual distillation; the trained models may optionally be merged for inference.Validation feedback guides task- and direction-specific distillation weights in the second stage.Source: §3 · p. 3

3.1 FRAMEWORK

  • Both models share data, architecture, and initialization; general instruction-following data enter SFT but are excluded from mutual distillation.The participating tasks are indexed by k.Source: §3.1 · p. 3
  • During warmup, Multi-task SFT samples by task frequency, while Smoothed SFT raises frequencies to the power 1/τs and renormalizes them to favor less frequent tasks.For each task and direction, the receiving model’s objective combines SFT loss with weighted KL-divergence distillation from its partner.Source: §3.1 · pp. 3–4
  • For comparison, FIXED-BI applies one fixed weight across participating tasks and directions; FIXED-UNI chooses the higher-validation-score model as teacher per task at each evaluation.FIXED-UNI keeps that teacher assignment until the next evaluation.Source: §3.1 · p. 4

3.2 ADAPTIVE-WEIGHT DISTILLATION

  • At each update, AMD tests globally decreasing, retaining, or increasing the task- and direction-specific weights in short probes started from the same current model pair; probe parameter changes are discarded.Increases and decreases occur in logit space, with a sigmoid mapping keeping bounded weights valid.Source: §3.2 · p. 4
  • AMD selects the highest-validation-score action separately for each task and transfer direction, combines the selected weights, and resumes training with the updated vector.Figure 2 illustrates the update for one transfer direction.Source: §3.2 · pp. 4–5

4 EXPERIMENTS

The experiments compare AMD with sampling-matched SFT, prior task-balancing methods, and fixed-distillation variants across three LLM backbones. Both Adaptive branches improve on the corresponding SFT baselines and evaluated task-balancing methods, while merging the branches yields the highest macro-average scores.

4.1 SETUP

  • The evaluation covers three LLM backbones, a four-domain training mixture, and six benchmarks, with comparisons against matched SFT, prior task-balancing methods, and fixed-strength distillation.Overall performance is the unweighted mean of six benchmark scores; the models share training examples, tokenizer, generation settings, and evaluation protocol within each backbone.Source: §4.1 · pp. 5–7, §4 · p. 5

4.2 MAIN RESULTS

  • Both Adaptive branches exceed their sampling-matched SFT baselines and all evaluated prior task-balancing methods across the three backbones.Gains over corresponding SFT baselines range from 2.20 to 2.50 points for empirical sampling and from 1.57 to 2.54 points for smoothed sampling.Source: §4.2 · p. 7
  • Adaptive yields the highest-scoring individual model among distillation variants on each backbone, with its best branch exceeding the best fixed-distillation branch by 0.26, 0.45, and 0.52 points.On Qwen3-0.6B-Base, Adaptive and Fixed-Uni tie under smoothed sampling.Source: §4.2 · p. 7
  • Merging the two Adaptive models improves over the better branch by 0.30, 0.51, and 0.50 points across Qwen3-0.6B-Base, Llama-3.1-8B, and Llama-3.2-1B; gains over multi-task SFT average 2.91 points.The merged models reach macro-average scores of 43.75, 47.40, and 25.53, respectively.Source: §4.2 · p. 7

5 ANALYSIS

The analysis tests the effects of sampling distributions, adaptive weight updates, training seeds, and model merging on Llama-3.2-1B. It finds that adaptive collaboration and subsequent merging improve benchmark scores, while the useful distillation weights vary by task and transfer direction.

5.1 EFFECT OF SAMPLING DISTRIBUTIONS

  • With different sampling distributions, FIXED-BI scores 24.07 and 23.92, outperforming both corresponding SFT baselines; its highest branch exceeds the best same-temperature result by 0.89 points.Under empirical sampling, same-temperature MD scores 22.50 and 22.45 versus an SFT baseline of 22.83; under smoothed sampling, it scores 23.18 and 22.22 versus 22.97.Source: §5.1 · pp. 7–8

5.2 EFFECT OF ADAPTIVE WEIGHT UPDATES

  • Online adaptation scores 25.03 and 24.54, above fixed final weights (23.41 and 22.01); separate direction and task updates also outperform their tied variants in both branches.Direction-Tied MD scores 24.59 and 23.38, while Task-Tied MD scores 24.33 and 23.69 under empirical and smoothed sampling, respectively.Source: §5.2 · p. 8
  • A globally shared update reaches 23.41 in its best branch versus 23.18 for fixed-weight Same-temp MD, improving one branch but slightly lowering the other; task-specific weight trajectories also differ across Math, Medicine, and Law.In Math both weights decrease overall, in Medicine their relative ordering changes, and in Law they decrease initially before rising toward training’s end.Source: §5.2 · p. 8

5.3 VARIABILITY ACROSS TRAINING SEEDS

  • Across three seeds, Adaptive averages 24.60 ± 0.38 under empirical sampling and 24.59 ± 0.10 under smoothed sampling; both branches beat their Fixed-Bi and Fixed-Uni baselines on average.The smoothed-sampling branch also exceeds both corresponding baselines in every reported run.Source: §5.3 · pp. 8–9

5.4 POST-TRAINING MODEL MERGING

  • Merging raises AMD’s score from 25.03 for its better branch to 25.53, exceeding the merged SFT baseline by 1.45 points; across three seeds, merged AMD scores 25.47 ± 0.17 and beats both constituents in every run.Parameter averaging combines the AMD models into one model for inference.Source: §5.4 · p. 9

6 CONCLUSION

AMD uses validation feedback to adapt distillation between models with different task-sampling distributions. Across three backbones and six benchmarks, both trained models outperform sampling-matched SFT, while merging further improves macro-average scores.

6 CONCLUSION

  • Merged AMD models achieve the highest macro-average scores among compared methods, exceeding multi-task SFT by 2.70–3.24 points.Source: §6 · p. 9
  • Both trained AMD models outperform their sampling-matched SFT baselines in macro-average performance across three backbones and six benchmarks.Source: §6 · p. 9
  • AMD enables models with different sampling distributions to learn from each other, with validation feedback guiding distillation strength by task and transfer direction throughout training.Source: §6 · p. 9

AI USE STATEMENT

The authors state that generative AI tools were used only to polish the manuscript's language.

AI USE STATEMENT

  • Generative AI tools were used only for language polishing, not translation or result interpretation.Source: AI USE STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

The reproducibility statement points readers to the framework, experimental settings, and seed-variability reporting; the appendix contains further experimental details and analyses.

REPRODUCIBILITY STATEMENT

APPENDIX

  • The appendix adds analyses of temperature-smoothed baselines, fixed distillation weights, weight-update schedules, and training budgets.Source: APPENDIX · p. 14

A.1 DATA MIXTURE AND EVALUATION

The appendix describes task-category data mixtures and evaluation, training configurations, and the rationale for shared probes with task-wise selection.

A.1 DATA MIXTURE AND EVALUATION

  • Table 7 summarizes each task category's source dataset and configured sample cap, and categories are consistent across sampling, baselines, and AMD policies.Source: §A.1 · p. 14, Table 7 · p. 14
  • The aggregate score is the macro-average over six benchmarks, with two benchmarks each for mathematical, medical, and legal reasoning; training, validation, and test data are disjoint.Source: §A.1 · p. 14
  • Adaptive validation uses fixed subsets of 100 mathematics, 100 medicine, and 200 law examples, shared across backbones; the checkpoint is selected by mean validation accuracy across both branches and three tasks.Source: §A.1 · p. 14

A.2 TRAINING CONFIGURATION

  • Main AMD runs pair empirical-mixture sampling at τ = 1 with smoothed sampling at τ = 3, begin distillation at step 3000, and use temperature T = 1.0.Source: §A.2 · p. 14
  • Each candidate probe uses the same next-batch sequence and restored model, optimizer, scheduler, iterator, and random-number-generator states; tied validation scores leave AMD weights unchanged.Source: §A.2 · p. 15
  • Comparisons keep effective optimizer-update scale comparable through differing batch structures, but are not exactly wall-clock-cost matched.Source: §A.2 · p. 15

A.3 RATIONALE FOR SHARED PROBING AND TASK-WISE SELECTION

  • Shared probes test adjustments to the full weight vector, while task-wise validation rankings select separate adjustments by task and transfer direction rather than imposing one aggregate action.Source: §A.3 · p. 15
  • The selected adjustments are recombined into a new vector and training resumes from the original state; this vector may differ from every probe, whose scores guide rather than determine its performance.Source: §A.3 · p. 15

B TEMPERATURE-SMOOTHED BASELINE SETTING

Under sampling-matched temperature smoothing, Adaptive achieves the highest six-benchmark macro-average in the table, exceeding both Smoothed SFT and the strongest prior baseline.

C SENSITIVITY TO FIXED DISTILLATION WEIGHTS

Adaptive outperforms the best fixed-weight result more clearly under empirical sampling than smoothed sampling; the standard square-root schedule scores highest among tested schedules, and Adaptive retains the highest score in a longer-budget comparison.

C SENSITIVITY TO FIXED DISTILLATION WEIGHTS

  • 25.03 versus 24.38 gives Adaptive a 0.65-point gain over the best fixed weight for τ = 1; for τ = 3, scores are 24.54 versus 24.51.Fixed-Uni peaks at α = 0.2 in both branches, and the adaptation gain is more pronounced under empirical sampling.Source: Figure 4 · p. 16, C SENSITIVITY TO FIXED DISTILLATION… · p. 16

D WEIGHT-UPDATE SCHEDULE SENSITIVITY

  • 25.03 is the standard square-root schedule's score, ahead of slower decay at 24.89, constant step at 23.96, and faster decay at 23.64.Both standard and slower schedules outperform constant and faster-decaying alternatives, supporting gradual reduction in update magnitude.Source: Table 9 · p. 17, D WEIGHT-UPDATE SCHEDULE SENSITIVITY · p. 17

E LONGER-BUDGET STRESS TEST

  • 25.03 is Adaptive's score, above six-epoch FAMO at 24.16, FIXED-BI at 24.07, and FIXED-UNI at 24.51.Single-model baselines train for six epochs, twice the three-epoch AMD protocol; their epoch budgets correspond to different computational costs because AMD also uses cross-model passes, probes, and validation.Source: Table 10 · p. 17, E LONGER-BUDGET STRESS TEST · p. 17

Related papers

Loading 2610.02856v1…