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

TL;DR

<3-5 sentence high-level summary>UDS tackles the challenge of efficiently fine-tuning LLMs by online batch selection that accounts for both data utility and data diversity without relying on external resources. It introduces two complementary signals derived from forward-pass logits: a nuclear-norm based intra-sample score capturing optimization utility and diversity, and a memory-buffer–driven inter-sample distance score to encourage global diversity, with a lightweight SRFT-based projection for efficiency. The method combines these scores to select informative, diverse samples, achieving state-of-the-art results across MMLU, ScienceQA, GSM8K, and HumanEval while reducing training time relative to full-dataset SFT. The work demonstrates robust ablations and scalability, suggesting practical impact for scalable SFT of large language models.

Abstract

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 \textbf{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. Code is available at https://github.com/gfyddha/UDS.

Figures (9)

• Figure 1: Schematic of the UDS Framework. In the forward pass, the LLM extract logits $\bm{L}(\bm{x}_t^i;\bm{\theta}_t)$ for each sample. Using $\bm{L}(\bm{x}_t^i;\bm{\theta}_t)$, we first calculate its nuclear norm for the intra-sample importance score $s_\text{intra}^{t,i}$, and then down-projected to calculate the distance $s_\text{inter}^{t,i}$ with historical samples in the memory buffer. Finally, we select the top-$K$ valuable samples to update the LLM.
• Figure 2: Strong correlation between $-\delta\ell(\bm{x}_t^i;\bm{\theta}_t)$ and $\|\bm{L}(\bm{x}_t^i; \bm{\theta}_t)\|_*$ during the training process. Each point represents the correlation coefficient calculated from $B$ pairs of values within a single batch. We extract the result using Qwen-2.5-7B on MMLU.
• Figure 3: Strong correlation between $rank(\bm{L}(\bm{x}_t^i; \bm{\theta}_t))$ and $\|\bm{L}(\bm{x}_t^i; \bm{\theta}_t)\|_*$ during the training process. Each point represents the correlation coefficient calculated from $B$ pairs of values within a single batch. We extract the result using Qwen-2.5-7B on MMLU.
• Figure 4: Illustration of Intra-sample Diversity.Left: Token sequence with high diversity, where the model predicts varied tokens (cat, dog, bird). Right: Token sequence with low diversity, where the model predominantly predicts a single token (cat).
• Figure 5: Performance across different data scales using Llama-3.1-8B on MMLU. We report accuracy when fine-tuning with varying proportions of training data. While all methods improve with more data, UDS consistently achieves the best accuracy and surpasses full-dataset fine-tuning.

Theorems & Definitions (2)

• Lemma 3.1
• Theorem 3.2