Improving Task Diversity in Label Efficient Supervised Finetuning of LLMs

Abstract

Supervised Finetuning (SFT) is essential for aligning Large Language Models (LLMs) with human instructions but is constrained by the high cost of manual annotation. Traditional data selection methods focus on geometric prompt diversity, which often ignores task-level semantics and incurs high computational costs. We present Weighted Task Diversity (WTD), a selection strategy that leverages existing task labels and model confidence to prioritize informative samples. Our results demonstrate that WTD can match or exceed full-dataset performance while reducing annotation requirements by up to 80%, achieving a 54.32% win rate against a 4.5x larger random baseline on AlpacaEval.

Introduction

The efficacy of Large Language Models (LLMs) relies heavily on Supervised Finetuning (SFT), where models are trained on high-quality instruction-response pairs. However, the primary bottleneck remains the acquisition of human-annotated data, which is both expensive and time-consuming. To mitigate this, researchers employ data selection strategies to identify a representative subset of the available data.

Current state-of-the-art methods typically emphasize "prompt diversity," using embedding-based geometry to ensure a broad coverage of the input space. While effective, these methods are often computationally prohibitive at scale, sensitive to outliers, and fail to utilize the inherent task-level structure present in many modern datasets.

Methodology

We propose a shift from complex embedding-based selection to an intuitive, task-centric approach. We introduce two primary algorithms:

1. Task Diversity (TD): A baseline method that allocates the annotation budget equally across all available task categories. By using a round-robin sampling strategy, TD ensures that every task—from coding to creative writing—is represented in the finetuning set.
2. Weighted Task Diversity (WTD): Our flagship approach, WTD, incorporates model-specific feedback. It calculates the pre-trained model's confidence on each task by measuring the joint probability of generated tokens. The algorithm then weights the allocation budget toward tasks where the model exhibits the lowest confidence, effectively targeting "hard" tasks while maintaining a minimum representation for "easy" tasks via round-robin sampling.

Experimental Setup

The proposed methods were evaluated using two major instruction-finetuning datasets: FLAN V2 (a 90K subset covering 1,691 tasks) and Databricks Dolly 15K (containing 15,000 human-generated examples across eight categories). Performance was measured across standard benchmarks, including MMLU for factual knowledge, Big-Bench Hard (BBH) for reasoning, and AlpacaEval for human-aligned preference judging via GPT-4 Turbo.

Key Results and Discussion

WTD demonstrates superior efficiency across all benchmarks. On the Dolly dataset, WTD achieved a 4% improvement in MMLU scores compared to training on the full dataset, despite using significantly fewer labels. This suggests that WTD not only selects informative data but also filters out "noisy" or redundant samples that might degrade performance.

Expand

Source, p. 8

Qualitative analysis reveals that WTD naturally prioritizes complex, open-ended tasks. The algorithm allocates a larger portion of the budget to tasks where the base model's confidence is typically lower, such as "Open QA" and "Brainstorming," while reducing the count for simpler, more structured tasks like "Classification" or "Summarization."

Expand

Source, p. 8

In the AlpacaEval benchmark, the WTD-selected 20K subset outperformed a model trained on a 90K random baseline, achieving a win rate of 54.32%. This underscores the effectiveness of prioritizing task difficulty over simple geometric diversity.

Conclusion

Weighted Task Diversity offers a computationally efficient and highly effective alternative to traditional data selection methods. By leveraging "free" task labels and model confidence, WTD reduces annotation costs by up to 80% without sacrificing performance. Future work will explore the scalability of this approach to larger model architectures and its application to datasets lacking pre-existing task labels through automated categorization.

@misc{arabelly2025improving,
  author = {Abhinav Arabelly and Jagrut Nemade and Robert D. Nowak and Jifan Zhang},
  title = {{Improving Task Diversity in Label Efficient Supervised Finetuning of LLMs}},
  year = {2025},
  eprint = {2507.21482},
  archivePrefix = {arXiv},
  url = {https://arxiv.org/abs/2507.21482},
}