Subset Selection for Fine-Tuning: A Utility-Diversity Balanced Approach for Mathematical Domain Adaptation

TL;DR

This work tackles the high cost of fine-tuning LLMs for mathematical reasoning by introducing a budgeted, utility-diversity subset selection framework. It combines a Pairwise Utility Score, which captures information gain from cross-sample context, with a diversity objective derived from embedding-based dissimilarity, and optimizes a submodular objective via greedy selection to guarantee near-optimal subsets. The approach is evaluated on LLaMA-3 8B and Phi-3 against random and DPP baselines, demonstrating competitive performance with substantially reduced data and compute. The findings suggest practical benefits for domain adaptation in mathematics, enabling near-full dataset performance with a fraction of the data and training time, while maintaining reasoning quality.

Abstract

We propose a refined approach to efficiently fine-tune large language models (LLMs) on specific domains like the mathematical domain by employing a budgeted subset selection method. Our approach combines utility and diversity metrics to select the most informative and representative training examples. The final goal is to achieve near-full dataset performance with meticulously selected data points from the entire dataset while significantly reducing computational cost and training time and achieving competitive performance as the full dataset. The utility metric incorporates both perplexity and Chain-of-Thought (CoT) loss to identify challenging examples that contribute most to model learning, while the diversity metric ensures broad coverage across mathematical subdomains. We evaluate our method on LLaMA-3 8B and Phi-3 models, comparing against several baseline approaches, including random selection, diversity-based sampling, and existing state-of-the-art subset selection techniques.