SMART : Submodular Data Mixture Strategy for Instruction Tuning (2024)

Abstract: Instruction Tuning involves finetuning a language model on a collection of instruction-formatted datasets in order to enhance the generalizability of the model to unseen tasks. Studies have shown the importance of balancing different task proportions during finetuning, but finding the right balance remains challenging. Unfortunately, there’s currently no systematic method beyond manual tuning or relying on practitioners’ intuition. In this paper, we introduce SMART (Submodular data Mixture strAtegy for instRuction Tuning) — a novel data mixture strategy which makes use of a submodular function to assign importance scores to tasks which are then used to determine the mixture weights. Given a fine-tuning budget, SMART redistributes the budget among tasks and selects non-redundant samples from each task. Experimental results demonstrate that SMART significantly outperforms traditional methods such as examples proportional mixing and equal mixing. Furthermore, SMART facilitates the creation of data mixtures based on a few representative subsets of tasks alone and through task pruning analysis, we reveal that in a limited budget setting, allocating budget among a subset of representative tasks yields superior performance compared to distributing the budget among all tasks. The code for reproducing our results is open-sourced at https://github.com/kowndinya-renduchintala/SMART.

## introduction one of the main goals of artificial intelligence (ai) research is to build machines that can communicate @xcite , and an essential part of communication is to understand and follow instructions. large language models (llms), which are pre-trained over massive text corpora on next-tokenprediction objective, can perform a wide range of nlp tasks via "prompting" @xcite @xcite @xcite . instruction tuning @xcite @xcite is an approach that further enhances the instruction-following ability and generalizability of pre-trained llms to unseen tasks. it involves fine-tuning an llm on a collection of instruction-formatted instances (encompassing multiple tasks) -each consisting of an instruction (or task description), an optional input, the corresponding output (the ground truth) and optionally a few demonstrations/examples. it is a special case of multitask learning where the llm is finetuned on a collection of instruction-formatted multitask datasets @xcite . finetuning on multiple tasks simultaneously, allows the model to share and transfer information across tasks, resulting in a better common internal representation that is preferred by all tasks while suppressing taskdependent noise @xcite . consequently, the model learns to generalize to unseen tasks by discerning helpful cues from both implicitly and explicitly related tasks that it has previously seen. the performance enhancement from instruction tuning is heavily contingent on data quality, data quantity, and task composition @xcite . studies by @xcite and @xcite have shown that while scaling the number of tasks is important, the relative proportion of various tasks (mixture weighting) merits as much attention for optimal instruction tuning. intuitively, we want the model to see enough data for a given task that it can perform well on it, but not to see so much data that it memorizes the training set @xcite . @xcite performed manual tuning of various benchmark proportions and decided on a final mixture, whereas @xcite studied the impact of removing each benchmark from the finetuning mixture and relied on their practioners' intuition from there on, to decide on the exact proportions of benchmarks. in this work, we would like to explore a more systematic approach to mixture weighting. specifically, we are motivated by the fact that in a large multitask dataset like flan 2022 @xcite , which has 1840 tasks, there will likely be many similar tasks leading to redundancies and not all of them may require sampling in equal proportions. for instance, there might be many tasks of the type natural language inference (nli), and it might be enough to sample relatively more instances from a few representative nli tasks and less from the others. furthermore, which samples we select from each task is also crucial because the samples should faithfully represent the task at hand. a random subset may fail to do this as it can miss out on essential corner cases. with this context, we focus on the following two fundamental research questions (rqs) that form the basis for our subsequent inquiry: as a first step towards addressing both of the above rqs, we first define a common subset selection problem (more formally stated in section 3) as follows -given a huge collection of m instructionformatted task datasets, a task budget m ′ ≤ m and a total budget (n ′ ) of (prompt, response) pairs, which m ′ tasks to select? and how many instances to select from each of these m ′ tasks and which instances to select? note that rq1 is an instance of this problem where m ′ = m . constrained submodular maximization (section 2) proves to be a good model for discovering representative subsets (or coresets) of a massive training dataset (or ground set) that acts as surrogate (i.e., achieves similar performance) and are much better than uniformly-at-random subsets. intuitively, this is because submodular functions model information in subsets, and hence maximizing a submodular function subject to a constraint yields non-redundant subsets of the ground set. an essential feature of this model is that it returns weighted subsets, i.e., each sample in the coreset comes with an associated score, which indicates how important the sample is. inspired by submodular functions, we propose our solution (section 3) to the above subset selection problem for instruction tuning that works in two stages. in the first stage, we select a weighted subset of tasks from the full dataset where the weights will determine how many samples to select from each task. in the next stage, we select samples from each task based on the assigned task budgets. note that the submodular functions used in each stage are not necessarily identical (section 4.8). the main contributions of our work can be summarized as follows: definition 1. (submodular function) a given set function f : 2 @xmath0, y ⊆ v, where x ⊆ y and for all v / ∈ y , the following inequality holds true: intuitively, the definition states that -adding an element to a smaller set x yields more value gain than adding the same element to a superset y of x. definition 2. (cardinality constrained submodular maximization) given a submodular function f : 2 @xmath1, the constrained submodular maximization problem involves finding s * such that the above cardinality-constrained submodular maximization problem is np-complete @xcite . however, if f is monotone submodular (i.e., f (x) ≤ f (y ) whenever x ⊆ y ), @xcite show that a simple greedy algorithm described in algorithm 1 can be used to find an approximate solution s greedy , with a guarantee that f (s greedy ) ≥ (1 -1/e)f (s * ). 1 algorithm 1 the naïve greedy algorithm 1 (also known as naïve greedy) requires o(n ′ .|v|) function evaluations which is costly in practice. accelerated greedy @xcite , also known as lazy greedy, can be used instead, which leverages submodularity and offers a more efficient heap-based implementation of the same algorithm. ## approach inspired by submodularity (section 2), in this section, we introduce a novel data mixture strategy, smart, as a technique to solve the following subset selection problem introduced in section 1: consider a collection @xmath2, . . . , t m } of m instruction-formatted task datasets where each smart models the above problem as a sequence of two cardinality-constrained submodular maximization problems. the first one is to select a weighted subset of m ′ tasks and the second one is to select instances -a total of n ′ instances from these tasks. the weights obtained in the first stage will be used to determine how many instances we sample from each task. ## conclusion & future work in this paper, we introduced smart -a novel data mixture strategy for instruction tuning that utilizes a submodular function to assign importance scores to tasks, determine the mixture weights, and also select non-redundant samples from each task. further, we also reveal that in a low-budget setting, splitting the budget among a small subset of representative tasks yields superior performance when compared to dividing it among all tasks, which suggests that task scaling should be done more judiciously. future work could explore making this method more model-specific (e.g., modify task embedding computation) and also possibly modify the approach for targeted instruction-data selection for creating expert llms that specialize in specific skills such as math, code etc.