Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Train on Validation (ToV): Fast data selection with applications to fine-tuning

Ayush Jain, Andrea Montanari, Eren Sasoglu

TL;DR

The paper tackles data-efficient fine-tuning by selecting training examples that most influence performance on a target distribution. It flips the conventional train/validation role, using a small target-validation set to gauge how training on a candidate sample would affect validation loss, and aggregates this across epochs to score the rest of the training pool. The authors formalize ToV, implement token-based adaptations for IT and NER, and demonstrate that ToV-based selection (Score+Random) frequently outperforms random, max-uncertainty, and LESS baselines, with strong empirical gains and a solid theoretical justification via influence-function-like linearizations. This approach offers a fast, gradient-free, and broadly applicable data-selection method that can substantially improve data efficiency in fine-tuning large models.

Abstract

State-of-the-art machine learning often follows a two-stage process: $(i)$~pre-training on large, general-purpose datasets; $(ii)$~fine-tuning on task-specific data. In fine-tuning, selecting training examples that closely reflect the target distribution is crucial. However, it is often the case that only a few samples are available from the target distribution. Existing data selection methods treat these target samples as a validation set and estimate the effect of adding or removing a single sample from the training pool by performing inference on the validation set. We propose a simpler and faster alternative that inverts the usual role of train and validation: we perform inference on the training pool before and after fine-tuning on the validation set. We then select samples whose predictions change the most. Our key insight is that the training samples most affected by fine-tuning on a small validation set tend to be the most beneficial for reducing test loss on the target distribution. Experiments on instruction tuning and named entity recognition tasks show that, in most cases, our method achieves lower test log-loss than state-of-the-art approaches. We support our findings with theoretical analysis.

Train on Validation (ToV): Fast data selection with applications to fine-tuning

TL;DR

The paper tackles data-efficient fine-tuning by selecting training examples that most influence performance on a target distribution. It flips the conventional train/validation role, using a small target-validation set to gauge how training on a candidate sample would affect validation loss, and aggregates this across epochs to score the rest of the training pool. The authors formalize ToV, implement token-based adaptations for IT and NER, and demonstrate that ToV-based selection (Score+Random) frequently outperforms random, max-uncertainty, and LESS baselines, with strong empirical gains and a solid theoretical justification via influence-function-like linearizations. This approach offers a fast, gradient-free, and broadly applicable data-selection method that can substantially improve data efficiency in fine-tuning large models.

Abstract

State-of-the-art machine learning often follows a two-stage process: ~pre-training on large, general-purpose datasets; ~fine-tuning on task-specific data. In fine-tuning, selecting training examples that closely reflect the target distribution is crucial. However, it is often the case that only a few samples are available from the target distribution. Existing data selection methods treat these target samples as a validation set and estimate the effect of adding or removing a single sample from the training pool by performing inference on the validation set. We propose a simpler and faster alternative that inverts the usual role of train and validation: we perform inference on the training pool before and after fine-tuning on the validation set. We then select samples whose predictions change the most. Our key insight is that the training samples most affected by fine-tuning on a small validation set tend to be the most beneficial for reducing test loss on the target distribution. Experiments on instruction tuning and named entity recognition tasks show that, in most cases, our method achieves lower test log-loss than state-of-the-art approaches. We support our findings with theoretical analysis.

Paper Structure

This paper contains 29 sections, 4 theorems, 65 equations, 10 figures, 2 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Train on validation: motivation and algorithm
  3. Related work
  4. Data selection for token-based learning
  5. Score computation for token-based learning
  6. Prediction tasks
  7. Experimental setting
  8. Experiments for instruction tuning
  9. Experiments for named entity recognition
  10. A formal justification
  11. Ideal scores, linearization, influence functions
  12. Train-validation duality
  13. Additional Experimental Details and Results
  14. Experiments for Token-Based Learning
  15. Expanded Results for Instruction Tuning
  16. ...and 14 more sections

Key Result

Proposition 1

Assume there exist $c_0, C_1, M>0$ such that $\nabla^2 \widehat{R}_U(\hat{\boldsymbol \theta}_{k}^{\hbox{\tiny\rm bas}})\succeq c_0{\boldsymbol I}_d$, $\|\nabla \ell(\hat{\boldsymbol \theta}_{k}^{\hbox{\tiny\rm bas}};{\boldsymbol x}_i)\|\le C_1$ for all $k$ and, for all ${\boldsymbol \theta}_1,{\bol

Figures (10)

  • Figure 1: Test log-loss improvement (%) over random selection for instruction tuning with $n = 8 \times 1024$ samples. Each group of bars represents a data-selection strategy (maximum-uncertainty and LESS as baselines); colors show target/training pool configuration (Table \ref{['table:IT_exp']}). Results use Method A (Algorithm \ref{['alg:tov']}) with the Score+Random strategy.
  • Figure 2: Test log-loss vs. number of selected samples $n$ for instruction tuning. (Due to space limits, Exp. 3 plot is in the Appendix.) Lines show mean log-loss over 10 runs; error bars are $\pm$1 standard error. Results use Method A with the Score+Random strategy.
  • Figure 3: Test log-loss improvement (%) relative to random selection for NER at $n = 8 \times 1024$. Each group of bars represents a data-selection strategy; colors show target/training pool configuration (Table \ref{['table:NER_exp']}). Results use Method A (Algorithm \ref{['alg:tov']}) with the Score+Random strategy
  • Figure 4: Expanded version of Figure \ref{['fig:IT_evol']} including the Experiment 3 plot.
  • Figure 5: Test log-loss vs. number of selected samples $n$ for NER. Lines show mean log-loss over 10 runs; error bars are $\pm$1 standard error. Results use Method A with the Score+Random strategy.
  • ...and 5 more figures

Theorems & Definitions (4)

  • Proposition 1
  • Theorem 1
  • Theorem 2
  • Theorem 3