SmallToLarge (S2L): Scalable Data Selection for Fine-tuning Large Language Models by Summarizing Training Trajectories of Small Models

Authors: yu  yang
siddhartha  mishra
jeffrey  chiang
baharan  mirzasoleiman
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 23:02:32

Abstract: Despite the effectiveness of data selection for pretraining and instruction fine-tuning
large language models (LLMs), improving data efficiency in supervised fine-tuning
(SFT) for specialized domains poses significant challenges due to the complexity
of fine-tuning data. To bridge this gap, we introduce an effective and scalable
data selection method for SFT, SmallToLarge (S2L), which trains a small
model, clusters loss trajectories of the examples, and samples from these clusters to
guide data selection for larger models. We prove that during fine-tuning, samples
within the same loss trajectory cluster exhibit similar gradients. Then, we show
that S2L subsets have a bounded gradient error w.r.t. the full data, hence guarantee
convergence to the neighborhood of the optimal solution. We demonstrate through
extensive experiments that S2L significantly improves data efficiency in SFT for
mathematical problem-solving, reducing the training data requirement to just $11$%
of the original MathInstruct dataset to match full dataset performance while
outperforming state-of-the-art data selection algorithms by an average of $4.7$%
across $6$ in- and out-domain evaluation datasets. Remarkably, selecting only 50K
data for SFT, S2L achieves a $32.7$% accuracy on the challenging MATH
benchmark, improving Phi-2 by $16.6$%. In clinical text summarization on the
MIMIC-III dataset, S2L again outperforms training on the full dataset using
only $50$% of the data. Notably, S2L can perform scalable data selection using a
reference model $100\times$ smaller than the target model, proportionally reducing the
computational cost.

Main paper: 
 introduction: 
 In recent years, large language models (LLMs) have revolutionized artificial intelligence by demonstrating an unprecedented ability to understand and generate human language [7] . Among all the contributing factors, the quality and selection of data is becoming increasingly recognized for its importance in training LLMs effectively. Recent research indicates that LLMs benefit more from training for additional epochs on carefully curated data rather than on larger, uncurated ones during pretraining [48] and instruction fine-tuning [67] , making data selection one of the most promising means to unlock the next level of LLMs' language capability. However, while generalist models obtained through pre-training or instruction fine-tuning excel in general language tasks, they may not deliver optimal outcomes in specialized domain, such as mathematics [3, 31, 63, 30, 64] , code [42, 32] , medicine [43, 44, 9] , or finance [57, 9] . These domains are not only critical for real-world applications but also hold substantial economic and societal impacts. 1 : Existing data selection methods depend heavily on the feature representations from a reference model, which makes their effectiveness vulnerable to the quality of training on the target domain [34] . For supervised fine-tuning (SFT), while pretrained models can effectively separate topics (shown in different colors) in natural language (Figure 1a ), they struggle with fine-tuning data that deviates from the pretraining distribution (Figure 1b ). Additionally, the cost of training a reference model escalates with model size (Figure 1c ), making existing data selection methods for large models prohibitively expensive.To maximize performance in specialized domains, models fine-tuned on domain data offer superior capabilities over generalist models [20] . Yet, maximizing the data efficiency in supervised fine-tuning (SFT) for specialized domains remains a challenging and under-explored problem. Firstly, heuristic approaches that are effective in the instruction fine-tuning stage, like manual curation [67] or using advanced models such as GPT-4 for dataset evaluation [8] , are less reliable due to the need for specialized knowledge and become costly with large volumes of uncurated fine-tuning data. Beyond these heuristic methods, other approaches rely on generating representations for each training example using a reference model, often utilizing metrics like perplexity [34] , confidence [46, 53] , or hidden states [1, 48, 61, 4] as features. However, these techniques also fall short in SFT for specialized domains for two reasons: (1) the significant shift between pretraining and SFT data can render these metrics less informative (Figure 1b ), and (2) the computation and memory demands associated with generating representations for each training example become prohibitive, as these specialized domains often require larger models, some with up to 540 billion parameters [10, 43] , leading to substantial scalability challenges (Figure 1c ).To tackle the challenges of data efficiency in SFT for specialized domains, we present SMALLTOLARGE (S2L), an effective and scalable data selection algorithm. S2L operates by first gathering training loss trajectories for each training example using a small model. These trajectories are then clustered, and similar number of examples are selected from these clusters uniformly at random. This process is grounded in our theoretical findings that examples within the same cluster exhibit similar gradients during training, thereby affecting the model similarly. Consequently, subsets sampled from these clusters have a bounded gradient error w.r.t. the full data, allowing for training a comparable model with only a subset of data. Furthermore, we provide a convergence rate analysis for training on these subsets, establishing a robust theoretical foundation for S2L's effectiveness and efficiency.To validate S2L's effectiveness, we applied it to the challenging tasks of SFT for (1) mathematical problem-solving and (2) clinical text summarization. Our experiments on MathInstruct [64] shows that S2L can significantly reduce the required training data size to just 11% of the original dataset size while still matching the performance levels of the full dataset, outperforming current state-of-the-art one-shot and online data selection algorithms by an average of 4.7% across 6 in-and out-domain evaluation datasets. Remarkably, on the MATH benchmark [19] , S2L attained a 32.7% accuracy with just 50K data points, improving the best open-sourced model under 3 billion parameters, Phi-2, by 16 .6%. For clinical text summarization tasks on the MIMIC-III [21] dataset, S2L outperforms training on the full dataset, using only half of the data. Unlike existing methods that require training and getting features from large models, S2L achieves superior data efficiency using a model with as few as 70 million parameters, which is 100× smaller than the largest target model we train with 7 billion parameters.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: The theoretical results provided in Section 4 (and Appendix A) and the experimental results provided in Section 5 support the main claims made in the abstract and introduction. We also conducted experiments with both different data scales and model scales (Section 5.2.1), in both the math and medical domains (Section 5.2 and Section 5.3), and ablation studies (Section 5.4), to reflect how much the results can be expected to generalize to other settings. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
