SFT-GO: Supervised Fine-Tuning with Group Optimization for Large Language Models

TL;DR

SFT-GO introduces a token importance-based grouping and optimization method for supervised fine-tuning of large language models, leading to improved performance and robustness across benchmarks.

Contribution

It proposes a novel token grouping and weighted loss optimization approach, enhancing fine-tuning effectiveness over existing methods.

Findings

Consistently outperforms baseline methods on LLM benchmarks.

Improves model robustness across different datasets and models.

Provides theoretical analysis of convergence rate.

Abstract

Supervised fine-tuning (SFT) has become an essential step in tailoring large language models (LLMs) to align with human expectations and specific downstream tasks. However, existing SFT methods typically treat each training instance as a uniform sequence, giving equal importance to all tokens regardless of their relevance. This overlooks the fact that only a subset of tokens often contains critical, task-specific information. To address this limitation, we introduce Supervised Fine-Tuning with Group Optimization (SFT-GO), a novel approach that treats groups of tokens differently based on their importance.SFT-GO groups tokens in each sample based on their importance values and optimizes the LLM using a weighted combination of the worst-group loss and the standard cross-entropy loss. This mechanism adaptively emphasizes the most challenging token groups and guides the model to better handle different group distributions, thereby improving overall learning dynamics. We provide a theoretical analysis of SFT-GO's convergence rate, demonstrating its efficiency. Empirically, we apply SFT-GO with three different token grouping strategies and show that models trained with SFT-GO consistently outperform baseline approaches across popular LLM benchmarks. These improvements hold across various datasets and base models, demonstrating the robustness and the effectiveness of our method.