PaperlayerView exact revision on arXiv

Source-linked AI summary

Every Batch Is Its Own Validation Set: Leave-One-Out Gradient Matching for Online Data Selection in LLM Fine-Tuning

Hongyu Chen, Xinyi Luo, Ming Zhao, Lin Tang, Zihan Xu, Jing Li, Yuxuan Wang, Haoran Deng, Wei Zhang

arXiv:2610.00436v1cs.LG

TL;DR

In-sample gradient matching can favor noisy examples and is capped at the full-batch solution because each candidate contributes to its own target. LOOM removes that diagonal to estimate update error without bias, then selects SNR-weighted subsets. Across four SFT tasks and seven backbones, it improves on full-batch training and in-sample matchers.

  • Problem

    Validation-free gradient matching targets the candidate batch using examples that also appear in that target, and in experiments its selectors fall between random selection and full-batch training.

    Source: §1 · p. 1

  • Method

    LOOM removes the gradient-Gram diagonal to form leave-one-out targets, yielding an unbiased estimate of update error and weights proportional to per-example gradient SNR.

    Source: §1 · p. 2

  • Results

    On four SFT tasks and seven backbones, LOOM improves on full-batch training by 2.3–2.4 points and on all in-sample matchers by 2.4.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The leave-one-out objective can favor a subset over the whole batch when per-example SNR heterogeneity exceeds the stated threshold.

    Source: §1 · p. 2

    LOOM assumes i.i.d. candidates, a metric fixed within a step, and finite noise variance; duplicated examples may require deduplication or grouped leave-one-out.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Online batch selection fine-tunes a language model on the most useful part of each candidate batch. Selectors that match the gradient of the candidate batch are attractive because they need no held-out data, yet they rarely beat training on the whole batch. We show why. In-sample gradient matching uses every example as part of its own target, so its objective credits each example with its own gradient noise. This is the covariance penalty that makes training error optimistic, now sitting on the diagonal of the gradient Gram matrix: it steers selection toward the noisiest examples and makes the full batch the best solution the objective can reach. The fix costs nothing. For each example, the other candidates form an independent sample of the data distribution, so removing the diagonal turns the matching objective into an unbiased estimate of the update's error with respect to the population gradient. The minimizer of this leave-one-out objective weights examples by their gradient signal-to-noise ratio (SNR), and whenever per-example SNR is heterogeneous enough, half of a batch yields a lower-error update than the whole batch; we give the exact condition. We build \method{} on this principle. It computes the Gram matrix in the metric of the Adam preconditioner during the ordinary backward pass, selects a weighted subset greedily with a $(1-e^{-\gamma})$ guarantee, and uses no held-out data. Across four fine-tuning tasks and seven backbones from 1.5B to 8B parameters, LOOM improves on full-batch training by 2.3 and 2.4 points on Llama-3.1-8B and Qwen2.5-7B, exceeds every in-sample gradient matcher by 2.4 points and the validation-guided GREATS and OPUS by 1.6--2.0, and selects injected label noise at under a fifth of its base rate.

1 Introduction

The section explains why in-sample gradient matching is capped by the full-batch update and proposes leave-one-out targets to estimate update error without held-out data. LOOM implements this principle and improves on full-batch training and existing selectors in the reported experiments.

1 Introduction

  • Existing in-sample matchers reproduce the candidate-batch gradient, making the full-batch update their objective’s ceiling.In the experiments, each in-sample matcher falls between random selection and full-batch training.Source: §1 · p. 1
  • The Gram-matrix diagonal adds each example’s self-similarity to its target, crediting its own gradient noise and steering selection toward noisy examples.The paper identifies this diagonal term as the covariance penalty behind optimistic training error.Source: §1 · pp. 1–2
  • Leave-one-out targets use the other candidates as an independent sample, making the estimated alignment with the population gradient unbiased.The resulting objective estimates update error; its optimal weights track per-example gradient SNR, and a subset beats the whole batch when SNR heterogeneity exceeds B/b.Source: §1 · p. 2
  • LOOM computes per-example LoRA gradients in the Adam-preconditioner metric during the ordinary backward pass and greedily selects a weighted subset.It provides a (1 − e^-γ) guarantee, adds 5% to step time, and requires no validation data.Source: §1 · p. 2
  • On four SFT tasks and seven backbones, LOOM improves on full-batch training by 2.3–2.4 points and on all in-sample matchers by 2.4.It is competitive with or better than validation-guided selectors and filters injected label noise.Source: §1 · p. 2

