PaperlayerView exact revision on arXiv

Source-linked AI summary

mSFT: Addressing Dataset Mixtures Overfitting Heterogeneously in Multi-task SFT

Woosung Koh, Jeyoung Jeon, Youngjin Song, Yujin Cheon, Soowon Oh, Jaehyeong Choi, Se-Young Yun

arXiv:2603.21606v5cs.LGcs.AI

TL;DR

Multi-task SFT assigns the same compute to tasks with different learning dynamics, allowing faster-learning sub-datasets to overfit while slower ones remain under-fitted. mSFT iteratively excludes the earliest-overfitting sub-dataset and rolls back to its optimal checkpoint; across 10 benchmarks and 6 base models, it outperforms 4 baselines and remains robust across dataset sizes and task granularities.

  • Problem

    Multi-task SFT forces sub-datasets with heterogeneous learning dynamics and optimal stopping points to use one global compute budget.

  • Method

    mSFT iteratively records each active sub-dataset’s optimal compute, excludes the earliest-overfitting one, and continues from the corresponding rollback checkpoint.

  • Results

    Across 10 benchmarks and 6 base models, mSFT outperforms 4 baselines and 2 ablative baselines on average.

  • Takeaways & Limitations

    mSFT’s gains remain robust across dataset sizes, task counts, and task granularities, while low compute budgets can improve performance and reduce training FLOPs.

  • Takeaways & Limitations

    The method incurs additional storage overhead, with dynamic pruning yielding an average storage footprint of approximately 4.44× SFT.

Abstract

from arXiv · show

Current language model training commonly applies multi-task Supervised Fine-Tuning (SFT) using a homogeneous compute budget across all sub-datasets. This approach is fundamentally sub-optimal: heterogeneous learning dynamics cause faster-learning tasks to overfit early while slower ones remain under-fitted. To address this, we introduce mSFT, an iterative, overfitting-aware search algorithm for multi-task data mixtures. mSFT trains the model on an active mixture, identifies and excludes the earliest overfitting sub-dataset, and reverts to that specific optimal checkpoint before continuing. Extensive evaluations demonstrate that mSFT consistently outperforms 4 baselines across 10 benchmarks and 6 base models. Further analysis confirms mSFT maintains robust gains across diverse dataset sizes, task granularities, and is insensitive to its single new hyperparameter (compute budget). Notably, at low compute budget, mSFT can improve performance while lowering training FLOPs. Ultimately, mSFT establishes a practical overfitting-aware algorithm for multi-task SFT that maximizes the potential of models across diverse data mixtures.

1 Introduction

Multi-task SFT commonly assigns the same compute to every sub-dataset, despite heterogeneous learning dynamics that make this approach sub-optimal. The paper proposes MSFT, which searches for overfitting points and reports robust gains across models, benchmarks, dataset sizes, and task granularities.

  • Nemotron’s code sub-dataset required less compute than others, yet its allocation remained coarse under multi-stage homogeneous training.
  • Homogeneous SFT trains all sub-datasets with the same compute, a status quo used by frontier open-weight models.
  • MSFT iteratively excludes the earliest-overfitting sub-dataset by rolling back to its overfitting checkpoint before continuing training.
  • MSFT’s average performance across 10 benchmarks outperforms 4 baselines across 6 base models.
  • MSFT gains are not concentrated in a few outlier tasks, as standard deviation decreases across benchmarks.
  • MSFT remains robust across dataset sizes of 9K, 18K, and 27K, task counts of 5, 10, and 15, and diverse task granularities.With low compute budget C, it can improve performance while reducing FLOPs.

2 Motivation: Dataset Mixtures Overfit Heterogeneously

Multi-task mixtures contain tasks with different optimal compute levels, but standard SFT forces every task to use one global budget. This causes rapidly converging tasks to overfit while slower-learning tasks remain under-fitted.

  • Sub-datasets exhibit substantially different overfitting dynamics, an observation demonstrated on Qwen3 8B and reported as consistent across other models.
  • For each task, optimal compute is defined as the stopping point where held-out test performance reaches maximum generalization.
  • Standard homogeneous training forces every task i to use the same global compute c_global.
  • Distinct task distributions and complexities produce widely varying convergence rates and optimal compute levels.
  • When c_global exceeds a rapidly converging task’s optimum, that task begins to overfit; when it falls below a slower task’s optimum, that task remains under-fitted.

3 MSFT: Heterogeneous Early-stopping for Multi-task Data Mixtures