related work: 
 Foundations of Data Selection. Data selection has been well studied for small models and classification tasks. There are one-shot algorithms that select data based on rankings of the proposed training statistics, for example, the L2-norms of error and gradient vectors (EL2N and GraNd) [39] , confidence and its variability across epochs [46] , and the number of times each example is learned but then forgot at the subsequent training step [49] . Besides these heuristic indicators, there are embedding-based pruning algorithms [45] and online selection algorithms with theoretical performance guarantees for efficiency [35, 23, 24, 40, 60] and robustness [59, 62, 16] . Coleman et al. proposed to use the intermediate feature representation of a small proxy model to select data for image classification. Most recently, data selection has shown great potential in more substantial training speedup when implemented on near-storage hardware [41] , and data selection beyond supervised learning of image data, e.g., for self-supervised learning [22] and multimodal learning [1, 33] , also emerged. The proposed algorithm, D4, first applies an embedding-based data de-duplication algorithm [1] and then discards data points that are the closest to the K-Means cluster centroids in the embedding space [45] to ensure diversity. On fine-tuning LLMs, existing work on data efficiency primarily focused on manually curating high-quality instructions [67] , or using strong closed-source models (e.g., GPT-4 [2] or ChatGPT) to rate the quality of each training example [18, 27, 8] . Bhatt et al. implemented an experimental design framework to evaluate the existing data selection methods for instruction fine-tuning of LLMs and found selecting facility locations based on hidden representations (i.e., embeddings) is the most effective. As the only data selection algorithm for specialized domains, SCIP [61] focuses on pruning low-quality code data for training code LLMs. Since it relies on breaking the code syntax to understand the characteristics of low-quality code in the embedding (i.e, hidden states) space, adapting SCIP to domains other than Python code data is non-trivial.Adapting Large Language Models for Specialized Domains. The rapid development of large language models (LLMs) gives rise to new state-of-the-art models in specialized domains. For mathematical reasoning, Galactica [47] , MINERVA [26] and Llemma [3] continue to train an LLM on large-scale math-related web data to improve a model's general scientific reasoning; WizardMath [31] and TinyGSM [30] fine-tune LLMs using supervised data. Similarly for medical LLMs, Cheng et al. continued training pre-trained LLMs on medical text, and [43, 44] fine-tuned PaLM with instruction prompt tuning on medical domain data.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: The paper discusses the limitations of the work in Section 6. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 3. Theory Assumptions and Proofs some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. Justification: Information about the GPU including the memory is provided in Appendix B.3.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
problem formulation: 
 LLM Fine-tuning Objective. Consider a transformer-based language model, parameterized by θ, and denoted as p θ . This model, when provided with a sequence of prompt tokens x = (x 1 , . . . , x M ), generates a sequence of response tokens y = (y 1 , . . . , y L ). The conditional probability of generating y given x is then formulated asEQUATIONNote that y 1:0 is an empty sequence. To adapt the pre-trained LLM for a specialized domain of distribution D, supervised fine-tuning (SFT) is usually employed with a domain-specific training dataset D train = {(x, y) i } n i=1 ∼ D containing pairs of prompt x and annotated response y. The fine-tuning objective is thus to minimize the following negative log likelihood loss, expressed as:EQUATIONData Selection Objective. In a general setting for data selection, we consider a target language model p θ with parameters θ. Given a fixed data budget B, which constrains the number of data points that can be used for training, our objective is to select a subset S ⊆ D train to train the target model, such that it obtains a superior generalization performance. In practice, the subset S is selected based on a reference model r ϕ parameterized by ϕ, which generates representations, confidence scores, or other metrics for each data point (x, y) i ∈ D train , denoted by r ϕ ((x, y) i ), which will be utilized by a data selection algorithm to produce S.In existing data selection algorithms, ϕ is commonly either weights of the pre-trained target model or a target model that has been fully trained on the dataset D train . However, as evidenced by Figure 1 , representations generated by the pretrained model may not always be good enough for data selection in specialized domains, and fine-tuning the target model significantly increases the computational cost of data selection. 