2 Related Work

Prior online selection methods use loss, gradient, reference-model, validation-gradient, or utility and diversity signals. Gradient-matching methods target the batch or dataset gradient; LOOM differs in the target it matches and connects self-influence with cross-example agreement.

2 Related Work

  • Online selectors span loss or gradient scoring, reference-model comparisons, validation-gradient alignment, and validation-free utility or diversity signals.LOOM is validation-free and gradient-based, but differs from these methods in its matching target.Source: §2 · p. 2
  • Gradient-matching and coreset methods target the in-sample batch or dataset gradient, the target identified as optimistic in the paper.Their guarantees rely on weak submodularity, which the paper says its objective retains.Source: §2 · p. 2
  • Prior work links self-influence to mislabeled or outlying examples and generalization to directions shared across examples.LOOM connects these views: the diagonal rewarded by in-sample matching signals noise, while removing it leaves cross-example agreement as signal.Source: §2 · pp. 2–3

3 Gradient Matching Overfits the Batch

The section shows that including each example in its own gradient-matching target adds diagonal optimism, favoring noisy examples and limiting the objective to the full-batch solution. It derives leave-one-out SNR weighting and conditions under which a subset lowers risk, and relates the estimated risk to convergence.

3 Gradient Matching Overfits the Batch

  • The estimated risk is the only selector-dependent term in the convergence bound for preconditioned SGD.The full batch has risk tr(MtΣt)/B, while leave-one-out matching minimizes an unbiased estimate of risk up to a constant.Source: §3 · pp. 3–4
  • In-sample matching underestimates selector error through a covariance penalty proportional to the weight placed on noisy examples.This penalty is the gap between training and test error; the diagonal term makes the objective prefer noisy examples.Source: §3 · p. 3
  • In-sample matching’s minimizer is the uniform full batch, which reproduces the batch gradient together with its noise.Facility location shares the flaw because each selected point covers itself with maximal similarity.Source: §3 · p. 4
  • Under the signal-plus-noise model, the leave-one-out objective estimates update risk without bias, and its optimal weights are proportional to per-example SNR with shrinkage as batch signal weakens.The attained reduction is monotone submodular, and the best size-b subset consists of the b examples with largest SNR.Source: §3 · p. 4
  • With normalized weights and equal signal loadings, the best b-subset risk is m/(b AM_b), versus m/(B HM_B) for the uniform full batch; real SFT batches have H1/2 values between 2 and 6.Monte Carlo risks match Theorem 1 within 8%.Source: §3 · p. 4

4 LOOM

LOOM forms per-example gradients and their Gram matrix in the Adam preconditioner’s metric, then uses leave-one-out targets to choose a weighted subset. Greedy selection optimizes the estimated risk before an unchanged AdamW step.

Gram matrix in the optimizer’s metric

  • One backward pass produces preconditioned per-example LoRA gradients and their Gram matrix, which LOOM uses to select a weighted subset before an unchanged AdamW step.The Adam metric uses the previous step’s second moment, fixed before the batch is drawn.Source: §4 · p. 4, Figure 2 · p. 5

Objective and greedy selection

  • LOOM’s leave-one-out targets exclude each candidate’s diagonal entry, validating it against the other B−1 candidates.The target vector and ridge-regularized objective define the estimated risk that selection seeks to reduce.Source: Figure 2 · p. 5, §4 · p. 5
  • Forward selection adds the candidate with the largest positive residual correlation and solves for nonnegative weights on the selected set.The ridge helps prevent weights from fitting noisy targets.Source: §4 · p. 5
  • The greedy guarantee is G(S) ≥ (1 − e^−γ) times the best gain achievable with a support of size at most b.The stated bound follows from the objective’s monotonicity and submodularity ratio.Source: §4 · p. 5

