Get more for less: Principled Data Selection for Warming Up Fine-Tuning in LLMs

• Pre-fine-tuning uses relevant open, unlabeled samples before targeted fine-tuning on curated task-specific data, which may be sparse, costly, and time-intensive to obtain.Source: §1 · p. 1
• Many existing selection methods are small-scale, while target-matching methods may select samples already represented in pre-training and yield marginal gains at low budgets.Source: §1 · p. 2
• The selection challenge includes improving task performance within constrained budgets, scaling to open language datasets, and generalizing across use cases.Source: §1 · p. 2
• GOT-D prioritizes samples that shift the pre-training distribution toward the target, with its validity proven under certain assumptions; common open sources can proxy for inaccessible pre-training data.Source: §1 · p. 2
• Optimal Transport distance gradients guide selection, with entropy regularization, momentum, and parallel GPU computation enabling selection from millions of samples within a few minutes on one GPU.Source: §1 · p. 2
• GOT-D outperforms existing selection methods across diverse tasks, reducing GPT-2 toxicity by 30% with 10K samples and improving average performance across 8 domain-specific tasks by 1.13% with 150K samples.Source: §1 · pp. 2–3

2.1 PROBLEM FORMULATION

• The setup selects unlabeled DU from candidate DS using target training data DR, then fine-tunes on DR, with the selection budget far below pre-training scale.DS is assumed to approximate the distribution of pre-training data; selected data is capped at N(DU) ≤ N0 ≪ N(DP).Source: §2.1 · p. 3

2.2 OPTIMAL TRANSPORT AND DATA SELECTION

• OT bounds the training–validation loss gap, and when training loss is near zero its distance between distributions serves as a proxy for validation performance.The cited background describes OT with L1 cost and notes analogous results for other distributional distances.Source: §2.2 · p. 4

2.3 DATA SELECTION FOR FINE-TUNING

• Light fine-tuning behaves like minimizing loss on λ·DU +(1−λ)·DP, so target-only matching can miss the pre-training component that selection should pull toward the target.This motivates minimizing distance between the effective mixture and the target rather than between DU alone and the target.Source: §2.3 · pp. 4–5
• Under light fine-tuning with N(DU) ≪ N(DP), OT-gradient selection minimizes the theoretical upper bound on target-task loss.The theorem assumes candidate data approximately matches the pre-training distribution and target training data is identically distributed as target test data.Source: §2.3 · p. 5
• Ranking samples by OT-distance gradients and taking the most negative values selects data that most rapidly decreases distance to the target, with millions of samples handled in minutes on a GPU.The gradient is obtained from the OT dual solution; selection requires one OT solve followed by ranking.Source: §2.3 · pp. 5–6
• The approach assumes candidate data approximates pre-training data and is not intended for small-scale adaptation to expertise radically outside that pre-training scope.The passage gives programming expertise after English-literature-only pre-training as an example where unsupervised fine-tuning at this scale will not be effective.Source: §2.3 · p. 6

3.1 MODEL DETOXIFICATION WITH UNLABELED DATA

• For 20K, GOT-D cut Expected Maximum Toxicity by 0.21 for toxic prompts and 0.12 for non-toxic prompts; Toxicity Probability fell 0.67 to 0.21 and 0.25 to 0.07, while average accuracy shifted 0.422 to 0.408.The authors also report the best performance under the Moderation API.Source: Table 1 · p. 7, §3.1 · p. 7

3.2 ADAPTATION TO DOMAIN-SPECIFIC TASKS

• With 150K selected samples, GOT-D beats other selection baselines on average, gaining around 1.2% over vanilla BERT and around 0.7%–0.9% over DAPT and DSIR.Gains were larger on smaller-target-data tasks, including an approximately 5% gain on ACL-ARC.Source: Table 2 · p. 8, §3.2 · p. 8
• In the 50K-selection, 5K-label setting, GOT-D significantly improves average performance over Vanilla BERT, DSIR, and TAPT/c.The study uses domain-specific classification tasks across eight datasets and four domains.Source: Table 3 · p. 8, §3.2 · p. 8

3.3 TASK-ADAPTION WITHOUT A PRE-DEFINED DOMAIN

• Across both GLUE settings, GOT-D consistently outperforms other selection methods on average and improves over vanilla BERT by 1.04% and 3.13%, respectively.The settings use 50K selected data, followed by fine-tuning with all GLUE training data or 5K samples per task.Source: Table 4 · p. 9, §3.3 · p. 9

• Pre-fine-tuning is introduced as a general paradigm for using open, unlabeled data to improve task adaptation.Source: §4 · p. 9
• GOT-D shifts the pre-training distribution toward the target distribution rather than merely aligning selected data with the target.Source: §4 · p. 9
• The paper reports that GOT-D improves performance across varied tasks and scales efficiently.Source: §4 · p. 9