The section shows that excluding overfitting sub-datasets changes the optimal compute of remaining tasks, making a single search rollout inaccurate. MSFT addresses this by iteratively rolling out on the active mixture, excluding the earliest-overfitting sub-dataset, and rolling back to its overfitting checkpoint.

  • 3.1 Limitation of a Naïve Solution: SRO SFT searches per-dataset optima once, then retrains from scratch while excluding each sub-dataset at its identified peak.For example, AQUA-RAT is excluded at epoch 1.25 and SciQ at epoch 2.75.
  • 3.1 Limitation of a Naïve Solution: SRO SFT becomes increasingly inaccurate because its search-stage optima reflect gradient interactions from the complete mixture, not later reduced mixtures.As exclusions accumulate, the active gradient sum diverges further from the original search dynamics, especially for late-stage tasks.
  • 3.1 Limitation of a Naïve Solution: Excluding one-tenth of the mixture significantly shifts the optimal stopping points of the remaining tasks.The optimal compute divergence averages an absolute shift of 0.91 epochs across model families and scales.
  • 3.2 Iterative Overfitting-Aware Search: MSFT aligns search and training by iteratively rolling out on active sub-datasets and rolling back after the earliest overfitting event.It records each active sub-dataset’s optimal compute, removes the earliest-overfitting dataset, and reverts the model to that dataset’s overfitting checkpoint.
  • 3.2 Iterative Overfitting-Aware Search: MSFT continues until no active sub-datasets remain, using compute budget C to obtain granular overfitting measurements.If no sub-dataset overfits within C, the algorithm updates the model and continues without rolling back.

4 Empirical Study

MSFT is evaluated against standard and adaptive SFT baselines across six models, multiple task domains, and robustness settings. It consistently improves performance, stability, scale robustness, granularity robustness, and compute efficiency, while outperforming naïve heterogeneous early-stopping alternatives.

  • Overall Performance and Robustness: MSFT consistently achieves the highest average accuracy across six evaluated models and improves all three major task domains.Reported gains are +0.7% in Science & Knowledge, +2.4% in Commonsense & Language, and +3.0% in Mathematical & Quantitative reasoning.
  • Consistency and Outlier Analysis: MSFT generally maintains the lowest benchmark standard deviation and achieves first place 26 times, compared with 13 times for IES.The authors interpret the lower standard deviation as evidence that gains are not driven by large outliers.
  • Ablation Study: MSFT outperforms both SRO SFT and Soft SRO SFT in average performance, showing that single-rollout heterogeneous early stopping is sub-optimal.The ablations compare MSFT with two naïve heterogeneous early-stopping alternatives averaged across six underlying models.
  • Further Analysis: MSFT remains robust across dataset sizes and task counts, yielding an average improvement of +5.4% over SFT across three configurations.The tested task counts are N ∈{5, 10, 15}, and the dataset-size analysis includes 9K, 18K, and 27K examples.
  • Further Analysis: At C = 1, MSFT improves performance by +3.4% while reducing average compute by 120.3 PFLOPs relative to SFT.The savings arise because MSFT adds no rollout overhead and dynamically excludes sub-datasets; gains do not degrade as C decreases.
  • Further Analysis: On 21 MedMCQA sub-categories, MSFT improves accuracy by +1.86% over SFT, compared with +0.29% for IES.The analysis demonstrates effectiveness at a more granular task decomposition.

5 Discussion

The discussion situates MSFT within prior work on dataset selection, mixture ratios, and task-imbalance mitigation, while identifying checkpoint storage as an operational trade-off.

  • Prior work studies SFT dataset inclusion, mixture ratios, and task imbalance through loss reweighting or gradient manipulation.
  • MSFT’s roll-out phase requires saving intermediate checkpoints, creating additional storage overhead.
  • 4.44× SFT is the average storage footprint after dynamic checkpoint pruning.
  • The authors consider this storage cost acceptable because disk space is rarely the primary bottleneck relative to compute.

A Computation of FLOPs Proportion

This appendix computes training FLOPs by stage using reported model, token, sequence-length, pair, and episode information, then summarizes the resulting stage proportions.

  • The FLOPs calculation adopts Kaplan et al.’s standard formula and extends it across training stages for proportional contributions.
  • Pre-training and mid-training token counts are assembled from OLMo 2 reports and model-sou trained annealing runs.
  • SFT uses 939,334 samples for 7B and 13B, 866,138 for 32B, a maximum sequence length of 4,096, and 2 epochs.
  • DPO uses 366,700, 377,700, and 377,900 preference pairs for 7B, 13B, and 32B, respectively, with 1 epoch and a 2,048-token maximum.
  • RLVR uses 10M episodes across model sizes, with PPO for 7B and 13B and GRPO for 32B.
  • Table 4 reports FLOPs by stage, including post-training totals and the SFT fraction of post-training compute.

