Improving Influence-based Instruction Tuning Data Selection for Balanced Learning of Diverse Capabilities (2025)

Abstract: Selecting appropriate training data is crucial for instruction fine-tuning of large language models (LLMs), which aims to (1) elicit strong capabilities, and (2) achieve balanced performance across different tasks. Influence-based methods show promise in achieving (1), by estimating the contribution of each training example to the model’s predictions, but often struggle with (2). Our systematic investigation reveals that this underperformance can be attributed to an inherent bias, where some tasks intrinsically have greater influence than others. As a result, data selection is often biased towards these tasks, not only hurting the model’s performance on others but also, counterintuitively, harming performance on these high-influence tasks themselves. To address this, we propose BIDS, a Balanced and Influential Data Selection algorithm. BIDS first normalizes influence scores of the training data, and then iteratively chooses the training example with the highest influence on the most underrepresented task. Experiments with both Llama-3 and Mistral-v0.3 on seven benchmarks spanning five diverse capabilities show that BIDS consistently outperforms both state-of-the-art influence-based algorithms and other non-influence-based frameworks. Surprisingly, training on a 15% subset selected by BIDS can even outperform full-dataset training with a much more balanced performance. Our analysis highlights the importance of both instance-level normalization and iterative optimization of selected data for balanced learning of diverse capabilities.