methodology: 
 Training a large target model to obtain feature representations for each example in D train can be computationally intensive. However, a recent finding demonstrates that the training dynamics of most examples are consistent across differently sized models of the same family, and this phenomena even generalizes across different model families [58] . Our proposed method, SMALLTOLARGE (S2L), leverages loss trajectories of training examples collected during fine-tuning a small reference model on the full or a subset of training data.Loss Trajectory. Let ϕ (t) be the parameters of a small LM during training on D train at times t q , q ∈ {1, ..., T }. S2L records the loss trajectory for each data point i at times t q during training the reference model[L proxy i (ϕ (t1) ), . . . , L proxy i (ϕ (t T ) )]whereEQUATIONand T is the length of the loss trajectory. Note that ϕ (t) is trained for a fixed number of iterations from ϕ (t-1) .Assume the parameter vector θ (t) represents the parameters of the target model at the time t. DefineL proxy i = [L proxy i (ϕ (t1) ), . . . , L proxy i (ϕ (t T ) )] and L target i = [L target i (θ (t1) ), . . . , L target i (θ (t T) )] as the training loss trajectory of the example i on the small proxy model and the large target model, respectively. Let H i ∈ R d×d be the Hessian matrix for each example i and assume that the loss function for each example during fine-tuning can be modeled by a second-order Taylor approximation with bounded curvature (c ≤ ∥H i ∥ ≤ C), a reasonable assumption in fine-tuning settings. The following lemma shows that examples with similar loss trajectories on the proxy model have similar gradients throughout the training of the target model. EQUATION)where ϵ ′ = ϵ + 2δ and ∥θ∥ ≤ D for all t.The proof of Theorem 4.1 can be found in Appendix A. EQUATIONwhere L Ci is the centroid of the loss trajectories in cluster C i , and d(•, •) is a distance metric, such as Euclidean distance, used for clustering. For datasets that contain different sources of data, we cluster each source separately. θ t+1 = θ t -α∇L target i (θ t ).(6) Corollary 4.2. Under the assumptions of Theorem 4.1, applying IG with stepsize α to subsets found by S2L, converges to the neighborhood of the optimal solution, as follows:EQUATION) where c ≤ ∥H∥, B = k • K is the total size of the subset, g max is the largest gradient norm of individual examples during training, r min = min j |C j |, r max = max j |C j |, R = min{d 0 , Bg max + ξ/c} and d 0 = ∥θ 0 -θ * ∥ is the initial distance to the optimal solution θ * , and ξ is given by:EQUATIONThe proof can be found in Appendix A.2. 
