Scalable Fine-tuning from Multiple Data Sources: A First-Order Approximation Approach (2024)

Abstract: We study the problem of fine-tuning a language model (LM) for a target task by optimally using the information from n auxiliary tasks. This problem has broad applications in NLP, such as targeted instruction tuning and data selection in chain-of-thought fine-tuning. The key challenge of this problem is that not all auxiliary tasks are useful to improve the performance of the target task. Thus, choosing the right subset of auxiliary tasks is crucial. Conventional subset selection methods, such as forward & backward selection, are unsuitable for LM fine-tuning because they require repeated training on subsets of auxiliary tasks. This paper introduces a new algorithm to estimate model fine-tuning performances without repeated training. Our algorithm first performs multitask training using the data of all the tasks to obtain a meta initialization. Then, we approximate the model fine-tuning loss of a subset using functional values and gradients from the meta initialization. Empirically, we find that this gradient-based approximation holds with remarkable accuracy for twelve transformer-based LMs. Thus, we can now estimate fine-tuning performances on CPUs within a few seconds. We conduct extensive experiments to validate our approach, delivering a speedup of 30 \times over conventional subset selection while incurring only 1% error of the true fine-tuning performances. In downstream evaluations of instruction tuning and chain-of-thought fine-tuning, our approach improves over prior methods that utilize gradient or representation similarity for subset selection by up to 3.8%.

## introduction fine-tuning a language model (lm) has emerged as an effective approach for knowledge transfer on text data. as the scale of lms keeps growing, efficient and scalable fine-tuning methods are in great demand. for instance, parameter-efficient fine-tuning with adapters or low-rank parameterization can significantly reduce the memory usage of fine-tuning @xcite @xcite . in many applications, besides solving a target task of interest, one can access several related data sources, which could be used for data augmentation. examples include multilingual systems such as neural machine translation @xcite , parsing @xcite , data selection @xcite , and targeted instruction tuning @xcite . a fundamental issue in these scenarios is selecting the beneficial data sources amongst all the available data sources, which can be formulated as a subset selection problem. in this paper, we develop efficient algorithms for subset selection in lm fine-tuning given n data sources, which can scale up to handle large numbers of n. classical subset selection methods, such as forward & backward stepwise selection, have been very effective for regression analysis practice @xcite . however, these classical methods are unsuitable for lm fine-tuning because they require repeated training of the base lm on many subsets of tasks, which is not feasible for large n. more recently, data selection methods have been proposed based on selecting influential samples by comparing gradients @xcite or feature representation during training @xcite . these methods depend on the training procedure, leading to noisy outcomes for approximating true fine-tuning performances. in summary, existing subset selection methods require expensive computation or rely on noisy measurements of task-relatedness @xcite @xcite . this paper presents a new approach to scale up subset selection by quickly estimating model finetuning losses without running the fine-tuning procedure. the key idea is first to perform multitask (1) perform multitask training on a base lm using all the samples, leading to a meta-initialization θ ⋆ ∈ r p . we store (randomly) projected gradients of θ ⋆ for every sample. (2) estimate model fine-tuning performances on a list of task subsets using projected gradients as features in logistic regression. we conduct extensive experiments to validate our approach using four benchmarks. we show that gradex can accurately approximate finetuned model losses within 1% errors across various lms. meanwhile, gradex can accelerate forward selection with 30× fewer flops and 25× less gpu hours by circumventing full model finetuning. gradex also speeds up random ensemble with 44× fewer flops and 46× less gpu hours. we consider targeted instruction tuning and chain-of-thought fine-tuning for downstream evaluations of gradex. our approach matches the accuracy of classical subset selection methods with full finetuning while costing only 0.5% of the computation. on the toxigen and truthfulqa benchmarks, gradex outperforms existing selection methods that rely on gradient or representation measurements @xcite ) by 3.8% on average. on strategyqa and commonsenseqa datasets, gradex improves chain-of-thought reasoning accuracy by 2.4% over prior methods. in conclusion, this paper introduces a new algorithm for subset selection in lm fine-tuning given multiple data sources. the algorithm can quickly estimate model fine-tuning performances without performing fine-tuning on each subset. we identify a linearization property of lms (after meta-training), which is empirically verified on twelve transformer-based lms. this empirical finding may be relevant for future developments of fast inference methods of llms. our estimation algorithm unlocks classical subset selection methods such as forward & backward selection to lms. our extensive experiments show that our approach can deliver significant speedups without losing downstream accuracy. in particular, our approach outperforms recent subset selection methods for instruction tuning and chain-of-thought fine-tuning. the code repository and detailed instructions for reproducing our findings are available at https://github.com/ virtuosoresearch/scalable-finetuning . ## problem formulation problem setup: consider fine-tuning a language model to solve a target task. suppose we also have access to n source tasks. because the information from different data sources may conflict with each other or hurt the target performance @xcite , we would like to select a subset of the n tasks. this often happens in transfer learning, for instance, working with low-resource neural machine translation @xcite . because there are 2 n possible choices of subsets, performing the best subset selection is very costly. in this paper, we develop a much faster approach for subset selection that is suitable for running on lms. concretely, given a base fine-tuning procedure on top of an lm (such as lora or qlora @xcite ), let f (s) be the model finetuning loss on the target task trained together with a given subset of auxiliary tasks s ⊆ {1, 2, . . . , n}. a lower value of f indicates that s is more relevant to the target task. best subset selection corresponds to finding a subset that minimizes f (s). ## our approach our algorithm involves two stages. first, we run a meta-training procedure to obtain an initialization θ ⋆ by fine-tuning an lm on all the samples, similar to performing multitask learning @xcite . second, we estimate model fine-tuning losses by solving a logistic regression problem to get θs i for every @xmath0, 2, . . . , m. notably, the second estimation stage can be run entirely on cpus, which will be very fast. the key idea of why this works is a first-order approximation property that we have found empiri- ## experiments we now validate gradex and its use cases across various datasets and models, focusing on the following key questions. does the estimation procedure accurately approximate the true model finetuning losses? how much computational cost does it save relative to classical subset selection methods that require repeated model fine-tuning? how effective are the subset selection methods using the estimation results in downstream evaluation? our experiments show that gradex approximates fine-tuned model losses within 1% error, tested on five different lms including llama-3-8b. gradex reduces the number of flops by up to 43× and gpu hours by 46× compared to conventional subset selection. next, we evaluate gradex for instruction tuning and chain-of-thought finetuning. in both cases, our algorithm performs on par with conventional subset selection that uses full fine-tuing results, while incurring less than 0.5% computation costs. our algorithm outperforms existing selection methods @xcite on toxigen by 3.8% and truthfulqa by 2.4%, while using comparable computation costs. ## conclusion this paper introduced a novel method for estimating lm fine-tuning performances with a first-order approximation approach. the method can significantly accelerate conventional subset selection, thus unlocking their applications to task/data selection for fine-tuning. evaluation across numerous datasets and lms demonstrates the benefit of gradex when applied to subset selection.