Holdout-Loss-Based Data Selection for LLM Finetuning via In-Context Learning

TL;DR

This work tackles the data selection problem in fine-tuning large language models by introducing an In-Context Approximation (ICA) that estimates the holdout loss a model would incur after training on a candidate example, using demonstrations from a small holdout set. ICA enables a principled, resource-efficient way to rank and reweight training examples without additional fine-tuning or a reference model. Empirically, ICA-based reweighting improves alignment across supervised fine-tuning and preference-based methods (SFT, DPO, SimPO) across diverse backbones and datasets, with only minimal overhead. The approach offers a practical path to data-efficient LLM alignment, while acknowledging limitations related to holdout set quality and on-policy dynamics, and it opens directions for future work on extending to on-policy settings and broader domains.

Abstract

Fine-tuning large pretrained language models is a common approach for aligning them with human preferences, but noisy or off-target examples can dilute supervision. While small, well-chosen datasets often match the performance of much larger ones, systematic and efficient ways to identify high-value training data remain underexplored. Many current methods rely on heuristics or expensive retraining. We present a theoretically grounded, resource-efficient framework for data selection and reweighting. At its core is an In-Context Approximation (ICA) that estimates the holdout loss a model would incur after training on a candidate example by conditioning on a small, curated holdout set in context. ICA requires no reference model and no additional finetuning. Under a local linearization, ICA is equivalent to a first-order update toward the holdout optimum, motivating its use as a proxy for data value. We derive per-example weights from ICA scores, dynamically reweighting gradient updates as model parameters evolve. Across SFT, DPO, and SimPO, and over diverse backbones and datasets, ICA-based reweighting consistently improves model alignment with minimal overhead. We analyze sensitivity to score update frequency and the choice of $k$ holdout examples for in-context demonstrations, and note limitations for rapidly drifting on-policy updates, highlighting directions for future work. Code and prompts will be released.

Figures (1)

• Figure 1: Average scores assigned to training examples from each domain (normalized to $[0,1]$). The target domain is highlighted in red, and different colors indicate different scoring methods, with our method in blue. Higher scores indicate stronger alignment with the target domain.