experiments: 
 In this section, we present the comprehensive experiments conducted to evaluate the efficacy of the proposed data selection method, SMALLTOLARGE (S2L), across two challenging domains (mathematical reasoning and clinical text summarization).We systematically compare S2L against a comprehensive set of open-sourced data selection methods. These methods are categorized based on the type of representation they use and selected as the most representative or best-performing methods as identified in prior work. These include: (1) Random Sampling; selecting examples with the (2) Least Confidence [4] or (3)Middle Perplexity [34] ; (4) High Learnability, determined by the loss decrease before and after full fine-tuning [68] ; and (5) Facility Locations selection based on hidden states [4] . Additionally, we incorporate one online selection techniques: (6) Confidence Curriculum proposed by Varshney et al., which selects examples with decreasing confidence during the training. Given that the optimal reference model may vary for each one-shot selection method, we ensure a fair comparison by adopting the approach used in [34] , which runs each method with both the fully fine-tuned target model and an early fine-tuning checkpoint as the reference model. We report the best results from these setups.Training Settings. We focus on fine-tuning using the MathInstruct dataset [64] with 262,040 training examples for its comprehensive coverage of diverse mathematical fields and its capability in training models to achieve state-of-the-art performance on the standard evaluation benchmarks.We employ the open-source model suites Pythia [6] , Phi-2 [28] , Llama-2 [50] as our base models to validate our S2L algorithm and directly compare its performance against the state-of-the-art.Evaluation Datasets. We follow the framework established in [64] for a comprehensive assessment using several well-regarded datasets, including in-domain and out-of-domain datasets. For the in-domain datasets, we consider GSM8K [11] , MATH [19] , and NumGLUE [36] . For the outof-domain datasets, we consider SVAMP [38] , Mathematics [13] , SimulEq [25] . These datasets collectively span a diverse range of mathematical subjects, such as algebra, probability, number theory, calculus, and geometry. Additionally, some questions in these datasets require the application of commonsense, reading comprehension, and multi-step reasoning. All questions are open-formed.Evaluation Metric. We use the standard evaluation metric for open-formed questions, exact match, which measures the model's accuracy by comparing its generated answers against the correct solutions.For an answer to be considered correct, it must match the reference solution precisely.More details about the settings and baseline implementations can be found in Appendix B. In the first setting, we standardize the number of training steps to correspond to 3 epochs on the full dataset, aligning with [64] . This allows us to maintain a consistent training schedule across different methods and data budgets, ensuring fair and accurate comparisons of the quality of data.SCALING THE DATA: SOTA algorithms fail with small data budgets while S2L stands out across data scales. In Figure 4 , we compare S2L against the baselines from Section 5.1 on Pythia-410M across varying data sizes. The training trajectories used by S2L are based on Pythia-70M, a model approximately 6x smaller than Pythia-410M. With the same number of training steps as the full training, S2L exceeds the full dataset's performance using only 30K examples, only 11% of the full dataset. It leads the runner-up baselines by an average of 4.7%, 4.6% and 2.4% with data budget 30K, 50K and 100K across all six evaluation datasets. While state-of-the-art data selection algorithms like Facility Locations [4] and High Learnability [68] have decent performance with a large enough data budget (i.e., 100K), all SOTA algorithms except S2L cannot even outperform the random sampling baseline when the allowed data size is small (i.e., 30K). Unlike the existing algorithms, S2L consistently outperforms all baselines and even full training across all data sizes. Note that compared to the runner-up algorithm in this setting, Facility Locations, the cost of S2L is much lower in both training the reference model and data selection stages (Figure 5 ), and therefore more scalable to both larger target models or larger data sizes. (Figure 16 ), and multiple choice questions that require multi-step reasoning (Figure 17 ), etc. Therefore, by sampling from different clusters, we make sure the selected examples cover the knowledge required for all topics and skills required for all types of questions.Loss Trajectories can Capture the Similarity Between Data Points As Much As Embeddings of a Fully Fine-tuned Model. We conducted a quantitative analysis to assess how effectively S2L identifies similar examples using loss trajectories from a small model. Assuming math problems under the same topic require similar knowledge and share question formats, we used unknown topic labels during S2L's data selection to check if each cluster predominantly contains a single topic. By calculating the fraction of the most common topic in each cluster and plotting this in Figure 6 (with K=100, excluding clusters of size one), we compared the loss trajectory clusters from S2L (in blue) against those from the embeddings of a fully fine-tuned Phi-2 model (in orange)-considered the ground truth for similarity. Results show that most clusters formed by S2L using the Pythia-70M model are based on a single topic and capture topic similarities more effectively than those from the Phi-2 model's embeddings. This analysis not only confirms the homogeneity within S2L clusters but also highlights the computational efficiency of using loss trajectories on small models to identify representative examples.S2L can improve data efficiency not only for fine-tuning data not only in mathematics but also in other specialized domains. This subsection explores its application to clinical text summarization within radiology reports. This task involves processing the detailed analysis and results listed in the findings section of a radiology report and distilling them into a concise impression section. Such summaries are crucial for providing clinicians with quick and actionable insights from radiological studies.Dataset & Setup. We use the MIMIC-III dataset [21] , a comprehensive collection of radiology reports and findings authored by attending physicians in routine clinical care. We use the same preprocessing procedures as [14, 15] to extract the findings and impression sections and remove invalid reports. Given that access to MIMIC-III requires specific credentials, we provide a synthetic example of a radiology report generated by GPT-4 [2] for illustrative purposes in Table 3 . We employ the Pythia-1B model and keep the training setting consistent with the mathematical reasoning task.Evaluation. Our evaluation of generated clinical summaries on the MIMIC-III dataset's test split employs three key metrics as recommended in [52, 51] : (1) BLEU [37] , which measures word sequence overlap between the generated and reference texts; (2) ROUGE-L [29] , assessing the longest common word sequence; and (3) BERTScore [65] , evaluating semantic similarity using BERT's contextual embeddings. These metrics together offer a comprehensive evaluation, ensuring our summaries are not only precise in language but also meaningful and coherent in the context of clinical information. We compare S2L to random selection, a surprisingly strong baseline as evidenced in Section 5.2, to check the validity of the data selection problem on this dataset and then compare it to training on the full dataset to assess its effectiveness.Results. We compare using 30K examples selected by random vs. selected through S2L. Even with only half of the data, the model trained with S2L selected data achieves similar BLEU and significantly higher ROUGE-L and BERTSCore compared to the model trained on the entire 61.5K data. Meanwhile, training on randomly selected 30K examples performs worse than training on the full dataset on all 3 metrics. These results together demonstrate S2L's effectiveness. We conduct ablation studies on MathInstruct and Pythia-410M to further understand the best practices for using S2L.S2L is robust w.r.t. the length of the trajectories but can benefit more from longer trajectories. Figure 8 The results, shown in Figure 10 , confirm that S2L remains effective when the proxy model is trained on a smaller subset of training data and therefore is scalable to larger datasets without a proportional increase in computational costs.S2L is robust across different clustering parameter values for K. We conducted detailed experiments varying the clustering parameter K, as shown in Figure 11 . The results demonstrate that S2L maintains high performance across different values of K, highlighting the robustness of our method to different clustering parameter choices. We chose K=100 for our experiments as it provided the best average accuracy across the evaluation datasets for the math reasoning task.S2L remains effective and efficient compared to using full data when trained for the same number of epochs. Figure 12 illustrates the relative accuracy to full data across different epochs, comparing S2L-selected data and full data with the same number of epochs. Both in-domain and overall average accuracy are shown. S2L demonstrates superior performance with fewer data and fewer training iterations.S2L supports a range of small models as effective proxies. To understand whether different small models could serve as effective proxies, we used GPT-2 (124M) and Pythia-160M as proxy models for data selection to train Pythia-410M. The results, illustrated in Figure 13 , show that both proxy models perform comparably in guiding the data selection, demonstrating the versatility and effectiveness of using different small models for S2L. In this work, we introduced SMALLTOLARGE (S2L), a scalable data selection method to improve the data efficiency of supervised fine-tuning (SFT) for large language models (LLMs) in specialized domains. By clustering data points based on their training dynamics on smaller models and balanced sampling from all clusters, S2L significantly reduces the required training data size without compromising performance compared to using the entire training dataset. Our comprehensive experiments across the mathematical problem-solving and clinical text summarization domains demonstrate the effectiveness of S2L. 