APPENDIX A EXTENDED RELATED WORK

• Prior language-selection methods designed for pre-training or continued pre-training often yield marginal gains below a million selected samples, while task-specific heuristics can require adjustments and fail to transfer.The cited examples include Xie et al. (2023) and Chowdhery et al. (2022).Source: Appendix A · p. 17
• Fine-tuning data selection remains sparsely addressed; curated-TAPT can select duplicated nearest neighbors, while data valuation and coreset approaches face scalability limits.The discussion cites Gururangan et al. (2020) for curated-TAPT and Shapley and LOO examples among valuation methods.Source: Appendix A · pp. 17–18

B.1 PROOF OF LEMMA 1

• Under light fine-tuning on small data, the fine-tuned model is equivalent to minimizing loss on a weighted mixture of task data and pre-training data.The lemma assumes N(DU) ≪ N(DP), with a mixture weight 0 < λ < 1.Source: §B.1 · p. 18
• The mixture weight λ depends on fine-tuning strength, including the learning rate and number of steps.Source: §B.1 · p. 18

B.2 PROOF OF THEOREM 1

• In the low-data regime, selecting data by the gradient of an optimal-transport problem minimizes a theoretical upper bound on expected target-task loss, under the theorem’s assumptions.The argument bounds the test–training loss gap by an OT distance and uses a first-order approximation when the selection budget is small.Source: §B.2 · pp. 18–19

C.1.1 MODELS

• GPT-2, BERT-base-uncased, and DistilBERT-base-uncased are among the models described for the experiments.The appendix notes GPT-2’s internet-text pretraining and BERT’s BookCorpus and Wikipedia pretraining.Source: §C.1.1 · pp. 19–20

C.1.2 DATASETS

• The experiments use OpenWebTextCorpus for text completion and seven text domains, totaling about 20M samples, for NLU tasks.The NLU domains include Amazon Reviews, PubMed, arXiv, OpenWebTextCorpus, RealNews, Wikipedia, and BookCorpus.Source: §C.1.2 · p. 20

C.1.3 EVALUATION METRICS

• The evaluation measures task-effectiveness gains, performance under selection budgets of 20K–150K, resource use, and consistency across scenarios.These correspond to the paper’s four stated metrics, M1–M4.Source: §C.1.3 · p. 21

C.2 IMPLEMENTATION FOR DATA SELECTION METHODS

• GOT-D selects samples with the largest negative OT gradients; the compared implementations include DSIR importance weighting, DAPT pretraining, TAPT/c nearest neighbors, and uniform domain sampling.GOT-D embeds candidates with distilled BERT and solves OT in embedding space; DAPT assumes domain-specific unlabeled data, while TAPT/c restricts selection to the target domain.Source: §C.2 · p. 21

C.3 RUNTIME ANALYSIS

• 1 hour to embed 2M examples, plus 1–5 minutes to solve OT, are the reported GOT-D costs on a single A100 GPU.The domain relevance test takes under a minute, and distilled-BERT fine-tuning takes 1–5 minutes.Source: §C.3 · pp. 21–22
• 1.2 hours for pretraining and 1.5–2.5 hours for embedding are reported for TAPT/c on a single A6000 GPU.Its selection time ranges from 5 minutes for 2.5K target samples to 1 hour for 393K.Source: §C.3 · p. 22
• 2 hours to compute features, 1.2 hours for importance weights, and 1.5 hours for selection are reported for CPU-only DSIR.DSIR uses multiple cores on an AMD EPYC 7763 64-core CPU.Source: §C.3 · p. 22

C.4 FURTHER DETAILS ON DETOXIFICATION EXPERIMENTS

• REALTOXICPROMPTS supplies toxic and non-toxic prompts, and this study uses the prompts while omitting their continuations.Toxicity scores are reported from both the Perspective and Moderation APIs.Source: §C.4 · p. 22, Table 5 · p. 22
• The Perspective API’s TOXICITY score and the maximum returned Moderation API score serve as the study’s toxicity measures.The Moderation API returns scores for categories including hate, threats, harassment, and violence.Source: §C.4 · p. 22
• Generation uses a 20-token maximum, temperature 1, and nucleus sampling with p = 0.9, with sentences truncated at EOS when generated.Batch generation is used across prompts.Source: §C.4 · p. 22
• GOT-D reduces toxicity and outperforms all baselines on Moderation API evaluation, consistent with Perspective API results.Neither data collection nor data selection used the Moderation API.Source: Table 6 · p. 23, §C.4 · p. 23
• The appendix lists eight zero-shot utility tasks and points to a table breaking down downstream accuracy across them.The tasks include ANLI, BoolQ, HellaSwag, LAMBADA, PIQA, RACE, WiC, and WinoGrande.Source: §C.4 · p. 23

