LESSER: Post-Training Data Selection with Output-Layer Gradients

Lyuxin David Zhang; Eric Wong; Surbhi Goel; Anton Xue

cs.LGOct 2, 2026

LESSER: Post-Training Data Selection with Output-Layer Gradients

Authors: Lyuxin David Zhang, Eric Wong, Surbhi Goel, Anton Xue

Organizations: University of Pennsylvania · University of Texas at Austin

Abstract

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\times$ for SFT and $3.0\times$ 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.

SFT Selector	Time (h)	Speedup
LESSER	$3.16$	$\mathbf{17.8}\times$
LESS	$56.16$	$1.0\times$

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.

Setting	Value
Candidates and selection	$5{,}120$ candidate problems; $256$ selected per round
Scoring responses	$16$ per candidate and per query problem
Training between selections	$10$ GRPO steps
Training batch and responses	$128$ problems per step; $128$ responses per problem
Optimizer	AdamW ( Loshchilov and Hutter, 2017 ) , learning rate $10^{-6}$
Sampling	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	$P_{\rm LoRA}$	$134{,}217{,}728$
Pool examples	$N$	$197{,}196$
Processed tokens	$T$	$86{,}343{,}509$
Response tokens	$T_{r}$	$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\times$
$5{,}000$	$86.8$	$98.5$	$8.58\times$
$10{,}000$	$76.7$	$97.0$	$7.69\times$

Appendix

Table 6: SFT computation costs for one target and training run.

Feature	FLOPs / scored token	Stored feature
LESSER	$3.11\times 10^{9}$	$32.8$ kB
Full gradients	$9.26\times 10^{9}$	$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\times$	$19.6$ / $12.3$
GSM8K, OLMo-2-1B, $10$ teachers	$5.34$ h	$0.99$ h	$5.4\times$	$17.0$ / $7.0$
MATH, Llama-3.2-3B, $10$ teachers	$14.68$ h	$3.09$ h	$4.8\times$	$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{\scriptstyle\pm}0.5$	$\mathbf{66.4{\scriptstyle\pm}0.2}$	$45.7{\scriptstyle\pm}0.4$	$53.7{\scriptstyle\pm}1.6$
Llama-3.2-3B	MMLU-Pro	$22.1{\scriptstyle\pm}2.3$	$22.8{\scriptstyle\pm}0.5$	$\mathbf{25.3{\scriptstyle\pm}0.8}$	$14.9{\scriptstyle\pm}8.9$
Llama-3.2-3B	GSM8K	$32.3{\scriptstyle\pm}0.9$	$\mathbf{34.2{\scriptstyle\pm}0.5}$	$32.3{\scriptstyle\pm}0.7$	$29.2{\scriptstyle\pm}0.7$
Llama-3.2-3B	Codex	$\mathbf{61.3{\scriptstyle\pm}2.4}$	$56.3{\scriptstyle\pm}3.0$	$60.6{\scriptstyle\pm}1.7$	$59.2{\scriptstyle\pm}7.8$
Llama-3.2-3B	BBH	$48.0{\scriptstyle\pm}0.2$	$47.6{\scriptstyle\pm}0.2$	$\mathbf{48.7{\scriptstyle\pm}0.5}$	$47.1{\scriptstyle\pm}0.5$
Llama-2-7B	TyDiQA	$50.2{\scriptstyle\pm}0.4$	$\mathbf{54.6{\scriptstyle\pm}0.4}$	$46.2{\scriptstyle\pm}0.6$	$49.4{\scriptstyle\pm}0.7$

Appendix

Table 9: Held-out performance at

k=1{,}000

, mean