conclusion and limitations: 
 G S M 8 K M A T H N u m G L U E In -d o m a in A v g S V A M P M a t h e m a t ic s S im u lE q A v g 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes]• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA]Justification: The paper does not release new assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: The paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: The paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Pythia-160M GPT-2(124M)Figure 13 : Per-dataset and average accuracy comparison between using different proxy models (Pythia-160M and GPT-2 (124M)) for data selection. Using both proxy models show comparable performance, demonstrating the effectiveness of different small models as reference models for S2L.Our study does come with its limitations. S2L has been only tested within two domains, mathematics and medicine, and on models up to 7 billion parameters, constrained by our computational resources. Additionally, our experiments employed a fixed training schedule across all methods without further optimization or hyperparameter tuning for each method, including S2L. This unified approach, while it ensures a fair comparison, may not fully capture the potential performance improvement that could be achieved with more tailored training strategies. We encourage further research to extend the application of S2L across a broader spectrum of domains and investigate the impact of hyperparameter tuning on its effectiveness.A.1 Proof of Theorem 4.1Proof. From the assumption that the loss trajectories of examples on the proxy and target models are close:EQUATIONSince i and j are in the same cluster C k based on the proxy model, we have:EQUATION) Using the triangle inequality:∥L target i -L target j ∥ ≤ ∥L target i -L proxy i ∥ + ∥L proxy i -L proxy j ∥ + ∥L proxy j -L target j ∥ ≤ 2δ + ϵ = ϵ ′ .(11) Therefore, at any iteration t:|L target i (θ (t) ) -L target j (θ (t) )| ≤ ϵ ′ , ∀t.(12) Assuming that the loss functions can be approximated by:EQUATIONwhere c i is the loss of example i at the beginning of fine-tuning, and dθ is the distance between the parameters of the pretrained model and those during fine-tuning. Similarly for L target j (θ). The loss difference between i and j is:EQUATIONGiven that|L target i (θ) -L target j (θ)| ≤ ϵ ′, we can write:EQUATIONLet us choose two different values, θ (1) and θ (2) , to generate two inequalities. For dθ (1) , we have: 1 2 (dθ (1) ) ⊤ (H i -H j )dθ (1) + (g i -g j ) ⊤ dθ (1) EQUATIONand for dθ (2) , we have:EQUATIONSubtracting these two inequalities, we get: 1 2 (dθ (1) ) ⊤ (H i -H j )θ (1) -(dθ (2) ) ⊤ (H i -H j )dθ (2) + (g i -g j ) ⊤ (dθ (1) -dθ (2) ) ≤ 2ϵ ′ .(EQUATION) This gives us:EQUATION)Assuming ∥dθ (1) -dθ (2) ∥ ≥ d, we get:EQUATIONA.2 Proof of Corollary 4.2Without loss of generality, assume we select k example from each cluster and we have k ≤ min j∈[K] |C j |. Then the error of the subset in capturing the full gradient will beEQUATIONwhere ḡj is the norm of the average gradient of the selected examples from C j . In practice, we can weight elements of the subset by r min /k, which has a similar effect to scaling the step size when training on the subset. Let g max = max j ∥g j ∥ be the maximum gradient norm during training, r max = max j |C j |, r min = min j |C j |. Then, we getEQUATIONEQUATIONThe first term in RHS of Eq (23) is the error of the subset selected from C j to capture its full gradient and the second term is due to selecting the same number of examples, k, from the larger clusters.Using the above error and following the proof of Theorem 1 in [35] , for a constant step size α ≤ 1/c we get:EQUATIONwhere c ≤ ∥H∥, and B = k • K is the total size of the subset, R = min{d 0 , Bg max + ξ ′ /c} and d 0 = ∥θ 0 -θ * ∥ is the initial distance to the optimal solution θ * .If k ≥ |C j | for any cluster C j , one can simply add (r min /k -1) • ĝj to ξ ′ for the corresponding clusters, where ĝj is the norm of the total gradient of cluster C j and we replace r min in Eq (23) with the size of smallest cluster that has larger than k examples.Pythia. The Pythia models [6] are a suite of large language models (LLMs) developed by EleutherAI licensed under the Apache License 2.0. These models range in size from 70 million to 12 billion parameters and are designed to enable controlled scientific research on transparently trained LLMs across various scales.Phi. The Phi models [28] developed by Microsoft are under the MIT License. Phi-1.5, a transformerbased model with 1.3 billion parameters, and its successor, Phi-2, with 2.7 billion parameters, have been trained on a diverse set of data sources, including synthetic texts and curated websites. The Phi models underscore the potential of small yet powerful language models in understanding and generating human language, empowering a range of NLP tasks. Phi-2, in particular, has raised the bar for reasoning and language understanding among foundation models, matching or even exceeding the performance of models 25 times its size on complex benchmarks.LLaMA 2. The LLaMA 2 models [50] , released by Meta AI and licensed under the LLaMA 2 Community License Agreement, are designed for improved natural language understanding and generation. LLaMA 2-7B, the smallest in this series with 7 billion parameters, has demonstrated competitive performance across various NLP benchmarks despite its moderate size.MathInstruct. The MathInstruct dataset [64] is compiled from 13 diverse math rationale datasets, using both chain-of-thought (CoT) and program-of-thought (PoT) rationales. It ensures comprehensive coverage across various mathematical fields in the 262K training examples, making it a popular resource for fine-tuning large language models (LLMs) for general math problem-solving. MathInstruct is licensed under the MIT license. We utilize the Faiss library [17] to perform efficient K-means clustering of loss trajectories with Euclidean distance with K = 100 and 20 iterations. The hyperparameter K is tuned in the range of {50, 100, 200} based on the average accuracy of the model trained on 30K selected data. We found K = 100 worked the best for both datasets we studied in this paper. Ablations studies on the length and the best time in the training to record the trajectories can be found in Section 5.4.Comparing Reference Models for the Baselines For one-shot selection methods (excluding S2L), we use representations from either step 1000 or the end of fine-tuning Pythia-410M on MathInstruct and reported the better result in Figure 4 and Table 1 . In Table 4 , we include the complete comparison between using early-fine-tuning vs. end-of-fine-tuning model checkpoints as the inference model. For Facility Locations, we further compared using the first hidden states as the feature representation as suggested in [4] to using the last hidden states [56] for the tasks we studied.The ranges for confidence, perplexity, and learnability are chosen according to the best-performing intervals reported in prior research (Section 5.1).Due to memory and computational constraints, for Facility Locations, we calculate pairwise similarity and perform greedy selection on a per-data-source basis. We found this per-source selection approach also yields benefits for S2L as different data sources within MathInstruct exhibit distinct common patterns in their training trajectories. Therefore, we implement S2L also on a per-source basis for MathInstruct, and recommend applying S2L per source when dealing with datasets composed of multiple data sources.Hyperparameters Following the setup used in [64] , we adopt a training regimen with a learning rate of 2e-5, a batch size of 128, a maximum length of 512, and a cosine scheduler with a 3% warm-up period.Experiments Compute Resources We fine-tune all the models with the Huggingface transformers library [55] with Fully Sharded Data Parallel (FSDP) [66] on 4 48G NVIDIA RTX A6000.Table 4 : Complete results used for selecting the best reference model for each one-shot data selection baseline. The choice of early-fine-tuning (step 1000) and end-of-fine-tuning checkpoint follows [34] . The best results selected for Figure 4 are highlighted in cyan. Datasets. We utilize 6 diverse datasets with open-formed questions for evaluating the mathematical reasoning capabilities of models trained with both the full MathInstruct dataset and selected subsets. These datasets, detailed in Table 5 , span a range of mathematical disciplines from early algebra to calculus and linear algebra, covering various types of questions such as multi-step reasoning, arithmetic word problems, and problems from mathematics competitions. This variety ensures a comprehensive assessment across both in-domain and out-domain tasks.Pipeline. We utilize the pipeline provided by [64] 2 , designed to first determine whether the model can be prompted to generate a code snippet. This code snippet, if successfully generated, should be executable and produce the correct answer when run. This code-based evaluation is also used for Phi models [28] . In cases where the model does not directly produce a viable code solution, we employ a "think step-by-step" prompting strategy [54] . This method prompts the model to break down its reasoning process, a technique that has been widely proven effective in fully exploiting the model's problem-solving capacity.Following [14, 15] , we include the six most common modality/anatomy pairs: CT head, CT abdomen, CT chest, MRI head, CT spine, and CT neck, and five less common pairs in the text data: MRI spine, CT sinus, MRI abdomen, MRI pelvis, and MRI neck in the evaluation. There are in total 13.7K test examples after data preprocessing and train-test splitting. Beyond qualitative examples from different clusters, we study how S2L changes the data distribution to outperform using the full fine-tuning dataset as well as using random subsets of the same size that have the same distribution as the original dataset. In Figure 18 , we can observe that S2L not only guarantees a thorough and balanced coverage across the spectrum of topics but also ensures sufficient representation of foundational topics, such as pre-algebra, which lays the groundwork for tackling more complex subjects.This paper introduces a data selection method for large language models (LLMs), aiming to enhance the data efficiency in the supervised fine-tuning (SFT) of these models. NeurIPS Paper ChecklistCode is available at https://github.com/BigML-CS-UCLA/S2L. 38th Conference on Neural Information Processing Systems (NeurIPS2024).https://github.com/TIGER-AI-Lab/MAmmoTH?tab=readme-ov-file#large-scale-evaluation 


