C luster UCB : Efficient Gradient-Based Data Selection for Targeted Fine-Tuning of LLM s (2025)

Abstract: Gradient-based data influence approximation has been leveraged to select useful data samples in the supervised fine-tuning of large language models. However, the computation of gradients throughout the fine-tuning process requires too many resources to be feasible in practice. In this paper, we propose an efficient gradient-based data selection framework with clustering and a modified Upper Confidence Bound (UCB) algorithm. Based on the intuition that data samples with similar gradient features will have similar influences, we first perform clustering on the training data pool. Then, we frame the inter-cluster data selection as a constrained computing budget allocation problem and consider it a multi-armed bandit problem. A modified UCB algorithm is leveraged to solve this problem. Specifically, during the iterative sampling process, historical data influence information is recorded to directly estimate the distributions of each cluster, and a cold start is adopted to balance exploration and exploitation. Experimental results on various benchmarks show that our proposed framework, ClusterUCB, can achieve comparable results to the original gradient-based data selection methods while greatly reducing computing consumption.

## introduction data selection has been a challenging problem in the supervised fine-tuning (sft) of large language models (llms) @xcite . some researchers propose to use the data influence approximation @xcite to select data samples with the highest influence on target loss optimization during the training process @xcite @xcite . the data influence at a certain training step is approximated as the inner products or the cosine similarities of the gradients of the training and target validation data samples. then, the one-step data influence approximations are computed after every short period and aggregated through the training process. although proven to be effective, the calculation of data influence approximation consumes many computing resources and can be infeasible in practice when the computing budget is restricted. to reduce resource consumption, one simple way is to compute the gradients of all data samples after longer training periods. previous work @xcite shows that the data influence approximated with gradients will soon lose its indication after multiple training steps. hence, simply extending the interval of every two times of gradient computation is likely to result in inferior selected data subsets. another way to lower the computation cost is to reduce the number of data samples needed in the calculation of one-step data influence approximation while maintaining the ability to select the data samples with the highest influences. from the derivation of data influence approximation, we come to an intuition that the training data samples with similar gradients tend to have similar influences on the same target loss optimization. hence, we first perform clustering on all training data samples at the beginning of training according to the similarities of their gradients. in this way, data samples with high influence tend to be concentrated into a few clusters. by picking out the clusters with higher probabilities to contain high-influence data samples, we can avoid the calculation over a large number of low-influence training data samples. with a constrained computing budget, we frame the data selection among clusters as a computing budget allocation problem, which we call the intercluster data selection. however, a challenge lies in the unknown influence distribution of each cluster. to tackle this challenge, we consider it as a multi-armed bandit problem with each cluster as one arm, and the reward of drawing this arm is the influence of a randomly chosen sample (without replacement) from this cluster. then, we adapt the most commonly used upper confidence bound (ucb) algorithm to the inter-cluster data selection problem. to better suit our objective, we record the historical rewards and directly estimate the distribution of each arm. a cold start is added to improve the initial estimations. combining clustering with the modified ucb algorithm, we propose an efficient gradient-based data selection framework clusterucb that can be applied to various gradient-based data selection methods. to verify the effectiveness of clusterucb, we evaluate it on four widely used benchmarks and two state-ofthe-art gradient-based data selection methods @xcite . experimental results demonstrate that clusterucb can achieve comparable results with the original gradient-based data selection methods while greatly reducing the computing consumption. individual data influence data influence function is proposed to evaluate the influence of data samples on model training @xcite . since the evaluation of different combinations of data samples is too costly, some researchers tend to evaluate the influence of individual data samples and treat the sum of these influences as the influence of a data subset. there are two branches in the individual data influence approximation: one is auxiliary model learning and simulation @xcite @xcite , and the other is gradient-based training dynamic approximation @xcite @xcite @xcite . in our work, we propose a framework to efficiently apply gradient-based individual data influence approximation in the sft data selection of llms. ## methodology our work focuses on the efficient data selection for supervised fine-tuning of large language models on a target task. given a pretrained model with parameter θ, the training loss l(•; θ), the training data pool x @xmath0, x 2 tr , ..., x n tr }, the target task t , and the validation data samples representing the target task ## conclusion in this paper, we aim to reduce the computational consumption used in gradient-based sft data selection for llms. our proposed framework first performs clustering over the training data pool based on the intuition that training data samples with similar gradients would have similar influences on target loss optimization. then, we frame the inter-cluster data selection as a computing budget allocation problem which is similar to the multiarmed bandit problem, and modify the ucb algorithm to solve it. combined with the state-of-the-art gradient-based data selection methods, experimental results show that our proposed framework can match the original methods while greatly reducing the computing consumption.