Utility-Diversity Aware Online Batch Selection for LLM Supervised Fine-tuning

Supervised fine-tuning (SFT) is a commonly used technique to adapt large language models (LLMs) to downstream tasks. In practice, SFT on a full dataset is computationally expensive and sometimes suffers from overfitting or bias amplification. This facilitates the rise of data curation in SFT, which prioritizes the most valuable data to optimze. This work studies the online batch selection family that dynamically scores and filters samples during the training process. However, existing popular methods often (i) rely merely on the utility of data to select a subset while neglecting other crucial factors like diversity, (ii) rely on external resources such as reference models or validation sets, and (iii) incur extra training time over full-dataset training. To address these limitations, this work develops UDS (Utility-Diversity Sampling), a framework for efficient online batch selection in SFT. UDS leverages the nuclear norm of the logits matrix to capture both data utility and intra-sample diversity, while estimating inter-sample diversity through efficient low-dimensional embedding comparisons with a lightweight memory buffer of historical samples. Such a design eliminates the need for external resources and unnecessary backpropagation, securing computational efficiency. Experiments on multiple benchmarks demonstrate that UDS consistently outperforms state-of-the-art online batch selection methods under varying data budgets, and significantly reduces training time compared to full-dataset fine-tuning.

Large language models often need to be fine-tuned on task-specific data before they work well in specialized settings, but training on all available data can be expensive and sometimes unnecessary. Some training examples are more useful than others, and repeatedly learning from redundant or poorly timed examples can waste computation. In this work, we propose UDS, a method that helps a model decide which examples in each training batch are most worth learning from. UDS looks at the model’s own predictions to identify examples that are both useful for learning and diverse from recently selected examples. Unlike some prior methods, it does not require an extra validation set or a separate reference model. Across several benchmarks, UDS improves model performance while using only part of the training data and reducing training time. This makes fine-tuning large language models more efficient and practical, especially when computation is limited.