Update

  • The selected weights are normalized so the update retains the scale of an ordinary AdamW step.The stored per-example gradients are used for the update without a second backward pass.Source: §4 · p. 5

5 Experiments

Across the evaluated fine-tuning tasks and backbones, LOOM outperforms full-batch training and gradient-matching baselines. Experiments examine the setup, ablations, mechanisms, efficiency, and additional settings.

5.1 Setup

  • Experiments use four SFT tasks, seven backbones, zero-shot evaluation, and four-seed means; selectors keep 64 of 128 candidates per step.The main backbones are Llama-3.1-8B and Qwen2.5-7B, with five others testing generality.Source: §5.1 · p. 6

5.2 Main Results

  • 56.5 on Llama-3.1-8B and 70.9 on Qwen2.5-7B are LOOM’s best averages, +2.3 and +2.4 over Regular, respectively.LOOM is also 2.9 and 2.8 points above GradMatch on those backbones.Source: §5.2 · p. 6
  • LOOM exceeds validation-guided GREATS and OPUS by 1.6–2.0 points without held-out data.On five further backbones from 1.5B to 8B parameters, it improves on Regular by 1.7–1.9 points and on UDS by 0.7–0.8.Source: §5.2 · pp. 6–7

5.3 Ablations

  • Keeping the diagonal loses 2.79 points and more than doubles gradient error, from 0.61 to 1.38.Two-fold cross-fitting loses 0.32 points; signed weights lower gradient error to 0.58 yet lose 1.51 points.Source: §5.3 · p. 7

5.4 Mechanism Analysis

  • LOOM’s update error stays below Regular for every b/B ≥ 1/4, reaching 0.61 at b=B/2 against a 4,096-example reference gradient.Across seven backbones and four tasks, gains over Regular rise with median SNR heterogeneity (Spearman ρ = 0.86).Source: §5.4 · p. 7, Figure 3 · p. 7
  • With 20% corrupted responses, LOOM selects 3.6% corrupted examples; at 40% corruption, it loses 3.2 points versus 6.9 for Regular and 7.4 for GradMatch.The authors attribute filtering to corrupted responses having a large diagonal entry but a small leave-one-out target.Source: Figure 3 · p. 7, §5.4 · p. 8

5.5 Efficiency and Generality

  • LOOM reaches Regular’s final accuracy after 57% of steps and 60% of wall-clock time, at 0.95× Regular’s throughput.It beats Regular at every ρ ≥ 1/4 and saturates at ρ = 1/2; gains grow from +1.12 to +2.41 as B increases from 16 to 256. Beyond single-task SFT, GREATS wins on TydiQA while LOOM leads on MMLU and BBH; adding validation gradients to LOOM is best overall.Source: §5.5 · p. 8, Figure 4 · p. 8

6 Conclusion

The conclusion identifies self-inclusion in validation-free gradient matching as the source of optimism and noisy-example preference. It presents leave-one-out matching as an unbiased alternative and reports LOOM’s improvements in the evaluated settings.

Conclusion

  • In-sample matching’s self-inclusion puts optimism on the Gram-matrix diagonal, favoring noisy examples and preventing the objective from beating the full batch.Removing the diagonal makes the objective an unbiased estimate of update error.Source: §6 · p. 8
  • The leave-one-out objective weights examples by signal-to-noise ratio, and a subset beats the whole batch when per-example SNR is sufficiently heterogeneous.The conclusion states this heterogeneity condition holds in LLM fine-tuning.Source: §6 · p. 8
  • LOOM uses one backward pass with 5% overhead and improves on full-batch training, in-sample matchers, and validation-guided selectors.The conclusion says the correction also applies to selectors whose targets are estimated from selected data, including offline coresets and data attribution.Source: §6 · p. 8

Limitations

LOOM estimates the training-distribution gradient, so it may not capture a different target task’s information. Its analysis and evaluation also have specific data, noise, and model-scale boundaries.