B Additional Figures for Heterogeneous Overfitting

Additional figures show that sub-datasets peak at different training steps, while SRO and Soft SRO search for per-sub-dataset optimal compute and exclude datasets as their budgets are exhausted.

  • B Additional Figures for Heterogeneous Overfitting: Across all models, each sub-dataset reaches maximum validation accuracy at a different training step, confirming heterogeneous overfitting dynamics.
  • SRO: Single roll-out search estimates each sub-dataset’s optimal compute as c∗_i = arg max_c acc(D_i, c).
  • SRO: SRO excludes sub-datasets that have exhausted their compute budgets during a new training run.
  • SRO: SRO iteratively finds the next stopping point among active datasets and rolls them forward by the corresponding compute increment.
  • Soft SRO: Soft SRO starts a new training run with a data mixture accounting for estimated optimal compute budgets.
  • Soft SRO: The Soft SRO initialization computes a normalization factor from optimal budgets and dataset sizes before constructing the new mixture.

D Further Experimental Results on ∆Optimal Compute

Further experiments examine how excluding datasets changes optimal compute, training-instance activity, hardware setup, and seed stability.

  • Excluding a small fraction of the training mixture shifts optimal stopping points for the remaining tasks, with ∆optimal compute varying across sub-tasks.
  • Experiments use B200, H200, RTX A5000, and RTX 3090 GPUs, with commonly available CPU and RAM because these did not bottleneck training.
  • Across models on IES, all training instances are active for the first three epochs, after which the active proportion consistently decreases.
  • Preliminary Qwen2.5 3B experiments across seeds 20, 30, and 40 produced virtually identical performance gains, while Table 6 reports low standard deviation and statistical significance.
  • Table 5 lists overlapping hyperparameters used across methods.

E.3 Method-specific Settings

The experiments use a 10-epoch SFT schedule because some datasets do not overfit within that duration, while continual SFT and MSFT use a 3-epoch compute budget.

  • Training and compute settings: SFT runs for 10 epochs because some datasets do not overfit even by epoch 10.The paper refers to Fig. 2 and Appendix B for this observation.
  • Training and compute settings: Continual SFT and MSFT use a compute budget of C = 3 epochs.
  • Baseline-specific settings: DynamixSFT uses β = 5000 after additional hyperparameter tuning improved performance in the authors’ environment.

F.1 Method-specific FLOPs

The FLOPs analysis compares standard, sequential, dynamic, instance-dependent, searched, and iterative overfitting-aware SFT procedures, including training, validation, search, and checkpoint-related costs. MSFT excludes overfitting subsets iteratively while reverting to earlier checkpoints without additional rollback FLOPs, and its total cost decreases as active mixtures shrink.

  • Baseline methods: Standard SFT trains on all sub-datasets for C units of compute budget.
  • Baseline methods: Continual SFT trains each sub-dataset independently for C units before moving to the next, summing costs across all N sub-datasets.
  • Baseline methods: DynamixSFT adds one-step look-ahead passes for all N sub-datasets at each update, incurring 8|θ| FLOPs per look-ahead token.The look-ahead batch size is B_look-ahead, and updates occur at steps t_u.
  • Baseline methods: IES excludes mastered samples from gradient updates, reducing training FLOPs as exclusions accumulate while validation continues to cover the full dataset.
  • MSFT: MSFT trains for C units on active subsets, adds overfitting subsets to the exclusion set, and rolls back to the earliest overfitting checkpoint without additional FLOPs.Active training and validation costs decline as more sub-datasets are excluded; excluded subsets are validated once at rollback checkpoints.
  • Empirical FLOPs comparison: DynamixSFT incurs 94.9% of training FLOPs on average in look-ahead overhead, while SRO SFT and Soft SRO SFT incur higher total costs because of their additional search phase.IES can cost less than SFT after samples are dropped from the third compute-budget unit onward.

I Disk Storage Footprint

The figures report checkpoints stored on disk across MSFT iterations, measured in model-size multiples, with an average utilization of 4.0 for mSFT.

  • Disk utilization: Checkpoints on disk are measured as multiples of model size |θ| across evaluation steps.Dashed vertical lines mark new roll-outs, while the orange horizontal line indicates average utilization.
Loading 2603.21606v5…