cs.LGMar 23, 2026

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

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

Organizations: KAIST AI · Yonsei University · Samsung Advanced Institute of Technology

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

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.

Figures & tables

MethodTypeEpochs
Magistral ( Rastogi et al., 2025 )Homogenous4
OLMo ( Groeneveld et al., 2024 )Homogenous3
OLMo 2 ( Walsh et al., 2025 )Homogenous2
OLMo 3 ( Ettinger et al., 2026 )Homogenous2
DeepSeek-V3 ( Liu et al., 2025 )Homogenous2
DeepSeek-R1 ( Guo et al., 2025 )Homogenous2-3
Table 1: Status quo. Frontier open-weight models continue to employ homogeneous SFT, where all sub-datasets are trained on the same amount of compute.
Figure 1: SFT is compute-light. Using OLMo 2 as an example, SFT is relatively compute-light, and therefore additional compute usage at this stage is negligible.
Figure 2: Heterogeneous learning dynamics. Multi-task SFT on Qwen3 8B demonstrates that underlying sub-datasets overfitting dynamics vary greatly. This observation is consistent across all other models; visualized in Appendix B .
Figure 3: Divergence of optimal compute upon dataset exclusion. Excluding a small fraction of the training mixture alters the optimization trajectory, shifting optimal stopping points for remaining tasks. (a) Δ\Delta optimal compute varies across individual sub-tasks. (b) This divergence is consistent across model families and scales, averaging an absolute shift of 0.91 epochs. Detailed decomposition across other models available in Appendix D
Model:OLMo 2Qwen2.5Qwen3 
Size:1B0.5B1.5B3B7B8BAverage
 Acc.Ep.Acc.Ep.Acc.Ep.Acc.Ep.Acc.Ep.Acc.Ep.Acc.Ep.
Science and Knowledge
Base32.4—26.1—54.6—12.1—4.0—24.6—25.6—
SFT47.99.7537.50.5065.83.0071.85.0074.52.0077.93.0062.53.88
Continual SFT48.51.9024.61.9566.62.0871.41.8072.91.4077.51.1560.2 -2.31.71
Table 2: Main results. Comparison of six methodologies across six underlying models ( OLMo 2 , Qwen2.5 , and Qwen3 ), evaluating performance across three major task categories. We report both accuracy (Acc.) and the epoch (Ep.) at which the best accuracy was achieved. Continual SFT’s Ep. is the average across benchmarks making values not in intervals of 1/4 epochs like others. The best scores are bolded , and second best underlined .
Figure 4: Further details of main results. [left] mSFT achieves the lowest levels of standard deviation across benchmarks (STD), indicating performance gains are not due to large outliers. [right] Across models, mSFT achieves 1st place the most. The 1st place count does not add up to 60 = 6 ⋅\cdot 10 (models ⋅\cdot benchmarks) as there are cases where 1st place is tied.
Figure 7
 Average
 Acc.Ep.
SFT61.93.88
SRO SFT63.43.75
Soft SRO SFT62.13.79
mSFT ( ours )63.74.12
Table 3: Ablation study results. Comparison of our proposed method ( mSFT ) against two naïve alternative heterogeneous early-stopping algorithms averaged across six underlying models.
Figure 7: Performance on further granular decompositions. Evaluating mSFT across MedMCQA sub-categories using Qwen2.5 3B demonstrates an average accuracy improvement of +1.86% over the SFT baseline, outperforming IES (+0.29%).
Figure 8: Decomposition of performance gains. mSFT ’s accuracy improvement over SFT is decomposed into overfitting prevention benefits and dataset exclusion effects. Minor catastrophic forgetting from hard exclusion is outweighed by gains from mitigating heterogeneous overfitting.
Figure 9: Training loss curve comparison at 8B . Smoothed with moving average with sliding window 10. Dashed vertical lines denote roll-back where a sub-dataset is excluded. Numerical annotation at the bottom indicate the number of remaining sub-datasets at each interval.
Model:Gemma 3 4B
 Science and KnowledgeCommonsense and LanguageMathematic and QuantitativeAverage
SFT59.269.038.559.0
mSFT ( ours )63.0 +3.874.1 +5.139.5 +1.062.7 +3.7
Table 4: Results on a different model family. mSFT ( C=3C=3 ) on Gemma 3 4B . Average denotes the average accuracy across 10 benchmarks. Subscripts denote Δ\Delta relative to SFT.
Model:OLMo 2Qwen2.5Qwen3 
Size:1B0.5B1.5B3B7B8BAverage
Average Accuracy Across 10 Benchmarks
SFT43.633.064.173.276.880.861.9
CoBa46.136.864.672.577.480.963.1 +1.2
mSFT ( ours )46.337.465.074.278.081.463.7 +1.8
Table 5: Comparison with loss-reweighting. mSFT against CoBa across six underlying models. Subscripts denote Δ\Delta relative to SFT. The best scores are bolded .
 SourceFLOPs