Limitations

  • When the target task differs from the training pool, validation gradients can carry information LOOM’s batch-based target does not.In pooled data with rare sources, leave-one-out targets favor the majority unless computed per source.Source: Limitations · p. 9
  • The analysis assumes i.i.d. candidates, a metric fixed within each step, and finite noise variance; duplicated examples may require deduplication or grouped leave-one-out.Source: Limitations · p. 9
  • Experiments cover SFT models up to 8B parameters; pretraining and reinforcement-learning post-training remain untested.Full fine-tuning of large models also requires ghost inner products or sketches because per-example gradients are cheap for LoRA.Source: Limitations · p. 9

Ethics Statement

The appendix situates LOOM among selection and gradient-matching work, then gives proofs connecting leave-one-out targets to unbiased risk estimation. It also shows why self-coverage can favor noisy examples and derives the SNR-based optimum.

Ethics Statement

  • LOOM may down-weight rare but legitimate content, including minority dialects or uncommon viewpoints, so practitioners should audit selected data.A stratified variant limits this effect for known sources.Source: Ethics Statement · p. 9

A Extended Related Work

  • Unlike prior gradient-matching methods that target an in-sample gradient, LOOM estimates the population gradient from the batch without its self-term.The appendix connects its leave-one-out target to cross-fitting and U-statistics.Source: A Extended Related Work · pp. 15–16

B.3 Proof of Proposition 2

  • The leave-one-out mean of cross-example Gram entries estimates alignment with the population gradient without bias.The risk-optimal weights are proportional to per-example signal-to-noise ratio, with shrinkage as batch signal weakens; the objective is submodular in the SNR set function.Source: §B.1 · p. 16, §B.3 · p. 17

B.1 Proof of Proposition 1

  • In-sample matching’s Gram-diagonal covariance penalty credits examples for their own gradient noise, favoring noisy examples and making the full batch optimal for that objective.Source: §B.1 · pp. 16–17

B.2 Facility Location and Self-Coverage

  • In-sample facility location lets each selected point cover itself, giving outliers with large noise energy high marginal gain.Leave-one-out facility location is monotone submodular and greedy selection attains (1 − 1/e) of the optimum, but its distance surrogate remains below the risk-based objective because it discards magnitude information.Source: §B.2 · p. 17

B.3 Proof of Proposition 2

  • The deployed Gram-matrix objective agrees with Proposition 2’s population objective under the signal-plus-noise model.Source: §B.3 · p. 17

B.4 Proof of Theorem 1

These proofs establish when heterogeneous SNR makes subset selection preferable, derive a greedy guarantee and convergence bound, and analyze estimation accuracy. Numerical checks support the analytic results under their stated settings.

B.4 Proof of Theorem 1

  • For log-normal per-example SNR, half a batch beats the whole batch exactly when its log-standard deviation exceeds σ⋆≈0.58.Real SFT batches have estimated log-SNR standard deviations from 0.62 on ScienceQA to 1.05 on MMLU, above the break-even point.Source: §B.4 · p. 18

B.5 Proof of Proposition 3

  • Greedy selection has a (1 − e^-γ) guarantee from the submodularity ratio, but the worst-case bound is loose.On exhaustively searchable batches, it averages 99.99% of the optimum; the authors do not claim a guarantee for the nonnegative variant.Source: §B.5 · p. 18

B.6 Proof of Theorem 2

  • The convergence argument applies pathwise to any selector, with the selector entering the bound only through its average risk in the preconditioner metric.A metric fixed before the batch is needed to interpret the leave-one-out estimate as unbiased risk.Source: §B.6 · p. 18

B.7 Uniform Accuracy of the Leave-One-Out Objective

  • Under the stated conditional sub-Gaussian assumption, the proof controls leave-one-out objective error uniformly over simplex weights, including data-dependent choices.This uniform control addresses LOOM choosing its weights using the same Gram matrix.Source: §B.7 · p. 19

B.8 Estimating Signal and SNR from the Gram Matrix

  • The Gram-matrix estimates of signal and per-example SNR converge in probability as batch size grows; on synthetic batches with B=128, estimated and true SNR have Spearman correlation 0.94.Source: §B.8 · p. 19

B.9 Relation to GREATS

  • GREATS uses an external validation target with fixed equal weights, whereas LOOM uses leave-one-out targets and optimizes the weights.For an in-distribution validation set, both targets are unbiased, but their variance depends on validation size versus B − 1; an out-of-pool target may carry information the batch lacks.Source: §B.9 · p. 19

