The choice of post-training data for large language models substantially affects downstream performance. Gradient-based data selection is a popular approach that ranks training data by how well their gradients align with those of a small validation set. However, ranking with full-parameter gradients requires an expensive backward pass on every sample, making computation intractable for large candidate pools. This raises a natural question: can we approximate full-gradient features at a fraction of the cost? Conveniently, we find that output-layer gradients suffice for effective data selection, yet require only the cheaper forward pass. We implement this as LESSER, a drop-in wrapper for selection methods that reduces the feature-extraction FLOP cost by 9.7× for SFT and 3.0× for RL benchmarks, while tracking full-gradient performance on downstream tasks. Empirically, we find that even when output-layer and full gradients rank individual samples differently, they select batches with aligned gradients.
Figures & tables
Figure 1: Efficient data selection with output-layer gradients. Given a candidate training pool and a small query set representative of downstream tasks, LESSER extracts output-layer gradient features to rank and select training examples. These features approximate the downstream behavior of full-gradient methods like LESS while substantially reducing feature-extraction FLOPs.
Figure 2: Data selection with output-layer gradient features in SFT and RL. Pool and query representations determine similarity scores, while the selection rule produces the training subset. LESSER replaces the gradient representation within this pipeline.
Figure 3: LESSER closely tracks LESS on downstream performance. Panels show Llama-2-7B; values are means over three training seeds, and error bars show one sample standard deviation. Across all 60 model–task–budget cells, the absolute score gap of LESSER averages 1.3 points against LESS , 2.2 against RDS+ , and 2.6 against Random.
Figure 4: Higher-similarity training examples lead to lower query loss. Query loss after SFT on Llama-2-7B, stratified by example score. For both LESS and LESSER , examples with higher query similarity generally yield lower query loss after fine-tuning.
Figure 5
SFT Selector
Time (h)
Speedup
LESSER
3.16
17.8×
LESS
56.16
1.0×
Table 1: Output-layer gradients are much cheaper to extract. Cost analysis on one A100 80GB for SFT feature extraction and RL scoring.
Figure 7: LESSER -selected batches induce gradients aligned with those from LESS . (a) Centered batch-gradient alignment between LESS batches and batches selected by each method, at the base model. Dots denote model–task pairs, and bars show medians. (b) Median share across model–task pairs of the query-loss decrease reached by each training step.
Appendix figures & tables17 assets
Supplementary material from the paper’s appendix.
Appendix
Task
Domain
Val. queries
Metric
TyDiQA ( Clark et al., 2020 )
multilingual QA
9
F1
MMLU-Pro ( Wang et al., 2024 )
knowledge / reasoning
70
accuracy
GSM8K ( Cobbe et al., 2021 )
grade-school math word problems
8
exact match
Codex ( Chen et al., 2021 )
Python code generation
16
pass@ 10
BBH ( Suzgun et al., 2023 )
challenging BIG-Bench reasoning
81
exact match
Appendix
Table 2: Target tasks. “Val. queries” is the query-set size used to specify each target; metrics and splits follow Nayak et al. (2026) , except that MMLU-Pro is evaluated five-shot.
Model
Task
LESS
Random
Llama-3.2-3B
TyDiQA
66.4 / 66.1
53.7 / 50.7
Llama-3.2-3B
GSM8K
34.2 / 34.0
29.2 / 27.6
Llama-3.2-3B
Codex
56.3 / 55.4
59.2 / 59.9
Llama-3.2-3B
BBH
47.6 / 47.2
47.1 / 47.1
Qwen3-4B-Base
TyDiQA
71.6 / 71.4
–
OLMo3-7B
TyDiQA
62.5 / 62.5
–
Appendix
Table 3: Pipeline replication, reported value / our rerun.
Temperature 0.7 ; maximum response length 3,072 tokens
Appendix
Table 4: RL selection, training, and evaluation settings.
Quantity
Symbol
Value
Model parameters
P
6,738,415,616
LoRA parameters
PLoRA
134,217,728
Pool examples
N
197,196
Processed tokens
T
86,343,509
Response tokens
Tr
25,430,302
Vocabulary size
∣V∣
32,000
Appendix
Table 5: Quantities used for the Llama-2-7B SFT feature-extraction FLOP calculation.
Extraction share / %
Budget k
LESSER
Four-checkpoint LESS
Extraction + training FLOP ratio
1,000
97.0
99.7
9.47×
5,000
86.8
98.5
8.58×
10,000
76.7
97.0
7.69×
Appendix
Table 6: SFT computation costs for one target and training run.
Feature
FLOPs / scored token
Stored feature
LESSER
3.11×109
32.8 kB
Full gradients
9.26×109
6.2 GB
Appendix
Table 7: Computation and per-problem storage for RL scoring.
Setting
Full gradients
LESSER
Speedup
Peak memory, GiB full / LESSER
GSM8K, Llama-3.2-1B, 14 teachers
6.49 h
1.65 h
3.9×
19.6 / 12.3
GSM8K, OLMo-2-1B, 10 teachers
5.34 h
0.99 h
5.4×
17.0 / 7.0
MATH, Llama-3.2-3B, 10 teachers
14.68 h
3.09 h
4.8×
39.8 / 24.2
Appendix
Table 8: GRACE feature-extraction cost on one A100 80GB. Extraction time is estimated from 32 timed responses per teacher and extrapolated to all 2,048 responses per teacher; peak memory is measured directly.
Model
Task
LESSER
LESS
RDS+
Random
Llama-3.2-3B
TyDiQA
63.3±0.5
66.4±0.2
45.7±0.4
53.7±1.6
Llama-3.2-3B
MMLU-Pro
22.1±2.3
22.8±0.5
25.3±0.8
14.9±8.9
Llama-3.2-3B
GSM8K
32.3±0.9
34.2±0.5
32.3±0.7
29.2±0.7
Llama-3.2-3B
Codex
61.3±2.4
56.3±3.0
60.6±1.7
59.2±7.8
Llama-3.2-3B
BBH
48.0±0.2
47.6±0.2
48.7±0.5
47.1±0.5
Llama-2-7B
TyDiQA
50.2±0.4
54.6±0.4
46.2±0.6
49.4±0.7
Appendix
Table 9: Held-out performance at k=1,000 , mean ± sample standard deviation over three seeds; each task uses its target metric on a 0 – 100 scale, higher is better.
Model
Task
LESSER
LESS
RDS+
Random
Llama-3.2-3B
TyDiQA
64.3±1.4
66.3±0.2
58.2±0.9
60.5±0.2
Llama-3.2-3B
MMLU-Pro
26.9±0.4
26.1±0.2
26.9±0.1
26.6±0.1
Llama-3.2-3B
GSM8K
34.6±0.1
34.5±0.6
40.0±0.6
33.3±0.9
Llama-3.2-3B
Codex
57.4±2.0
58.1±3.1
61.0±0.4
59.0±1.0
Llama-3.2-3B
BBH
48.5±0.3
49.6±0.4
49.7±0.3
48.5±0.5
Llama-2-7B
TyDiQA
51.7±0.4
52.9±0.1
50.4±0.4
51.2±0.2
Appendix
Table 10: Held-out performance at k=5,000 .
Model
Task
LESSER
LESS
RDS+
Random
Llama-3.2-3B
TyDiQA
66.0±0.2
67.2±0.1
60.3±0.1
63.6±0.4
Llama-3.2-3B
MMLU-Pro
26.5±0.2
27.1±0.4
27.3±0.1
26.9±0.3
Llama-3.2-3B
GSM8K
35.2±0.4
36.7±0.9
39.6±0.3
35.2±0.9
Llama-3.2-3B
Codex
57.4±0.7
58.3±2.1
60.1±2.0
58.3±4.3
Llama-3.2-3B
BBH
49.3±0.5
49.1±0.6
49.4±0.4
48.6±0.2
Llama-2-7B
TyDiQA
52.8±0.4
53.6±0.1
50.7±0.6
52.1±0.2
Appendix
Table 11: Held-out performance at k=10,000 .
Figure 8: Llama-3.2-3B: query loss across similarity bins. Each bin uses one training run; bin 0 is closest to the queries.
Figure 9: Qwen3-4B-Base: query loss across similarity bins. Each bin uses one training run; bin 0 is closest to the queries.
Figure 10: OLMo3-7B: the closest LESSER bin has the lowest query loss on all five tasks. Loss is not monotone across bins on MMLU-Pro. Each bin uses one training run.
Task
Student
Teachers
Spearman
Top-ranked teacher
GSM8K
Llama-3.2-1B
14
0.938
Qwen2.5-1.5B-Instruct
GSM8K
OLMo-2-1B
10
0.964
Qwen2.5-3B-Instruct
MATH
Llama-3.2-3B
10
0.976
Qwen2.5-3B-Instruct
Appendix
Table 12: Teacher-ranking agreement between LESSER and full-gradient GRACE. Spearman correlation uses the complete ranking.
Model
TyDiQA
MMLU-Pro
GSM8K
Codex
BBH
Mean Jaccard
Llama-2-7B
491 / 0.052
505 / 0.053
324 / 0.033
214 / 0.022
228 / 0.023
0.037
Llama-3.2-3B
407 / 0.042
433 / 0.045
441 / 0.046
266 / 0.027
186 / 0.019
0.036
Qwen3-4B-Base
594 / 0.063
173 / 0.018
337 / 0.035
269 / 0.028
143 / 0.015
0.032
OLMo3-7B
1,004 / 0.112
190 / 0.019
250 / 0.026
476 / 0.050
182 / 0.019
0.045
Appendix
Table 13: LESSER and LESS select substantially different examples. Each cell gives the number of shared examples and the Jaccard similarity.
Figure 11: Per-example rankings depend on gradients throughout the network. Spearman correlation with the full-gradient cosine ranking as decoder blocks are added to the output layer from last to first. Candidates are ranked by their cosine with the query gradient over the included parameters. Thin lines denote model–task pairs, colored by model, and the black line shows the median.
Figure 12: Query loss across training steps. Each panel is one model–task pair.
Data attribution methods using gradient similarity are widely used to analyze and select training data for large language models, but what gradient similarity actually measures is debated. Some interpret it as identifying task-relevant skills, while other work reports that surface form is the main factor. We resolve this debate for supervised fine-tuning examples by varying task and answer format independently. Specifically, we render benchmarks in different answer formats, such that datasets can share a task without a format or a format without a task. We find that gradient alignment follows the answer format, as benchmark pairs sharing an answer format align strongly (disattenuated cosine near 0.4), while same benchmarks rendered with different answer format classes show no alignment (near 0.0). We demonstrate that this ordering holds from the earliest pretraining checkpoints through post-training, and across model scales and families. We then analyze the released selections of LESS, a gradient-based data selection method for instruction tuning, and find that each target's selections over-represent the target's own answer format. Hence, we demonstrate that gradient-based attribution methods track format similarity more than task semantics, meaning that such methods, as well as the semantic interpretation of the gradient, should be tested on data where answer format and task vary independently for greater robustness and reliability.
Data selection methods address a critical challenge in LLM post-training: effectively leveraging scarce, high-fidelity target data alongside abundant but imperfectly aligned general training data. In this work, we move beyond the data-selection framing and introduce Dr. Post-Training (Data-Regularized Post-Training), a novel framework that reconceptualizes general training data as a data-induced regularizer that prevents overfitting to the scarce target objective, rather than serving as a pool for selection. Specifically, our framework proposes that, at each training step, we construct a feasible set of model update directions using the general training data, and project the model update direction specified by the scarce target data onto that feasible set. Standard training and existing data selection methods arise as special cases with different choices of the data-induced regularizer, and these methods correspond to different points on a bias-variance spectrum with different regularization strengths. Building on this view, we propose a family of methods offering a richer design space and more flexible bias-variance tradeoffs. For practical LLM-scale use, we introduce careful system optimizations that realize these methods with minimal overhead. Extensive experiments across SFT, RLHF, and RLVR show that our methods outperform standard training and state-of-the-art data selection baselines, with the largest gains where the general and target data are most misaligned, and system benchmarks confirm their efficiency.
Pingbang Hu, Xueshen Liu, Z. Morley Mao +1
University of Illinois Urbana–Champaign · University of Michigan
LoRA fine-tuning adapts small language models (SLMs) to heterogeneous instruction data within a low-rank update subspace, making it vulnerable to three structural problems: conflicting gradients that cancel, static data selection that cannot track evolving learning dynamics, and subspace saturation that causes later updates to overwrite useful directions. We argue that effective adaptation therefore requires controlling which data-induced gradients enter the LoRA subspace and when. We propose GRADE (GRadient-Aligned Data-centric rEcipe), a data-centric framework combining two mechanisms: a state-aware selector that continually admits samples aligned with the evolving multi-task gradient field, and a self-calibrating step-level gate that rejects updates likely to cause destructive overwrite near saturation. Across three current-generation backbones and a heterogeneous seven-dataset instruction pool, GRADE outperforms strong data-selection and PEFT-stabilization baselines in accuracy and robustness. It is the only method to improve consistently over standard LoRA on every architecture, while producing more coherent gradient trajectories and less destructive overwrite. These results show that successful SLM adaptation depends not only on which data are selected, but also on which gradients are allowed to enter and persist in the constrained update subspace.
Hongyu Cao, Yanchi Liu, Kunpeng Liu +5
Arizona State University · Clemson University · Meta AI +1