Decision: This paper provides a method for dataset selection for fine-tuning small LMs by looking at the the training trajectory. The experimental evidence is quite compelling (especially the extra results show in the rebuttal) which shows that gains from their method seem to improve/maintain with scale. While it would be interesting to see this method applied to pretraining dataset selection, but I understand this is out of scope for this work. The reviewers agree that this method is novel and interesting, and the experimental evidence is compelling, additionally the most negative review seems to ask for baselines that are known to not work well in the literature. For these reasons I vote to accept the paper.

Reviews: review: paper_summary: paper_summary: This paper addresses the challenge of data selection in supervised fine-tuning (SFT) of pretrained language models by introducing a novel method called SmallToLarge (S2L). The S2L method involves collecting loss trajectories from a smaller model, clustering these trajectories, and resampling the SFT data to ensure balanced representation across clusters within the given data budget. The authors tested S2L on tasks such as mathematical reasoning and clinical data summarization, demonstrating superior results compared to several existing offline and online data selection methods.
strength_weakness: strength_weakness: 1. The proposed method may require training the model for multiple rounds, which may not scale up to datasets with millions of examples.
2. The proposed method introduce several additional hyperparameters such as the K in k-means clustering, the training epochs, but the effect of these components are not studied. 
3. The connection between the theory and the proposed method is not clear. More explanation on this matter would be appreciated.
questions: questions: See above
limitations: limitations: Yes
paper_summary: paper_summary: This study presents an innovative technique aimed at improving data efficiency in the supervised fine-tuning of large language models for niche domains. The method leverages the training trajectories of smaller models to inform the data selection for larger models, thereby maximizing the utility of the available data. Experimental outcomes on tasks such as mathematical problem-solving and clinical text summarization demonstrate that this technique can significantly reduce the training dataset size while delivering better performance than other standard approaches.
strength_weakness: strength_weakness: - Overclaim: The paper showcases experiments on only two specialized domain datasets, yet it is presented as if the method is universally applicable. While the introduction offers some context, the title remains overly broad, and the methodology lacks domain-specific adaptations. The authors should refine the framing of their contribution to more accurately reflect its actual scope.

