SFTMix: Elevating Language Model Instruction Tuning with Mixup Recipe

  • 2024-10-07 18:52:21
  • Yuxin Xiao, Shujian Zhang, Wenxuan Zhou, Marzyeh Ghassemi, Sanqiang Zhao
  • 0

Abstract

To induce desired behaviors in large language models (LLMs) forinteraction-driven tasks, the instruction-tuning stage typically trains LLMs oninstruction-response pairs using the next-token prediction (NTP) loss. Previouswork aiming to improve instruction-tuning performance often emphasizes the needfor higher-quality supervised fine-tuning (SFT) datasets, which typicallyinvolves expensive data filtering with proprietary LLMs or labor-intensive datageneration by human annotators. However, these approaches do not fully leveragethe datasets' intrinsic properties, resulting in high computational and laborcosts, thereby limiting scalability and performance gains. In this paper, wepropose SFTMix, a novel recipe that elevates instruction-tuning performancebeyond the conventional NTP paradigm, without the need for well-curateddatasets. Observing that LLMs exhibit uneven confidence across the semanticrepresentation space, we argue that examples with different confidence levelsshould play distinct roles during the instruction-tuning process. Based on thisinsight, SFTMix leverages training dynamics to identify examples with varyingconfidence levels, then applies a Mixup-based regularization to mitigateoverfitting on confident examples while propagating supervision signals toimprove learning on relatively unconfident ones. This approach enables SFTMixto significantly outperform NTP across a wide range of instruction-followingand healthcare domain-specific SFT tasks, demonstrating its adaptability todiverse LLM families and scalability to datasets of any size. Comprehensiveablation studies further verify the robustness of SFTMix's design choices,underscoring its versatility in consistently enhancing performance acrossdifferent LLMs and datasets in broader natural language processingapplications.

 

Quick Read (beta)

loading the full paper ...