Stage7B13B32B7B13B32B
Pre-trainingpaper §2.31.64×10231.64{\times}10^{23}3.90×10233.90{\times}10^{23}1.16×10241.16{\times}10^{24}
Mid-trainingpaper §4.56.30×10216.30{\times}10^{21}4.68×10224.68{\times}10^{22}1.04×10231.04{\times}10^{23}
SFTHF dataset2.97×10192.97{\times}10^{19}5.52×10195.52{\times}10^{19}1.24×10201.24{\times}10^{20}
DPOHF dataset1.89×10191.89{\times}10^{19}3.61×10193.61{\times}10^{19}1.28×10201.28{\times}10^{20}
RLVR-gradtulu3.md script7.15×10197.15{\times}10^{19}1.33×10201.33{\times}10^{20}3.27×10203.27{\times}10^{20}
Table 6: OLMo 2 training FLOPs by stage. “Post” denotes the sum of SFT, DPO, RLVR-grad, and RLVR-roll. Post/Total is the ratio of total post-training FLOPs to total training FLOPs. SFT/Post is the fraction of post-training compute spent on SFT.
Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 10: Heterogeneous learning dynamics. Multi-task SFT demonstrates underlying sub-datasets overfitting dynamics vary greatly.
Figure 11: Heterogeneous learning dynamics. Multi-task SFT demonstrates underlying sub-datasets overfitting dynamics vary greatly.
Figure 17
BenchmarkAssigned CategoryDomainLicense
SciQ ( Welbl et al., 2017 )Science and KnowledgeCrowdsourced science examCC BY-NC 3.0
ARC-Easy ( Clark et al., 2018 )Science and KnowledgeGrade-school science examCC BY-SA 4.0
OpenBookQA ( Mihaylov et al., 2018 )Science and KnowledgeElementary science factsApache 2.0
MedMCQA ( Pal et al., 2022 )Science and KnowledgeMedical entrance examApache 2.0
CommonsenseQA ( Talmor et al., 2019 )Commonsense and LanguageGeneral commonsenseMIT
HellaSwag ( Zellers et al., 2019 )Commonsense and LanguageGrounded sentence completionMIT
Appendix
Table 7: Benchmarks used in our experiments. All datasets are English. Assigned Category denotes the grouping used to report results in Tab. 2 . Each sub-dataset contributes 1,800 training and 200 validation examples.
HyperparameterValue
Learning Rate1×10−51\times 10^{-5}
Learning Rate ScheduleConstant
Batch Size64
Seed20
Sub-dataset Size (train)1800
Sub-dataset Size (validation)200
Appendix
Table 8: Overlapping hyperparameters.
Model:Qwen2.5
Size:3B7B
 SFTmSFTSFTmSFT
Average Accuracy Across 10 Benchmarks
Seed 2073.2574.2576.8078.05
Seed 3073.0574.0576.9078.85
Seed 4072.6573.2576.8078.05
Appendix
Table 9: Seed stability on Qwen2.5 3B and 7B . Subscripts in the Mean row denote Δ\Delta relative to SFT. pp -values are from a two-sided paired tt -test against SFT ( ∗p<0.05{}^{*}p<0.05 ).
Figure 13: Training instances across epochs on IES. Percentage of active training instances per epoch, relative to the initial dataset size at Epoch 1. All models process the complete dataset for the first three epochs, after which the proportion of active instances consistently decreases.
      SoftmSFT
ModelSFTCont.DynamixIESSROSRO(C=1)(C=3)
OLMo 2 1B157.19124.52263.22112.05218.66273.3574.06210.11
Qwen2.5 0.5B58.1047.41103.5437.1980.35120.9331.24105.78
Qwen2.5 1.5B207.90308.98363.97135.80374.98443.58144.21416.26
Qwen2.5 3B461.48727.13760.17313.30652.51879.80320.14770.21
Qwen2.5 7B1089.42965.961742.91663.911717.302057.52–1654.26
Appendix
Table 10: Total PFLOPs for each method across model scales.

Explore similar work

CardsList

  1. Beyond Full-Model Rollback: AuroSFT for Adapter-State Multi-Task Fine-Tuning

    Aug 5, 2026Yue Han, Ziniu LiuSupervised Fine-TuningMulti-Task Learning

    College of Systems Engineering, National University of Defense Technology · College of Computer Science and Technology, National University of Defense Technology

  2. Learning Task Mixtures from Task Affinities: A Probabilistic Graphical Model for Supervised Fine-Tuning

    Jul 16, 2025Prateek Chanda, Saral Sureka, Parth Pratim Chatterjee +3Supervised Fine-TuningMulti-Task Learning

    IIT Bombay · IBM Research · Red Hat AI Innovation +1

  3. SFT Conflicts, RL Coexists: A Theoretical and Empirical Analysis of Multi-Task Learning for LLMs

    Aug 4, 2026Kejian Zhu, Zhuoran Jin, Shangqing Tu +5Supervised Fine-TuningGradient Interference

    The Key Laboratory of Cognition and Decision Intelligence for Complex Systems, Institute of Automation, Chinese Academy of Sciences, Beijing, China · School of Artificial Intelligence, University of Chinese Academy of Sciences, Beijing, China · Tsinghua University, Beijing, China