- Marginal Novelty: Data selection using proxy models is not a new concept, as highlighted by reference [12]. Additionally, the use of training trajectories has been explored in prior studies, such as [1].

- Limited Evaluation: It would be beneficial to conduct experiments with larger language models beyond the 7B scale to validate the proposed method's effectiveness. The results only suggest that the method may be useful for smaller language models. Also, it would be better to calculate the win-rate for the summarization dataset.

[1] LESS: Selecting Influential Data for Targeted Instruction Tuning
questions: questions: 1. I am still not sure why using a small fraction of data is better than the performance of using full data on NUMGLUE and MATHEMATICS. Can authors provide more insights on this?
2. How to determine the number of K for clustering?
limitations: limitations: Yes
paper_summary: paper_summary: This paper introduces a data selection method called Small to Large (S2L), which uses the training trajectory of a small model to build clusters and select data based on these clusters. The method is shown to be effective in two domains, math and medicine, with superior performance than other data selection methods and, oftentimes, better than training on the full dataset. A theoretical guarantee is also provided for the convergence of training based on this data selection method. Extensive analysis demonstrates S2L's effectiveness under various compute budgets and across different model series.
strength_weakness: strength_weakness: - As shown in Figure 4, it is clear that for many methods, the SFT has not yet converged. It would be interesting to see how each method affects the convergence of the SFT process and what the relative accuracy looks like, i.e., whether better performance can be further achieved.
- The work studies the Pythia model as the small reference model and multiple models from different model families for task evaluation. It is intriguing to know whether other small models can serve as good reference models, and how the alignment between the pretraining data of the small and large models will affect the effectiveness of this data selection strategy.
- In the LLM era, the size of SFT data scales up quickly. It will be interesting to see how this method scales with the SFT data size, for instance, at M/B token scale.
- It is still tricky to determine whether the effect of such clusters will be influenced by different learning rate schedulers and optimizers.
questions: questions: Please refer to weaknesses.
limitations: limitations: The work addresses the limitations of the limited domain tested and the limited size (7B) of the model used for evaluation.
paper_summary: paper_summary: This paper proposes a data selection approach to reduce the sample size required to conduct supervised fine-tuning (SFT) of LLMs for specific domains. The method achieves it by approximating the training gradients on the full data using only a subset of data.