C Numerical Verification of the Theory

  • Table 3’s synthetic checks find the in-sample risk estimate too low by 11.5%, while the leave-one-out estimate is unbiased.The covariance-penalty prediction matches within 0.01%, and Monte Carlo risk ratios at half-batch size agree with Theorem 1 across three heterogeneity levels.Source: C Numerical Verification of the Theory · pp. 19–20, Table 3 · p. 20

D Implementation Details

LOOM forms Adam-metric per-example gradients and their Gram matrix during the backward pass, then uses greedy selection for a weighted update. The implementation also describes distributed, ghost-gradient, and forward-only variants.

D Implementation Details

  • Per-example LoRA gradients are formed in backward hooks with a batched matrix product costing the same FLOPs as the usual weight gradient, then stored in bfloat16.The per-example loss is a token mean over response tokens.Source: D Implementation Details · p. 20
  • Gradients are scaled by AdamW’s bias-corrected second moment from the previous step, making the metric fixed before the batch is drawn.The Gram matrix is accumulated in float32; the first-step metric is the identity.Source: D Implementation Details · p. 20
  • Greedy selection updates an incremental Cholesky factor, while distributed training exchanges gradient chunks and all-reduces the Gram matrix without a second backward pass.LOOM-fast sketches LM-head gradients in a forward-only pass and reaches 1.17× Regular throughput, but requires forwarding and back-propagating the selected half for the update.Source: D Implementation Details · p. 21

E Experimental Setup

The experiments use four SFT tasks, shared evaluation protocols, and common candidate batches to compare online selectors. LOOM’s hyperparameter is fixed across tasks and backbones, and the study reports reproducibility and compute details.

Tasks and evaluation

  • The study trains on MMLU, ScienceQA, GSM8K, and CodeAlpaca-20k, with zero-shot evaluation and means over four seeds.ScienceQA uses textual context and options; prompts follow the cited training recipe, with loss on response tokens only.Source: E Experimental Setup · pp. 21–22
  • Evaluation uses greedy decoding with at most 512 new tokens, task-specific answer parsing, and pass@1 for HumanEval.Standard deviations over four seeds are reported separately.Source: E Experimental Setup · p. 21

Hyperparameters

  • LOOM’s single hyperparameter, λ̄, was selected on a held-out MMLU auxiliary split with Qwen2.5-1.5B and then fixed across tasks and backbones.The passage reports a broad optimum for λ̄.Source: E Experimental Setup · p. 21, Table 5 · p. 22

Baselines

  • All selectors receive the same candidate batches in the same order and retain 64 examples, with baseline-specific configurations and resources.GREATS uses five development examples per task, while OPUS uses 256 held-out training examples for its proxy direction.Source: E Experimental Setup · p. 22

Mechanism experiments

  • Mechanism experiments assess reference-gradient alignment, SNR heterogeneity, and label-noise selection using fresh candidates and corrupted-response settings.The noise protocol replaces 20% or 40% of responses with responses from other randomly chosen examples of the same task.Source: E Experimental Setup · p. 22

Compute

  • A Llama-3.1-8B MMLU run takes 2.0 hours for Regular and 2.1 hours for LOOM on 8×A100, excluding evaluation.The full study used about 1.1 × 10^4 A100-hours.Source: E Experimental Setup · p. 22

F Additional Results

Additional experiments extend the comparisons across backbones and settings, and examine variability, label noise, selection behavior, and compute. Results also identify cases where validation information or source-aware targets change the comparison.

Seed variability

  • Across four seeds, LOOM has the lowest standard deviation on five of eight columns; its margins over the strongest baseline exceed the larger of the two standard deviations on MMLU, ScienceQA, and HumanEval.Across five further backbones, LOOM is 1.7–1.9 points above Regular and 0.7–0.8 above UDS.Source: F Additional Results · pp. 22–23, Table 6 · p. 24

Further backbones

  • LOOM is 1.7–1.9 points above Regular and 0.7–0.8 above UDS on five further backbones.The reported pattern holds across the additional backbone scales summarized in Table 7.Source: F Additional Results · p. 23, Table 7 · p. 24