\pm

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{\scriptstyle\pm}1.4$	$\mathbf{66.3{\scriptstyle\pm}0.2}$	$58.2{\scriptstyle\pm}0.9$	$60.5{\scriptstyle\pm}0.2$
Llama-3.2-3B	MMLU-Pro	$\mathbf{26.9{\scriptstyle\pm}0.4}$	$26.1{\scriptstyle\pm}0.2$	$\mathbf{26.9{\scriptstyle\pm}0.1}$	$26.6{\scriptstyle\pm}0.1$
Llama-3.2-3B	GSM8K	$34.6{\scriptstyle\pm}0.1$	$34.5{\scriptstyle\pm}0.6$	$\mathbf{40.0{\scriptstyle\pm}0.6}$	$33.3{\scriptstyle\pm}0.9$
Llama-3.2-3B	Codex	$57.4{\scriptstyle\pm}2.0$	$58.1{\scriptstyle\pm}3.1$	$\mathbf{61.0{\scriptstyle\pm}0.4}$	$59.0{\scriptstyle\pm}1.0$
Llama-3.2-3B	BBH	$48.5{\scriptstyle\pm}0.3$	$49.6{\scriptstyle\pm}0.4$	$\mathbf{49.7{\scriptstyle\pm}0.3}$	$48.5{\scriptstyle\pm}0.5$
Llama-2-7B	TyDiQA	$51.7{\scriptstyle\pm}0.4$	$\mathbf{52.9{\scriptstyle\pm}0.1}$	$50.4{\scriptstyle\pm}0.4$	$51.2{\scriptstyle\pm}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{\scriptstyle\pm}0.2$	$\mathbf{67.2{\scriptstyle\pm}0.1}$	$60.3{\scriptstyle\pm}0.1$	$63.6{\scriptstyle\pm}0.4$
Llama-3.2-3B	MMLU-Pro	$26.5{\scriptstyle\pm}0.2$	$27.1{\scriptstyle\pm}0.4$	$\mathbf{27.3{\scriptstyle\pm}0.1}$	$26.9{\scriptstyle\pm}0.3$
Llama-3.2-3B	GSM8K	$35.2{\scriptstyle\pm}0.4$	$36.7{\scriptstyle\pm}0.9$	$\mathbf{39.6{\scriptstyle\pm}0.3}$	$35.2{\scriptstyle\pm}0.9$
Llama-3.2-3B	Codex	$57.4{\scriptstyle\pm}0.7$	$58.3{\scriptstyle\pm}2.1$	$\mathbf{60.1{\scriptstyle\pm}2.0}$	$58.3{\scriptstyle\pm}4.3$
Llama-3.2-3B	BBH	$49.3{\scriptstyle\pm}0.5$	$49.1{\scriptstyle\pm}0.6$	$\mathbf{49.4{\scriptstyle\pm}0.4}$	$48.6{\scriptstyle\pm}0.2$
Llama-2-7B	TyDiQA	$52.8{\scriptstyle\pm}0.4$	$\mathbf{53.6{\scriptstyle\pm}0.1}$	$50.7{\scriptstyle\pm}0.6$	$52.1{\scriptstyle\pm}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.

Training Data Selection

Training data selection studies how to choose, filter, weight, or schedule examples to improve learning while reducing data and compute costs.

See all 186 papers

Form Over Content In Gradient-Based Data Attribution Methods
Sep 17, 2026Sunwoo Kim, Seokwon Jung, Sohyung Kim +2Gradient-Based Attribution Instruction-Tuning Data Selection
KAIST
Dr. Post-Training: A Data Regularization Perspective on LLM Post-Training
May 8, 2026Pingbang Hu, Xueshen Liu, Z. Morley Mao +1Supervised Fine-Tuning Training Data Selection
University of Illinois Urbana–Champaign · University of Michigan
Which and When to Admit: Gradient Admission for Data-Centric Small Language Model Finetuning
Oct 6, 2026Hongyu Cao, Yanchi Liu, Kunpeng Liu +5Gradient Interference Low-Rank Adaptation
Arizona State University · Clemson University · Meta AI +1

Sep 17, 2026J/K move · Enter open · S save

Form Over Content In Gradient-Based Data Attribution Methods

Sunwoo Kim, Seokwon Jung, Sohyung Kim, Seong Joon Oh+1

KAIST

Gradient-Based Attribution Instruction-Tuning Data Selection Training Data Attribution

Open paper AskPDF arXiv

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.

Abstract

Training Data Selection

Explore similar work

Form Over Content In Gradient-Based Data Attribution Methods

Dr. Post-Training: A Data Regularization Perspective on LLM Post-Training

Which and When to Admit: Gradient Admission for Data-Centric Small Language Model Finetuning