C.5.1 UNSUPERVISED PRE-TRAINING

• Pretraining comparisons use GOT-D and baselines at 50K and 150K selection budgets, with a reported ideal corpus token size of 295.A learning rate of 1e-4 was ideal in most cases; larger rates did not lower training loss.Source: §C.5.1 · pp. 23–24, Table 7 · p. 24, Table 8 · p. 24

C.5.2 SUPERVISED FINE-TUNING

• Supervised fine-tuning varies epochs by target-set size, using 3 epochs for larger datasets and 10 for datasets under 5K samples.The appendix specifies F1 metrics, with macro F1 for CS, News, and Reviews, and micro F1 for Biomed.Source: §C.5.2 · p. 24

C.6.1 EXPERIMENTAL DETAILS AND HYPERPARAMETERS

• GLUE evaluation uses eight tasks and excludes WNLI; RTE and MRPC use a BERT-base model already fine-tuned on MNLI.The appendix says these hyperparameters are applied across tasks.Source: §C.6.1 · p. 24, Table 10 · p. 25, Table 9 · p. 25

C.6.2 ADDITIONAL RESULTS

• Additional results are provided for a restricted budget of 20K pre-fine-tuning samples and 5K labeled target samples.The appendix directs readers to Table 12 for those results.Source: §C.6.2 · p. 25

D.1 ANALYSIS ON PERSPECTIVE API AND MODERATION API

• Perspective API scores may be biased or inaccurate and can change over time, making direct comparisons across dates inconsistent.The authors report significant discrepancies between original REALTOXICPROMPTS scores and scores obtained again in September 2023.Source: §D.1 · p. 25
• Across 10k instances, the APIs’ toxicity scores correlate at 0.5977, with substantial divergences in some cases.The paper points to Table 14 for examples of cases where the APIs diverge.Source: §D.1 · p. 25, Table 14 · p. 27

D.2 GENERALIZATION AND IMPLEMENTATION DISCUSSION

• GOT-D assumes candidate data approximates pre-training data and is not intended for tasks needing wholly out-of-scope expertise, such as programming after English-literature-only pre-training.The discussion notes that candidate and pre-training sources can still overlap in language or domain knowledge despite differing compositions.Source: §D.2 · pp. 25–26
• To avoid extreme selection ratios, the method samples candidate sources by OT relevance, then embeds and selects from the reweighted pool; the full process takes under an hour on one GPU.Downsampling trades some selection-pool size for resources that can support stronger embedding schemes.Source: §D.2 · p. 26

E.1 EXPERIMENTAL DESIGN

• The study tests GPT-2 XL (1.5B) and GPT-neo (2.7B) on AG News and BoolQ, using 5K and 9K target examples, respectively.Source: §E.1 · p. 27
• GPT-2 XL selection uses OWTC; GPT-neo uses a seven-domain substitute pool, with GOT-D and TAPT/c restricted to the closest domain while DSIR spans all domains.Source: §E.1 · p. 27
• 13.9% on AG News and 6.6% on BoolQ: GOT-D gains after pre-fine-tuning with 40k samples and visibly outperforms baseline methods.Source: §E.1 · p. 28
• Zero-shot labels are selected by minimum whole-instance perplexity over candidate labels, with task prompts illustrated in Table 15.Source: §E.1 · p. 28, Table 15 · p. 28

E.2 RESULTS FOR DATASET AGNEWS

• 4% improvement with 5k instances, rising to over 13% at 80k; GOT-D outperforms random selection, DSIR, and TAPT/c across all tested budgets.Source: §E.2 · pp. 28–29, Table 16 · p. 29
• Sentence-transformer performs mostly on par with distilled-BERT, suggesting embedding choice is not critical, while token space is not a proper paraphrase-invariant OT embedding.Source: §E.2 · p. 29, Table 17 · p. 29
• DSIR and TAPT/c select target-matching samples that overlap heavily with OWTC, whereas GOT-D favors underrepresented formal business news that better aligns pre-fine-tuning with AG News.Source: §E.2 · pp. 29–30, Figure 5 · p. 30

E.3 RESULTS FOR DATASET BOOLQ

• GOT-D outperforms baselines on BoolQ with GPT-neo (2.7B) at a 40k-sample budget.Source: §E.3 · p. 30, Table 18 · p. 30
• The BoolQ result contrasts with continued-pretraining settings described as using selection budgets from millions to billions of samples.Source: §E.3 · p. 30
• The reported comparison uses zero-shot accuracy on 1,000 randomly selected BoolQ test samples, with the original model as the zero-data baseline.Source: Table 18 · p. 30