Label noise and selection behavior

  • At 20% corrupted responses, LOOM is protected by leave-one-out targets without clean data, while in-sample matchers fare worse than Random.LOOM’s weights rise with loss through the seventh decile, then drop sharply in the top decile, where corrupted and malformed examples concentrate.Source: F Additional Results · p. 23, Table 9 · p. 25

Diagnostics and cost

  • Selection adds 499 ms to a 9.2 s step, while the Euclidean metric concentrates two thirds of tr(K) in the top eight layers and Adam spreads it across depth.LOOM’s weights scale with estimated SNR at slope 0.97 on a log–log plot; Figure 5 shows the per-task training curves, with early saturation and a small final gap on ScienceQA.Source: F Additional Results · p. 23, Figure 6 · p. 24, Figure 5 · p. 24

GREATS protocol

  • In the mixed-pool GREATS protocol, GREATS leads TydiQA by 1.5–1.8 points, LOOM leads MMLU and BBH, and their average scores are within 0.2 points.LOOM+val gives the best average on both backbones in this setting.Source: F Additional Results · p. 23

Selection ratios

  • At task-specific selection ratios, LOOM averages 71.8 versus UDS at 71.1 on Qwen2.5-7B, leading by 1.4 on HumanEval and 0.7 on MMLU at 12.5%.The tested ratios are 12.5% for MMLU and 25% for GSM8K.Source: F Additional Results · p. 23

Pooled training sets

  • On pooled training sets, plain LOOM under-weights GSM8K and loses on GSM8K and HumanEval, while source-specific targets in proportion to source shares give the best average.GSM8K makes up 5.3% and CodeAlpaca 14.3% of the pooled data.Source: F Additional Results · pp. 23–24

Full fine-tuning and transfer

  • With full fine-tuning of Qwen2.5-1.5B, LOOM improves on Regular by 1.9 points; it also transfers best to MATH500 and SVAMP from GSM8K-trained Llama-3.1-8B models.UDS is slightly ahead in-distribution in the transfer comparison.Source: F Additional Results · p. 24

G Discussion

The discussion explains why an out-of-sample target can outperform full-batch training, distinguishes online selection from data-pruning coresets, and outlines settings where LOOM’s assumptions or targets need adjustment.

Ablations

  • The in-sample target scores 53.74 average versus 56.21 for 2-fold cross-fitting; the factored Adam metric scores 56.44 versus 55.71 for the Euclidean metric.The discussion’s ablation table reports these variants and averages.Source: G Discussion · p. 26, Table 13 · p. 26

Why validation-free selection can beat the full batch

  • When example noise levels differ, inverse-variance weighting can beat the sample mean, and shrinkage can help in high dimension; leave-one-out targets make these gains accessible without known noise levels.Theorem 1 gives the condition, and the paper states that LLM fine-tuning satisfies it with a wide margin.Source: G Discussion · p. 26

When LOOM is not the right tool

  • Under training-to-target distribution shift, LOOM needs a validation signal; duplicates require grouped leave-one-out or deduplication, and rare sources require stratified targets.When Hρ ≤ 1/ρ, selection cannot beat the whole batch, though reweighting with b=B is not worse in expectation.Source: G Discussion · p. 26

Targets beyond the current batch

  • Leave-one-out target variance falls as 1/B, but gradient buffers across steps can be stale; the momentum target is worse than the fresh leave-one-out target.The proposed buffer would need age-based down-weighting or correction for parameter drift.Source: G Discussion · p. 26

Coresets and broader impact

  • Online selectors seeking to outperform full-data training need an out-of-sample target, unlike pruning coresets that correctly approximate the full-data gradient.LOOM favors content agreeing with the batch majority, so data owners should audit selected data when fairness across sources matters.Source: G Discussion · pp. 26–27

Practical recommendations

  • The recommendations are to use λ̄ in [0.03, 0.3], prefer larger candidate batches when memory permits, monitor H1/2, and stratify pooled-data targets by source.When whole-batch backpropagation is too expensive, LOOM-fast retains most of the gain at throughput above Regular.Source: G Discussion · p. 27

Related papers

Loading 2610.00436v1…