The experiments were done for SFT tasks, including (1) math problem-solving and (2) clinical text summarization. It was found that with 11% of the original dataset, the yielded model can be comparable to the one trained on the full dataset.
strength_weakness: strength_weakness: 1. Data selection for machine learning has been explored in the literature. Some related works which this paper omits but may need to compare and discuss include:

- influence function based approaches [1,2]

- reinforcement learning based approaches [3]

- data Shapley [4]

2. The generalizability of the proposed theory remains elusive. A small model is used to approximate the gradients. There are two gaps

- The small model can have a distinct loss landscape from the large model. Approximating full gradients of a small model not necessarily extrapolate to a large and different models.

- There is no theoretical guidance on how this small model should be like, which architecture, which parameter size, to be competent of approximating the gradients of the large model. 

Considering the black-box nature of deep learning, especially for LLMs, it seems very hard to build such theoretic foundation for the proposed method.


[1] Wang, Z., Zhu, H., Dong, Z., He, X., & Huang, S. L. (2020, April). Less is better: Unweighted data subsampling via influence function. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 34, No. 04, pp. 6340-6347).

[2] Kong, S., Shen, Y., & Huang, L. (2021, October). Resolving training biases via influence-based data relabeling. In International Conference on Learning Representations.

[3] Yoon, J., Arik, S., & Pfister, T. (2020, November). Data valuation using reinforcement learning. In International Conference on Machine Learning (pp. 10842-10851). PMLR.

[4] Ghorbani, A., & Zou, J. (2019, May). Data shapley: Equitable valuation of data for machine learning. In International conference on machine learning (pp. 2242-2251). PMLR.
questions: questions: From Figure 4, it is surprising that the proposed method and some baselines can outperform the full data performance with around 38% of the samples. How much benefit can the method bring if we continue to scale the data?
limitations: limitations: do not apply

Comments: comment: Dear Reviewer dErU,

We hope our recent rebuttal has been helpful in addressing your concerns. If you have any remaining questions, please don't hesitate to let us know—we’d be more than happy to discuss further. We truly appreciate the time and effort you’ve put into reviewing our work, so we want to ensure we’ve adequately responded to your feedback.

Best regards,

SmallToLarge (S2L) Authors
Dear Reviewer 5vbK,

We hope our recent rebuttal has been helpful in addressing your concerns. If you have any remaining questions, please don't hesitate to let us know—we’d be more than happy to discuss further. We truly appreciate the time and effort you’ve put into reviewing our work, so we want to ensure we’ve adequately responded to your feedback.

Best regards,

SmallToLarge (S2L) Authors
Dear Reviewer mU1m,

We hope our recent rebuttal has been helpful in addressing your concerns. If you have any remaining questions, please don't hesitate to let us know—we’d be more than happy to discuss further. We truly appreciate the time and effort you’ve put into reviewing our work, so we want to ensure we’ve adequately responded to your feedback.

Best regards,

SmallToLarge (S2L) Authors
Thank you very much for taking the time to carefully review our rebuttal. We’re glad to hear that our responses addressed your concerns and appreciate the increase in your rating and confidence.

Your feedback has really helped us improve our work. Thanks again for your valuable input.
Thanks to the authors for the detailed reply. All of my concerns have been addressed.