## introduction instruction tuning plays a crucial role in eliciting strong capabilities from llms. typically, a pretrained llm is finetuned on a mixture of different instruction datasets, to learn to follow human instructions and achieve strong and balanced performance on a wide range of capabilities (ouyang * work done during internship at uiuc. @xcite @xcite . the importance of data quality for instruction tuning @xcite has spawned many works on data selection @xcite @xcite . among them, influence-based methods, which estimate each training example's influence on the model's prediction on a downstream task @xcite , have shown promise. thanks to recent advances, they have also been scaled to llm-level computations and demonstrated strong potential in selecting highquality instruction tuning data @xcite @xcite . however, influence estimation methods are typically designed to measure the data influence for a single task @xcite . in this study, we demonstrate that existing influence-based data selection algorithms @xcite struggle to balance capabilities across diverse tasks, which is crucial in real-world applications. 1 specifically, our analysis reveals that the influence scores for some tasks can be larger than others' by several magnitudes. this can introduce systematic biases in the data selection process when cross-task influence scores are directly compared, as done in many existing works @xcite . this leads to a couple of pitfalls. first, biasing towards some tasks hurts the model's performance on others, making it more challenging for llms to achieve balanced capabilities. second, perhaps counterintuitively, it may even hurt the model's performance on the very task that the data is biased towards. these issues call for an influence-based selection algorithm designed for training llms to achieve balanced capabilities across diverse tasks. bids, our proposed algorithm, addresses these challenges with two key designs. given a training dataset to select from and a validation dataset representing diverse target tasks, we describe the influence pattern between them with a matrix, where each column consists of the influence scores of different training examples on a specific validation instance. bids first applies column-wise normalization to this matrix, thus setting the influence for different validation instances on the same scale. then, in contrast to prior methods that simply select instances with the highest influence values, bids applies an iterative selection algorithm. at each iteration, this algorithm compares the influence of each candidate with the average influence of those already selected, and selects the one that can provide the largest improvement for the current selected subset. if the current dataset falls short in influence on certain validation instances, then our algorithm will intuitively favor training examples that have high influence on the specific tasks represented by these validation data, and thus promotes balanced multi-task learning. we evaluate bids on ultrainteract @xcite , a strong instruction tuning dataset suite, with llama-3-8b @xcite and @xcite . across seven tasks spanning five capabilities including coding, math, logical inference, world knowledge and instruction following, bids consistently outperforms both influence-and non-influence-based baselines. surprisingly, a 15% subset selected by bids even outperforms full-dataset training in average performance. further analysis highlights the importance of the key design choices of bids, and reveals completion length as a potential factor for biased influence. the contributions of this work include: (1) through systematic analysis, we identify inherent biases in current influence-based instruction tuning data selection algorithms. (2) we propose bids to address this issue and promote balanced learning of multiple capabilities. (3) extensive experiments confirm the strong and consistent performance of bids, and provide valuable insights on what makes a balanced set of influential data. ## background and preliminaries influence-based instruction tuning data selection. estimating the influence of training examples on model predictions is critical for understanding model behavior. traditional methods, including retraining-based @xcite @xcite and gradientbased @xcite ones, have proven effective but are computationally prohibitive when scaling to llms, as they require either retraining on a large number of subsets, or computing at least a forward and backward pass for each training example in order to obtain its gradient @xcite . several recent advances seek to address these challenges by extending gradient-based approaches to scale more efficiently. given a large training dataset to select from and a validation set representing certain targeted capability, less @xcite models the influence between each pair of training and validation examples through lora-based low-dimensional gradient similarity, and then selects training points with the highest influence on the validation set. logra @xcite leverages a low-rank gradient projection algorithm to further improve the efficiency of this process. mates @xcite formulates the pointwise data influence between each training point and the whole validation set, and uses a small model to learn this pointwise influence. upon closer inspection, these llm-scale influence-based data selection methods share a similar problem formulation. they all need a validation set to represent a targeted data distribution and require pointwise data influence between pair of training and validation instances. we aim to extend this to more general-purpose instruction tuning, where the model simultaneously learns multiple diverse capabilities that may require training data from drastically different distributions. concretely, since only less directly targets instruction tuning among the three llm-scale approaches, we ground our study on the specific formulation of less. but we emphasize that due to the highly similar influence modeling patterns , the results of our work should also provide useful insights for other influence-based selection methods. problem setup and notations. assume an instruction tuning dataset d, a validation dataset v, which spans m diverse tasks that we want to optimize the llm's performance for: @xmath0, and an influence estimation method that can compute the influence of each training example on each validation instance. we first compute the influence score between each pair of training and validation data, yielding a |d| × |v| matrix a. each row of a corresponds to an individual training example, and each column a validation instance. element a ij indicates the influence of i-th example from d on j-th instance from v. we dub a an attribution matrix (am) as it reveals the overall attribution pattern from the training set to all target tasks, and each row a i the influence distribution of the i-th training example. our goal is to design a data selection algorithm that can effectively select a subset t from d with size under a pre-defined budget, based on the information presented in a. finetuning the llm on t is supposed to achieve strong and balanced performance on all targeted tasks. the evaluation tasks are specifically chosen to have minimal overlap in terms of the capabilities they benchmark. the size of validation set for each task is also kept the same to avoid bias in the selection process. ## existing influence-based selection fails at balancing diverse tasks we first show that less leads to suboptimal and unbalanced performance in general-purpose instruction tuning. this is quantitatively revealed by our analysis framework, which identifies inherent biases in the scale of influence values across different tasks. insights drawn in this section pave the way for the design choices of bids in §4. ## bids: selecting influential data for balanced capability learning this section introduces bids to address the issues identified in §3. bids has two key design choices. instance-level normalization. this technique addresses the scale differences in influence values across different validation instances, by applying a column-wise normalization to the attribution matrix. specifically, for a validation instance v j , the influence of each training example t i is normalized by a norm @xmath1, where µ j and σ j are the sample mean and standard deviation of all values in column j of a. this normalization ensures the influence scores across different columns are comparable in scale. in other words, if two influence scores of different columns have similar intra-column rankings, then they should also have similar values. iterative selection favoring underrepresented tasks. we further propose an iterative greedy selection algorithm (figure 3 and algorithm 1) training set (𝓓𝓓 = 𝒕𝒕 𝒊𝒊 𝒊𝒊=𝟏𝟏 |𝓓𝓓| ) |𝓥𝓥| |𝓓𝓓| (𝒋𝒋) 𝑨𝑨 𝒊𝒊𝒋𝒋 𝒏𝒏𝒏𝒏𝒏𝒏𝒏𝒏 = 𝑨𝑨 𝒊𝒊𝒋𝒋 -𝜇𝜇 𝑗𝑗 𝜎𝜎 𝑗𝑗 (2) bids selected (𝓣𝓣) candidates (𝓓𝓓 \ 𝓣𝓣) (𝒊𝒊) . . . 𝚫𝚫 𝟏𝟏 (𝒊𝒊) 𝚫𝚫 𝟒𝟒 (𝒊𝒊) 𝚫𝚫 |𝓥𝓥| (𝒊𝒊) 𝚫𝚫 𝟐𝟐 (𝒊𝒊) 𝚫𝚫 𝟑𝟑 (𝒊𝒊) influence validation index component-wise difference 𝜟𝜟 (𝒊𝒊) = max 1≤𝑗𝑗≤|𝒱𝒱| 𝜟𝜟 𝒋𝒋 (𝒊𝒊) = 𝑨𝑨 𝒊𝒊𝒋𝒋 -1 |𝒯𝒯| � 𝑘𝑘: 𝑡𝑡𝑘𝑘∈𝒯𝒯 𝑨𝑨 𝒌𝒌𝒋𝒋 select candidate 𝒕𝒕 𝒊𝒊 * with highest 𝜟𝜟 (𝒊𝒊 * ) candidates ( 𝓓𝓓 \ (𝓣𝓣 ∪ 𝒕𝒕 𝒊𝒊 * ) ) (𝒊𝒊 * ) selected (𝓣𝓣 ∪ 𝒕𝒕 𝒊𝒊 * ) mean(@xmath2: comparison between bids and the task-wise max algorithm used by less. for convenience, we represent the training set d with its attribution matrix (am), in which the i-th row is the |v|-dimensional influence distribution of the i-th training example in d. bids differs from less, which naively selects top-b examples with the highest task-wise max utility scores, in mainly two aspects. §4 contains a detailed walkthrough. to promote the balance across different capabilities. it begins with an empty set. in each iteration, it first computes the average influence distribution of the current selected subset t , denoted as a t ≜ 1 |t | k:t @xmath3. then it iterates through each training example t i in d \ t , and calculates a component-wise difference between a i and a t . the utility ∆ (i) of candidate t i is then defined as the largest component of a i -a t , and the candidate with the highest utility is selected. algorithm 1 bids: iterative selection 1: input: d: the set of all training examples; v: the set of validation examples; b: the number of examples to be selected; @xmath4: the attribution matrix between d and v. 2: initialization: @xmath5, @xmath6: while |t | < b do 4: ∆ (i) = max 1≤@xmath7: i * = arg max @xmath8: @xmath9: end while 8: return: t : selected training examples. in other words, bids prioritizes training examples that maximize influence improvement for the most underrepresented task in the selected data. it fundamentally differs from less, which scores each training example independently and selects the top-ranked ones. ## ablation this section presents ablation studies of the two key components of bids, and analyzes their effects on the composition of selected data. ## completion length as a task-specific factor for spurious influence this analysis aims to explain the counterintuitive phenomenon identified in §3, where the performance of high-influence tasks actually degrades (e.g., mmlu). we focus on the completion length of target tasks, which has been shown important for gradient-based influence and data selection @xcite @xcite . ideally, if influence has a salient correlation with a task-specific factor that is not strongly correlated with downstream performance, then this can be explained, as such spurious influence no longer translates into better performance. figures 6 and 7 show instance-and task-level correlations between influence and completion lengths of validation data. on the task-level, there is a strong negative correlation (@xmath10.769) for 6/7 tasks. the shortest completion length of mmlu indeed translates into the highest average influence. only ifeval, the only task that does not have a single objective answer, contradicts this negative correlation, possibly suggesting other influencing factors. however, the instance-level correlation pattern becomes more intricate, with a global correlation of only -0.083. even for tasks that show strong negative correlations in figure 7 such as mmlu, humaneval, mbpp and gsm-plus, there is no similar trend for instances inside the same task, not to mention ifeval, which even shows a slightly positive correlation. in summary, completion lengths can partly explain the counterintuitive phenomenon, and should be considered for improving the fidelity of influence. but it is still far from the whole picture, suggesting the potential impact of other task-specific factors such as sample style or complexity @xcite , which are harder to quantify and deserve more focused investigation in future work. influence estimation. influence estimation is generally classified into gradient-based and retraining-based approaches @xcite . gradient-based influence estimation focuses on the gradient trace of each training point, and assesses the gradient alignment between training and validation examples @xcite @xcite . retraining-based estimation trains a large number of models on different subsets, and then inspects how their performance changes when a training example is added to these subsets @xcite @xcite . recently both lines of works have been extended to llm-scale applications, covering various aspects including pretraining @xcite @xcite @xcite and instruction tuning @xcite @xcite . there is also a popular line of work aimed at designing efficient proxies for expensive gradient-based or retraining-based estimation by taking inspiration from in-context probing @xcite and uncertainty estimation @xcite . ## conclusion we introduce bids, an influence-based instruction tuning data selection algorithm specifically designed for learning multiple capabilities. motivated by the observation of an inherent bias in influence across different tasks, bids first applies column-wise normalization to the attribution matrix, and then an iterative selection algorithm favoring underrepresented tasks. bids consistently outperforms various baselines as well as full-dataset training with much more balanced performance. our analysis provides insights into the properties of an influential dataset with balanced capabilities.