Source-linked AI summary
Scalable, Transferable Meta-network for Data Selection Requires a Different Loss (and Why the Obvious Choice is Problematic)
Zilin Du, Bowen Yang, Boyang Albert Li
TL;DR
Training-data selection for large language models needs to handle heterogeneous corpora while transferring to unseen examples, but existing MTS approaches struggle with that transfer. TESS trains a selection network on sample-level loss-difference pseudo-labels, and the paper reports transfer across datasets, corpus scales, and model sizes.
Problem
Existing MTS methods face a trade-off between fine-grained example valuation and transfer to unseen data, limiting their fit for large heterogeneous corpora.
Source: §1 · pp. 1–2Method
TESS trains a data-selection network to match sample-level pseudo-labels derived from the loss difference between training-only and validation-guided models.
Source: §1 · p. 2, §3.3 · p. 4Results
TESS shows strong transfer across datasets, corpus scales, and model sizes, including subset-to-full-corpus and smaller-to-larger-model settings.
Source: §1 · p. 2, §6 · p. 10Takeaways & Limitations
The trained selector can score unseen examples in a candidate pool without further optimization, and TESS reduces peak memory footprint by more than 20%.
Source: §1 · p. 2The analysis strictly applies to sequential training in TESS; the observed dynamics under joint ScaleBiO training are qualitatively similar.
Source: §4.1 · p. 5
Abstract
from arXiv · showhide
Data selection is critical for training large language models on massive and heterogeneous corpora. Meta-learning for Training-data Selection offers a principled alternative to heuristic scoring by learning data weights from a target validation objective, but existing methods face a trade-off between fine-grained valuation and transferability to unseen data. A natural solution is to replace per-sample weights with a selection network. However, we find that directly incorporating such a network into existing MTS objectives leads to unstable optimization and poor generalization, caused by weight suppression and persistent reliance on easy-to-learn features. To address these issues, we propose Transferable Example Scoring and Selection (TESS), a scalable data-selection framework built on a Pointwise Value Matching objective (PVM). Experiments on LLM safety and targeted instruction tuning demonstrate strong transfer across datasets, from subsets to full corpora, and from smaller to larger models.
1 INTRODUCTION
MTS learns training-data values from their effects on a validation objective, but scaling it to LLMs leaves both computational cost and transfer to unseen data unresolved. TESS uses sample-level pseudo-labels to train a selection network and reports transfer across datasets, corpus scales, and model sizes.
1 INTRODUCTION
- MTS learns training-data values from their effects on validation loss, offering an alternative to manually designed scores that do not dynamically adapt to new scenarios.The introduction presents data selection as important because heterogeneous, massive corpora often contain redundant and low-quality examples.Source: §1 · pp. 1–2
- Exact MTS hypergradients are costly for LLMs; even one-step unrolling with finite differences requires three backward operations per weight-gradient step.ScaleBiO reduces backward calls to two, while SEAL reduces them to one but incurs discernible performance losses.Source: §1 · p. 2
- Per-sample weights cannot generalize to unseen points, while domain-level weights miss within-domain variation and cannot transfer to new domains.The cited methods include separate weights for each point and for each domain.Source: §1 · p. 2
- Directly training a meta-network with the existing objective suppresses most predicted weights and reinforces easy features, impairing sample discrimination and encouraging shortcut learning.The paper identifies these as two limitations of applying the existing MTS objective to a shared selection network.Source: §1 · p. 2
- TESS uses sample-level pseudo-labels to train a selection network and reports transfer across datasets, from subsets to full corpora, and from smaller to larger models.It also reduces peak memory footprint by more than 20%.Source: §1 · p. 2
2 RELATED WORK
Prior data-selection work addresses LLM-scale optimization and transfer with per-example weights, domain weights, or meta-networks, but none resolves the combination of fine-grained scoring and transfer to unseen data. The paper analyzes meta-network failures and proposes PVM as a practical alternative.
Meta-Learning Data Selection for LLMs
- Traditional methods differentiate through model updates, while ScaleBiO and SEAL target model-scale efficiency; generalizing from small training sets to unseen data could also improve scalability to large datasets.The related-work discussion says unseen-data generalization remains difficult.Source: §2 · p. 3
- Per-example methods learn no weights for new points, domain weights are too coarse for within-domain differences, and prior meta-networks had not reported generalization to unseen data.DataRater trains a separate scorer for each dataset, suggesting cross-dataset generalization remains challenging even with large training data.Source: §2 · p. 3
Shortcut Learning
- ScaleBiO reinforces easy-to-learn features, whereas PVM promotes harder, more transferable features; easy cues can become shortcuts when their correlations weaken or reverse under distribution shifts.The passage distinguishes feature predictivity from availability, or ease of extraction.Source: §2 · p. 3
3 METHODOLOGY
The methodology formulates selection as learning a network that scores candidate examples to improve a target validation objective. TESS replaces direct optimization of that network under ScaleBiO with loss-difference pseudo-labels and a sequential training procedure.
3.1 PROBLEM FORMULATION
- The formulation uses a selection network to predict positive weights for training examples and score candidates in a pool that may overlap the training set or be disjoint.The network is optimized so that training with its predicted weights minimizes validation loss.Source: §3.1 · p. 3
3.2 PRELIMINARIES ON SCALEBIO
- ScaleBiO reformulates the bilevel constraint as equality of training losses and derives a penalty-based minimax objective, but simultaneous training of two LLMs and weights has substantial memory overhead.SEAL removes the auxiliary LLM at an accuracy penalty; ScaleBiO selects points that benefit from validation-guided training relative to training-only loss.Source: §3.2 · p. 4
3.3 TRANSFERABLE EXAMPLE SCORING AND SELECTION (TESS)
- Replacing independent weights with a selection network yields the SBO objective, but its direct optimization produces unstable training and poor generalization.The paper introduces Pointwise Value Matching to address these empirical problems.Source: §3.3 · p. 4
- PVM uses each training example’s loss difference between a training-only model and a validation-guided model as a utility pseudo-label for the selector.A larger difference indicates greater loss reduction under the validation-guided model and suggests greater relevance to the target validation objective.Source: §3.3 · p. 4
- TESS trains the models and selector sequentially, reducing peak memory, permitting independent optimization schedules, and allowing the trained network to score unseen candidates without further optimization.The LLMs are frozen before pseudo-label computation and selector training.Source: §3.3 · pp. 4–5
4 WHY THE SCALEBIO LOSS PERFORMS POORLY
The ScaleBiO objective can suppress most example weights and erode sample discrimination, while also favoring easy-to-learn features that may act as shortcuts. The section contrasts these dynamics with PVM and examines their effects on unseen-data selection.
4.1 DATA WEIGHT SUPPRESSION AND LOSS OF DISCRIMINATION
- ScaleBiO’s competitive updates can concentrate nearly all normalized weight on one example, driving the other weights toward zero and weakening sample discrimination.With a unique example having the largest pseudo-label, its normalized weight tends to 1 while the others tend to 0; for a positive lower bound, the remaining mass is O(t^-1/3).Source: §4.1 · p. 5
- As ScaleBiO training progresses, selection accuracy can deteriorate even on seen data; the interval when weights separate higher-pseudo-label examples is transient, making early stopping trial-and-error dependent.The passage reports that seen and unseen accuracy first improve and eventually fall, while Figure 1 illustrates the transient separation period.Source: §4.1 · pp. 5–6
- PVM maintains nonvanishing weight quantiles and reaches 98.58% seen accuracy and 97.96% unseen accuracy, while ScaleBiO accuracy eventually falls to around 60%.The reported PVM accuracies are for the controlled safe–unsafe classification experiment.Source: Figure 2 · p. 6, §4.1 · p. 6
4.2 EASY-FEATURE DOMINANCE AND SHORTCUT LEARNING
- ScaleBiO continues favoring easy features, whereas PVM can shift learning toward harder features that explain remaining prediction error, a distinction relevant when easy features are unreliable on unseen data.The analysis states that PVM’s easy–hard feature gradient gap decreases exponentially, while ScaleBiO’s remains nondecreasing.Source: §4.2 · pp. 6–7
- With shortcut features, PVM improves unseen accuracy and reduces the seen–unseen gap from early to late training, while ScaleBiO has lower unseen accuracy and larger gaps over training.All variants show larger average gaps under strong shortcuts.Source: Table 1 · p. 7, §4.2 · p. 7
5 EXPERIMENTS
TESS is evaluated on harmful-data selection and targeted instruction tuning, including transfer across datasets, corpus sizes, and model scales. Across these settings, results favor TESS over the reported baselines, while the loss ablation highlights the role of PVM.
5.1 LLM SAFETY
- 16.27% over SBOw and 11.04% over Bi-Anchor, TESS’s reported Training-setting margins, accompany its best result in 14 of 18 comparisons.Bi-Anchor remains stronger on Llama3-Alpaca; Table 2 reports ASR for Alpaca and Dolly, ranked separately within Training and Generalization settings.Source: Table 2 · p. 8, §5.1 · p. 8
- 22.76 and 24.58 points in average ASR over SEALϕ and SBOϕ, respectively, accompany TESS wins in all 18 Generalization evaluations and a 5.13% drop from its Training performance.The Generalization setting trains on one candidate dataset and applies the selector to the other.Source: §5.1 · pp. 7–8
- 5.65–6.60× lower training time accompanied retention of 77.53% and 83.67% of large-model performance in Training and Generalization, respectively, when selection transferred from Qwen2.5-0.5B to Qwen2.5-7B.The passage identifies Table 5 in Appendix F as the source of these results.Source: §5.1 · p. 8
5.2 TARGETED INSTRUCTION TUNING
- 4.92 points overall over Random, TESS achieved the best average on both targeted tasks and ranked first in five of six settings.Table 4 reports Llama-2-7B performance after fine-tuning on selected Tulu V2 examples; the comparison uses a 50K training subset and full selection pool for TESS, SEALϕ, and SBOϕ.Source: Table 4 · p. 9, §5.2 · pp. 9–10
5.1 LLM SAFETY
- 21.60% and 16.24% higher Training-setting ASR than SBO Loss and KL Loss, respectively, and 18.15% higher than Pseudo-Label Ranking were reported for PVM.In Generalization, PVM exceeded SBO Loss and KL Loss by 24.61% and 25.07% on average, with a 5.65-point ASR decrease from Training; the ablation is on Qwen2.5-7B-Instruct.Source: §5.1 · p. 9, Table 3 · p. 9
5.2 TARGETED INSTRUCTION TUNING
- 13.4 hours for TESS versus 59.8 hours for SBOw and 39.6 hours for SEALw; those sample-specific-weight baselines train on the full 200K pool, while TESS trains on 50K.SBOw also uses 188.2 GB versus 156.8 GB for TESS.Source: §5.2 · pp. 9–10
6 CONCLUSION
The conclusion presents TESS as a transferable selector trained from sample-level pseudo-labels and identifies weight suppression and easy-feature preference in the ScaleBiO objective as problems it addresses. The appendix passages define the per-sample LLM loss and describe dependencies in the pseudo-label weighting formulation.
6 CONCLUSION
- TESS learns a selector from sample-level pseudo-labels and is reported to transfer across datasets, corpus scales, and model sizes.The conclusion says it addresses weight suppression and easy-feature preference associated with the ScaleBiO objective using PVM.Source: §6 · p. 10
A GRADIENT OF THE SCALEBIO META-NETWORK OBJECTIVE
- ℓ(θ; x_i) denotes the token-averaged cross-entropy loss of an LLM parameterized by θ on sample x_i.The passage introduces this per-sample loss in the appendix discussion of the ScaleBiO meta-network objective.Source: A GRADIENT OF THE SCALEBIO… · p. 16
- The appendix states that per-sample pseudo-labels are constant with respect to ϕ and data weights, while their weighted average depends on the weights through p_j.The weighted-average passage introduces the pseudo-labels and their weighted average.Source: A GRADIENT OF THE SCALEBIO… · p. 16
B PROOF OF THEOREM 1
Under the proof’s assumptions, the projected weight dynamics concentrate normalized weights on the unique maximal pseudo-label. Nonmaximal weights reach the imposed floor in finite time, yielding finite-time concentration when that floor is zero.
B PROOF OF THEOREM 1
- The proof assumes nonnegative weights with initial values above a floor δ0, and a unique index K with the largest fixed pseudo-label.The weights have no imposed upper bound, and the dynamics use projected gradient descent that prevents weights from falling below δ0.Source: B PROOF OF THEOREM 1 · pp. 16–17
- The weighted mean of pseudo-labels converges to the largest pseudo-label, and the normalized weights concentrate on its unique maximizer K.The proof establishes positivity and growth bounds for the normalization before deriving convergence of the weighted mean.Source: B PROOF OF THEOREM 1 · pp. 17–18, §3 · p. 18
3. The normalized weights concentrate on the unique maximizer. For i̸ = K, define
- Every nonmaximal weight reaches the floor δ0 in finite time and remains there thereafter.The proof first establishes finite-time arrival at the floor, then applies the projected dynamics to keep each such coordinate there.Source: §3 · p. 18
- When δ0 = 0, only the maximal weight remains positive after the nonmaximal weights hit the floor, giving finite-time convergence of normalized weights.The proof states p_K(t) = 1 and p_i(t) = 0 for every i ≠ K after T0.Source: §3 · p. 18
C PROOF OF THEOREM 2
The proof compares feature-gradient dynamics for ScaleBiO and PVM in a fixed affine surrogate, with easy and hard features distinguished by their covariance with the target label. It reports a persistent ScaleBiO gradient gap and exponential contraction of the PVM gap under the stated comparison conditions.
C PROOF OF THEOREM 2
- The comparison defines the easy feature e by its stronger covariance with the target label than the hard feature h, using a fixed affine meta-network surrogate.The analysis assumes standardized features and fixed pseudo-labels; it considers training steps for which all meta-network outputs remain positive.Source: C PROOF OF THEOREM 2 · pp. 19–20
- ScaleBiO continues learning easy features more than hard features, and its gradient gap does not decrease over the stated comparison interval.The proof derives the ScaleBiO gradient with respect to parameters while holding pseudo-labels fixed.Source: C PROOF OF THEOREM 2 · pp. 20–21
- PVM exponentially contracts the absolute squared-gradient gap, whereas the ScaleBiO gap remains no lower than its initial value under the comparison conditions.The PVM contraction follows the stated step-size condition.Source: C PROOF OF THEOREM 2 · pp. 21–22
D.1 DATASETS
The appendix evaluates safety selection on Alpaca and Dolly and targeted instruction tuning on Tulu V2, then describes implementation and meta-network inputs.
D.1 DATASETS
- Safety selection uses Alpaca and Dolly, with evaluation on DirectHarm4, HarmBench, and HEx-PHI.The three evaluation sets cover harmful instructions, harmful behaviors, and prohibited-use categories, respectively.Source: §D.1 · p. 23
- Targeted instruction tuning uses a roughly 200K-example Tulu V2 corpus, with the meta-network trained on a random 50K subset and applied to the full pool.GSM8K is evaluated using a shared validation set of 1,000 examples; CodeX uses 16 query examples and 148 test examples.Source: §D.1 · p. 23
D.2 IMPLEMENTATION DETAILS
- The meta-network is a three-layer MLP with sigmoid output, optimized with AdamW at 10^-3; LLM learning rates and the trade-off coefficient are tuned over supplied grids.The two LLMs use LISA for up to three epochs; α is tuned over {0.01, 0.1, 1.0}, and LLM learning rates over {5 × 10^-5, 10^-4}.Source: §D.2 · p. 23
- Downstream safety fine-tuning applies LoRA to query and key projections, while targeted instruction tuning fully fine-tunes Llama-2-7B.Safety fine-tuning uses rank 8 for three epochs; targeted tuning runs two epochs with a maximum sequence length of 2,048 tokens.Source: §D.2 · p. 23
D.3 INPUT FEATURES TO THE META-NETWORK
- Meta-network inputs combine representation- and gradient-space distances with sample loss and gradient norm into a 10-dimensional feature vector.Distances are measured to each of the three nearest validation examples and to the mean validation embedding in each space.Source: §D.3 · pp. 23–24
- For safety selection, the feature vector also includes a binary indicator for numbers or keywords such as “list,” “convert,” and “calculate.”The authors motivate this feature by findings that some benign list-format and mathematical examples can degrade safety alignment after fine-tuning.Source: §D.3 · p. 24
D.4 EVALUATION DETAILS
Evaluation uses task-specific metrics for safety, math, and code, and the ablation’s KL variant matches network outputs to loss-gap-reweighted pseudo-labels.
D.4 EVALUATION DETAILS
- Safety performance is reported as attack success rate (ASR), the percentage of evaluated responses classified as unsafe; lower is better.Responses are generated with greedy decoding and classified using Llama Guard, also with greedy decoding.Source: §D.4 · p. 24
- GSM8K performance is exact-match accuracy of extracted final numerical answers on 1,319 test problems.The eight chain-of-thought demonstrations serve as individual query examples for data selection and are all included in each evaluation prompt.Source: §D.4 · p. 24
- CodeX performance is pass@10 on 148 heldout problems, estimating the chance that at least one of ten sampled completions passes all unit tests.Completions are sampled at temperature 0.8.Source: §D.4 · p. 24
D.5 KL LOSS IN ABLATION STUDY
- The KL ablation matches normalized meta-network outputs p to a distribution q reweighted by loss-gap pseudo-labels.Source: §D.5 · p. 24
E TOY EXPERIMENT IN SEC.4
The toy experiment tests selection under shortcut reversal using seen and unseen splits, measuring unsafe-example identification and generalization across the reversal.
E TOY EXPERIMENT IN SEC.4
- The experiment builds disjoint seen and unseen splits, each with 2,575 safe GSM8K examples and 1,125 unsafe BeaverTails examples.Validation uses 100 safety-related examples pairing harmful prompts with refusal responses.Source: E TOY EXPERIMENT IN SEC.4 · p. 24
- The meta-network trains on the seen set for 30 epochs with AdamW and a cosine schedule decaying the learning rate to zero.Training uses batch size 256 and an initial learning rate of 5 × 10^-4.Source: E TOY EXPERIMENT IN SEC.4 · p. 24
- For each split, the 1,125 lowest-weight examples are classified as unsafe, and identification accuracy equals both precision and recall.The equality follows because the predicted unsafe set and ground-truth unsafe set have the same size.Source: E TOY EXPERIMENT IN SEC.4 · p. 24
- Three synthetic features encode a class-dependent shortcut whose means reverse on unseen data, with separate strong- and weak-shortcut settings.The settings use a = 0.5 and a = 0.3, respectively, with independently sampled noise of standard deviation 0.3.Source: E TOY EXPERIMENT IN SEC.4 · p. 24
- Shortcut associations are summarized by mean absolute covariance and mean absolute Pearson correlation before standardization.The reported statistics summarize each of the three shortcut features’ association with the class label.Source: E TOY EXPERIMENT IN SEC.4 · p. 25
- Generalization is summarized by the seen–unseen accuracy gap, where a smaller gap at comparable seen accuracy indicates better performance under shortcut reversal.Source: E TOY EXPERIMENT IN SEC.4 · p. 25
F SMALL-TO-LARGE MODEL GENERALIZATION
TESS transfers data selection from a smaller to a larger model with lower training cost, and the conclusion reports generalization gains across safety and targeted instruction tuning.
F SMALL-TO-LARGE MODEL GENERALIZATION
- 5.65–6.60× training speedup retains 77.53% and 83.67% of large-model performance in the Training and Generalization settings, respectively.The meta-network is trained with Qwen2.5-0.5B-Instruct and selects data for Qwen2.5-7B-Instruct.Source: F SMALL-TO-LARGE MODEL GENERALIZATION · p. 25
G CONCLUSION
- TESS uses sample-level pseudo-labels to train a data-selection network that generalizes to unseen data.Source: G CONCLUSION · p. 25
- TESS outperforms the strongest learnable-weighting baselines by 22.75 points on LLM safety and 3.99 points on targeted instruction tuning in generalization.Source: G CONCLUSION · p. 25