DeGMatch: Target-Aware Reasoning Data Selection via Deconfounded Gradient Matching

Changho Lee; Janghoon Han; Rajdeep Mukherjee; Yeda Song; Lening Nick Cui; Xinya Du; Honglak Lee

cs.CLOct 18, 2026

DeGMatch: Target-Aware Reasoning Data Selection via Deconfounded Gradient Matching

Authors: Changho Lee, Janghoon Han, Rajdeep Mukherjee, Yeda Song, Lening Nick Cui, Xinya Du, Honglak Lee

Organizations: University of Michigan · University of Texas at Dallas · LG AI Research

Abstract

Reasoning post-training on long trajectories has substantially improved the reasoning capabilities of large language models, but scaling such training to large datasets is increasingly expensive. We study target-aware reasoning data selection, where a small reference set specifies a target capability and guides the selection of training trajectories for specialization. We find that existing target-aware methods can exhibit limited target discrimination on long reasoning trajectories, while dominant population-level structure in candidate gradients can obscure finer target-dependent variation during gradient matching. To address these limitations, we introduce Deconfounded Gradient Matching (DeGMatch), which first constructs a target-relevance gate and then reranks candidates using model-specific gradients. DeGMatch centers gradients using the candidate-population mean and removes leading principal directions before target matching, reducing the influence of dominant population-level structure and improving target discrimination. Across mathematics, science, and code, DeGMatch outperforms strong heuristic, target-aware, and reasoning-specific data selection baselines, while remaining effective across selection budgets, model families, and scales. Our code is publicly available at https://github.com/CHLee0801/degmatch.

Figures & tables

Figure 1: Dominant gradient structure obscures fine-grained target variation. Each point represents a reference subset positioned by the candidate ranking it induces. (a) Rankings using only the top-50 principal-component subspace. (b) Rankings after removing the population mean and top 50 principal components. The corresponding five-way target-identification accuracy, computed in the original ranking space, increases from 56% to 71%. See Appendix B.7 for details.

Figure 2: Overview of DeGMatch . DeGMatch constructs a target-relevant gate, deconfounds the candidate gradient geometry, and ranks candidates using target-conditioned directional alignment and target-subspace magnitude to select the final training subset.

Method	# Ex.	Tokens	AIME24	AIME25	AIME26	AMC23	HMMT	MATH500	GK24	GK25	Avg.
Base Model	–	–	1.2	1.9	0.2	12	0.2	17.3	9.4	7.2	6.18
OpenR1-Qwen-7B ‡	93K	561.2M	47.9	34.8	35.6	84.2	24.0	92.5	72.6	78.8	58.80
Heuristic selection
Random	10K	60.3M	39.2	30.8	34.4	77.7	18.3	89.6	70.2	80.0	55.03
Length	10K	142.3M	45.0	33.3	40.4	82.3	21.7	91.8	74.3	81.0	58.73
MaxPPL	10K	117.2M	41.7	34.6	38.5	82.0	20.3	92.4	73.1	80.8	57.93

Table 1: Main results on mathematical reasoning. # Ex. denotes the number of training examples, and Tokens denotes the number of training tokens consumed per epoch. For target-aware methods, we report the average token count across target-specific training runs, with per-target statistics provided in Appendix E.2 . † Select2Reason uses the released Qwen-7B checkpoint, with tokens estimated from reported statistics. ‡ OpenR1-Qwen-7B is a full-data reference rather than a budget-matched baseline. Bold and underlined values indicate the best and second-best results among selection methods, respectively.

Science
Method	GPQA-D	JEE	# Ex.	Tokens
Random	46.8	45.8	5K	12.2M
Length	48.2	46.9	5K	35.8M
RDS+	46.4	51.6	5K	19.9M
TEMP	47.5	46.6	5K	11.6M
DeGMatch	49.0	51.8	5K	22.1M

Table 2: Cross-domain results on scientific reasoning and code generation. Tokens are reported per epoch and averaged across target-specific subsets for target-aware methods. Bold and underlined values denote the best and second-best results among selection methods.

Variant	A24	A25	A26	AMC	Avg.
Random	30.4	27.3	32.3	75.5	41.38
Relevance only	32.7	30.2	32.3	79.1	43.58
Raw-gradient dir.	36.9	30.8	34.0	78.0	44.93
Centered dir.	33.3	33.1	35.6	79.1	45.28
Deconfounded dir.	34.6	31.9	36.7	80.3	45.88
Deconfounded mag.	32.5	32.5	33.3	80.5	44.70

Table 3: Component validation for target relevance, candidate-population deconfounding, and the final matching score. Avg. is the unweighted mean over AIME24, AIME25, AIME26, and AMC23. Bold and underlined values indicate the best and second-best results, respectively.

Figure 3: Budget scaling by selected examples (top) and training tokens (bottom), measured by the average score on AIME24, AIME25, and AIME26. The dashed line denotes full-data training on 93K trajectories (561.2 M tokens). Detailed benchmark-level results are provided in Appendix E.3 .

Method	A24	A25	A26	AMC	Avg.	Tokens
Qwen2.5-3B-Base
Random	7.1	5.2	6.7	39.2	14.55	18.2M
Length	7.5	7.7	5.8	42.2	15.80	45.8M
RDS+	7.7	10.4	6.2	38.7	15.75	29.1M
DeGMatch	8.8	7.7	7.5	40.2	16.05	14.7M
Qwen3-1.7B-Base

Table 4: Cross-model robustness under a fixed 3K selection budget. Avg. is the unweighted mean over AIME24/25/26 and AMC23. Tokens denotes the total number of training tokens in each selected subset. Bold and underlined values denote the best and second-best results among selection methods.

Appendix figures & tables40 assets

Supplementary material from the paper’s appendix.

Appendix

Pool	# Examples	Tokens	Mean Len.
Math ( Open-R1 Team, 2025c )	93,733	561.2M	5,986
Science ( Open-R1 Team, 2025a )	53,195	131.9M	2,480
Code ( Open-R1 Team, 2025a )	26,061	337.0M	12,931

Appendix

Table 5: Candidate-pool statistics.

Domain	Target	Reference source	Evaluation set
Math	AIME	AIME 1983–2023 ( Neubig, )	AIME24/25/26 ( Mathematical Association of America, 2025a )
Math	AMC	AMC12 2018–2022 ( edev2000, )	AMC23 ( Mathematical Association of America, 2025b )
Math	HMMT	Omni-MATH HMMT ( KbsdJames, )	HMMT Feb. 2025 ( MathArena, 2025 )
Math	MATH	MATH train ( Hendrycks et al., 2021 )	MATH500 ( Lightman et al., 2024 )
Math	Gaokao	Gaokao 2023 ( MARIO-Math-Reasoning, 2023 )	Gaokao 2024/25 ( FrankieYao, )
Science	GPQA	GPQA non-Diamond ( Rein et al., 2024 )	GPQA-Diamond ( Rein et al., 2024 )

Appendix

Table 6: Target-reference sources and corresponding evaluation sets. Reference examples specify the target capability and are not directly included as supervised training examples.

Method	Implementation source	Preserved selection rule	Adaptation to our setting
RDS+	Published procedure	Position-weighted student hidden-state representations, cosine similarity, and round-robin target retrieval	Initial student representations computed over complete sequences up to 32K tokens
LESS	Official pipeline	Adam-preconditioned gradients across warmup checkpoints with learning-rate-weighted target alignment	Common LoRA warmup; per-example gradients computed over complete sequences up to 32K tokens
GIST	Official pipeline	Target-gradient Gram decomposition, 95% spectral-mass subspace, and maximum target-reference cosine	Common LoRA warmup; gradients computed over complete sequences up to 32K tokens
TAROT	Published method description	Multi-checkpoint gradient features, whitening, and fixed-budget OT-based selection	Common LoRA warmup; gradients computed over complete sequences up to 32K tokens
RSR	Official code	Rank–surprisal ratio from student-model response-token probabilities	Global fixed-budget ranking rather than within-problem trajectory selection
TEMP	Official code + adaptation	Difficulty filtering, perturbation-based loss trajectories, clustering, and learnability	OpenR1-Distill-7B as directional endpoint; source allocation follows the published equations

Appendix

Table 7: Implementation sources and adaptations for the principal comparison methods. For methods producing selected subsets in our pipeline, we preserve the published selection rule while using the same candidate pool, final selection budget, and downstream fine-tuning recipe as DeGMatch.

Method	Setting	Target	Reason.	Student	Primary signal	Grad.	Preparation
AlpaGasus ( Chen et al., 2024 )	Instruction	No	No	No	LLM-estimated quality	No	LLM scoring
InsTag ( Lu et al., 2024 )	Instruction	No	No	No	Instruction complexity / diversity	No	Tagging model
DEITA ( Liu et al., 2024a )	Instruction	No	No	No	Quality, complexity, diversity	No	Scoring models
RHO-LOSS ( Mindermann et al., 2022 )	General	No	No	Yes	Reducible loss / learnability	No	Reference-model fit
IFD ( Li et al., 2024b )	Instruction	No	No	Yes	Instruction-following difficulty	No	Model scoring
SuperFiltering ( Li et al., 2024a )	Instruction	No	No	No	Surrogate-model difficulty	No	Surrogate scoring

Appendix

Table 8: Extended comparison of representative data-selection methods. Target indicates explicit conditioning on an external target or reference set; Reason. indicates that the method is designed specifically for reasoning data; Student indicates use of student-model signals rather than surrogate-only scoring. Grad. distinguishes gradient-based signals from forward-only gradient proxies. The final column summarizes auxiliary fitting, scoring, or feature construction required for selection rather than quantitative compute.

Benchmark	# Problems	$K$ (Avg@ $K$ )	Scoring
AIME24/25/26 ( Mathematical Association of America, 2025a )	30 / 30 / 30	16	final-answer equivalence
AMC23 ( Mathematical Association of America, 2025b )	40	16	final-answer equivalence
HMMT Feb. 2025 ( MathArena, 2025 )	30	16	symbolic answer verification
MATH500 ( Lightman et al., 2024 )	500	4	math-verify
Gaokao 2024/25 ( FrankieYao, )	26 / 26	16	math-verify
GPQA-Diamond ( Rein et al., 2024 )	198	8	multiple-choice accuracy

Appendix

Table 9: Evaluation protocols. We report Avg@

K

, the average correctness across

K

sampled generations per problem.

Method	Auxiliary fitting	Reusable candidate representation	Approx. GPU-h
RDS+	None	Student hidden states (full-sequence forward)	$\sim$ 12
LESS	Four-epoch LoRA warm-up	Optimizer-aware LoRA gradients (4 checkpoints)	$\sim$ 200
GIST	LoRA warm-up	LoRA gradients at warm-up checkpoint	$\sim$ 57
TAROT	Four-epoch LoRA warm-up	Projected LoRA gradients (4 checkpoints)	$\sim$ 200
DeGMatch	None	Full-parameter CountSketch gradients at initialization	$\sim 58.7$

Appendix

Table 10: Approximate one-time selection-feature construction cost for the mathematical candidate pool under our 32K-context implementations. Candidate features can be reused across target capabilities. The DeGMatch value includes the Qwen3-Embedding-8B problem embedding of the pool used for the relevance gate (

\sim

0.3 GPU-h), the full-pool gradient pass, and one 50-example target-gradient computation. LESS, GIST, and TAROT use LoRA gradients in our implementations, whereas DeGMatch computes gradients with respect to all student parameters. These values characterize our implementations and hardware rather than algorithm-independent complexity.

Target	Input	Output	Cost
AIME	5.5K	918.2K	$1.98
AMC	4.9K	523.5K	$1.31
HMMT	5.6K	1.08M	$2.71
MATH500	3.3K	256.5K	$0.64
Gaokao	4.4K	310.9K	$0.78
GPQA	9.4K	486.0K	$1.22

Appendix

Table 11: Token usage and monetary cost for generating verified target-reference trajectories. Each target contains 50 reference examples, for 450 references in total. Costs reflect the API pricing used at the time of generation.

Target	Top-1	Top-5	Top-10	Top-50	Top-200
AIME	0.343	0.732	0.794	0.873	0.917
AMC	0.355	0.747	0.806	0.881	0.922
HMMT	0.326	0.717	0.782	0.867	0.913
MATH500	0.391	0.773	0.826	0.894	0.930
Gaokao	0.385	0.769	0.822	0.892	0.929

Appendix

Table 12: Cumulative variance explained by the leading principal directions of centered candidate gradients in each target-specific relevance gate.

Panel of Figure 4 — Figure 4: Candidate gradients exhibit strongly concentrated population-level structure. Top: cumulative variance explained by the leading principal components of centered gradient sketches in the 30K AIME relevance gate, together with representative associations between PC1 and gradient norm and between PC2 and trajectory length. Bottom: absolute Spearman correlations between the first five PC scores and trajectory length, student NLL, and gradient norm. The top 50 components explain 87.3% of the variance; PC1 is strongly associated with gradient norm ( $|\rho|=0.86$ ), and PC2 with trajectory length ( $|\rho|=0.66$ ).

Property	PC1	PC2	PC3	PC4	PC5	# PCs with $\|\rho\|>0.3$
Log trajectory length	0.61	0.66	0.28	0.02	0.11	2
Student NLL	0.60	0.63	0.20	0.27	0.08	2
Gradient norm	0.86	0.46	0.09	0.13	0.05	2
Olympiad source	0.01	0.12	0.07	0.12	0.06	0
Algebra	0.16	0.22	0.18	0.19	0.02	0
Geometry	0.08	0.10	0.10	0.15	0.35	2

Appendix

Table 13: Absolute Spearman correlations between scores along the leading gradient PCs and trajectory properties. The final column counts components among the first 50 for which

|\rho|>0.3

Set	$k=10$	$k=25$	$k=50$	$k=100$
Random examples	0.774	0.827	0.861	0.884
AIME references	0.658	0.720	0.777	0.817
AMC references	0.687	0.749	0.805	0.841
HMMT references	0.654	0.714	0.781	0.819
MATH500 references	0.809	0.856	0.886	0.906
Gaokao references	0.795	0.835	0.865	0.889

Appendix

Table 14: Fraction of centered-gradient energy contained in the leading candidate-gradient subspace. The gate statistics are computed from 50 randomly sampled gate examples.

PC range	Across targets	Pool-split ref.	Random
1–5	0.93	0.94	0.0003
1–50	0.79	0.87	0.0031
51–100	0.45	0.60	0.0031

Appendix

Table 15: Normalized subspace overlap across target-conditioned candidate populations after excluding candidates shared between relevance gates. The pool-split reference uses non-overlapping random subsets of the mathematical pool and reflects finite-sample stability.

Figure 5: Why DeGMatch removes the dominant candidate-gradient directions. (a) The raw AIME ranking is almost completely reproduced using only the dominant population component (the mean direction and top 50 centered candidate PCs), whereas (a ′ ) the orthogonal residual alone produces a substantially different ranking. (b–b ′ ) A length- and NLL-matched non-target reference set induces a ranking highly correlated with the AIME ranking in the raw gradient space, but this correspondence is sharply reduced after dominant-component removal. (c) The effect is specific to the leading principal directions: after population centering, removing the top 50 PCs reduces the matched-control rank correlation from $0.87$ to $0.28$ , whereas removing 50 random directions or lower-ranked PC blocks leaves it essentially unchanged.

Representation / correction	Matched-control $\rho$	Random-ref. $\rho$	Overlap w/ deconf.	Top-3K tokens	$\rho$ (score, len.)
Raw direction	0.86	0.58	0.18	28.5M	+0.72
Score residualization (linear)	0.79	0.62	0.02	13.2M	$-0.17$
Score residualization (cubic)	0.70	0.56	0.02	10.4M	$-0.11$
Gradient residualization	0.78	0.73	0.12	17.2M	$-0.16$
Mean + PC50	0.28	0.08	–	30.8M	+0.34
Covariate residualization + PC50	0.24	0.09	0.63	30.0M	+0.32

Appendix

Table 16: Comparison with explicit residualization of measured trajectory covariates within the fixed 30K AIME relevance gate. Matched-control

\rho

is the Spearman correlation between rankings induced by the AIME references and by non-target references matched in trajectory length and student NLL. Random-ref.

\rho

reports correlation with rankings induced by random reference sets. Overlap w/ deconf. is the fraction of the top-3K subset shared with the default mean-centered, top-50-PC-removed directional ranking. The regression variants use log trajectory length, student NLL, and log gradient norm as covariates.

Figure 6: Shared gradient structure becomes more prominent with reasoning-trajectory length. Left: mean-direction energy after recomputing gradients for the same reasoning subset truncated to 256, 1,024, and 4,096 response tokens, together with the full trajectories. Right: sample-to-mean gradient cosine across deciles of the original reasoning-trajectory length distribution (20,000 trajectories). The truncation analysis holds the underlying examples fixed while varying the observed response length, whereas the decile analysis measures the corresponding association in the original data.

Panel of Figure 7 — Figure 7: Target sensitivity varies substantially across target-aware selectors. Top: mean pairwise Jaccard similarity among the 10K subsets independently selected for AIME, AMC, HMMT, MATH500, and Gaokao. The dashed line denotes the expected Jaccard similarity between two independently sampled 10K subsets from the 93,733-example mathematical pool. Bottom: pairwise Jaccard matrices across the five targets. Higher overlap indicates that the selected subset changes less when the supplied target references change.

	IT	Short Math	R1 Reasoning
Median response tokens	33	159	4,826
Mean-direction energy	0.016	0.100	0.438
Sample–mean cosine	$-0.02$	0.40	0.71
Raw random-ref. $\rho$	0.35	–	0.75

Appendix

Table 17: Shared gradient structure across instruction-tuning examples, short mathematical solutions, and long R1 reasoning trajectories (4,000 examples each). All gradients are computed using the same student model and CountSketch representation. Mean-direction energy is the fraction of gradient energy contained in the population mean direction, and sample–mean cosine measures the alignment of individual gradients with that direction. Raw random-ref.

\rho

is the Spearman correlation between rankings induced by the target references and randomly sampled reference trajectories.

Representation	Matched-control $\rho$	AMC	HMMT	MATH500	Gaokao	$\rho$ (score, len.)
Raw gradients	0.86	0.92	0.96	0.90	0.89	0.78
Mean-centered	0.87	0.68	0.88	0.45	0.45	0.58
Mean + PC10 removal	0.39	0.75	0.78	0.62	0.66	0.56
Mean + PC25 removal	0.36	0.74	0.79	0.58	0.59	0.63
Mean + PC50 removal	0.28	0.72	0.78	0.56	0.59	0.69
Mean + PC100 removal	0.24	0.71	0.75	0.57	0.58	0.69

Appendix

Table 18: Target-specificity diagnostics on the AIME relevance gate. Matched-control

\rho

is the Spearman correlation obtained after replacing the AIME references with a length- and NLL-matched non-target control. The four middle columns report top-10K intersection fractions after replacing the AIME references with another mathematical target.

Panel of Figure 8 — Figure 8: Removing dominant candidate directions increases target sensitivity. Top left: Spearman correlation between the ranking induced by the AIME references and rankings induced by length- and NLL-matched non-target controls as the gradient representation changes (mean $\pm$ s.d. over three draws). Mean-centering leaves the matched-control correlation essentially unchanged, whereas principal-component removal sharply reduces it. Top right: mean pairwise intersection fraction among the top-10K subsets induced by five mathematical targets when all targets score the same 30K AIME candidate set. Bottom: pairwise intersection fractions for raw gradients, mean-centered gradients, mean plus top-50-PC removal, and the final DeGMatch selections using target-specific relevance gates.

Figure 9: Dominant-component removal reveals finer target-conditioned structure in candidate rankings. Each point represents the ranking induced by a disjoint 10-example reference subset over the same fixed 30K candidate population. From left to right: raw gradients, population centering alone, projection onto the top 50 principal components after centering, and the residual representation after removing the population mean and top 50 components. Distances are

1-\mathrm{Spearman}

and are visualized using classical MDS; each panel is scaled independently for visualization. Quantitative target-identification results are computed in the original ranking space rather than from the two-dimensional projection.

Figure 10: Deconfounding reduces concentration on the same highly ranked candidates across targets. t-SNE visualization of the top-1K candidates selected for each of five mathematical targets under raw-gradient matching (left) and after mean-centering and removal of the top 50 candidate principal directions (right). Black points denote candidates selected by at least two targets. Raw matching yields a union of only 1,317 distinct candidates across the five top-1K subsets, of which 54.0% are selected by all five targets. After deconfounding, the union expands to 2,950 distinct candidates, of which only 2.4% are selected by all five targets. The reported selection statistics are computed from the original selected sets rather than from the two-dimensional projection.

$k$	Var.	Ctrl. $\rho$	Overlap	Tokens
10	0.794	0.39	0.63	96.8M
25	0.845	0.36	0.69	96.0M
50	0.873	0.28	1.00	98.3M
100	0.897	0.24	0.73	98.7M
200	0.917	0.20	0.67	99.1M

Appendix

Table 19: Diagnostics as the number of removed principal components varies. Var. is cumulative explained variance; Ctrl.

\rho

is the matched-control correlation from Table 18 ; Overlap is measured against the default selection with

k=50

Figure 11: Target-ranking structure is robust to the number of removed principal components. Target identification is computed from candidate rankings induced by disjoint reference subsets, following Figure 9 . Blue removes the leading

k

components after population centering; orange retains only those components. Dashed horizontal lines denote raw-gradient and mean-centered baselines.

$k$	AIME24	AIME25	AIME26	Avg.
0	33.3	33.1	35.6	34.00
10	37.1	32.0	33.0	34.03
25	36.0	32.2	34.2	34.13
50	37.9	31.0	35.6	34.83
100	35.0	33.0	34.6	34.20

Appendix

Table 20: Downstream sensitivity to the number of removed principal components under a fixed 3K selection budget.

k=0

corresponds to population centering without principal-component removal. Avg. is the unweighted mean over AIME24, AIME25, and AIME26.

Estimator	Overlap w/ default
Full 30K gate PCA	1.00
6K gate sample, seed 0	0.87
6K gate sample, seed 1	0.89
6K gate sample, seed 2	0.90
6K gate sample, seed 3	0.90
Full-pool PCA	0.88

Appendix

Table 21: Sensitivity to the population used to estimate the dominant gradient subspace. Overlap is measured against the default selection obtained from PCA over the complete 30K relevance gate.

Target	Gate	$d_{i}<0$	Pos./Budget	Min. sel. $d_{i}$
AIME	30K	0.16%	29,951/10K	0.017
AMC	30K	0.03%	29,991/10K	0.014
HMMT	30K	0.13%	29,962/10K	0.025
MATH500	30K	0.04%	29,987/10K	0.023
Gaokao	30K	0.06%	29,983/10K	0.021
GPQA	15K	0.15%	14,978/5K	0.015

Appendix

Table 22: Directional-alignment statistics within each relevance gate. Pos./Budget reports the number of candidates with

d_{i}>0

relative to the final selection budget, and Min. sel.

d_{i}

is the minimum directional alignment among selected examples.

Gate	Size	Overlap w/ default	Tokens	Median len.	Median NLL	AIME sim.
$1\times B$	10K	0.36	68.9M	5,968	1.26	0.531
$2\times B$	20K	0.69	90.0M	8,326	1.49	0.525
$3\times B$	30K	1.00	98.3M	9,196	1.56	0.519
$5\times B$	50K	0.64	107.4M	10,378	1.61	0.507
No gate	93,733	0.47	115.8M	11,642	1.66	0.493

Appendix

Table 23: Effect of relevance-gate size on the final AIME subset. Overlap w/ default is the fraction shared with the subset selected using the default

3B

gate.

Gate size	AIME24	AIME25	AIME26	Avg.
9K	34.6	32.3	36.2	34.37
15K	37.5	31.5	34.0	34.33
30K	37.9	31.0	35.6	34.83
Full pool	37.3	29.7	34.7	33.90

Appendix

Table 24: Downstream sensitivity to the relevance-gate size under a fixed 3K selection budget. Full pool removes the relevance restriction and allows gradient reranking over the complete mathematical candidate pool. Avg. is the unweighted mean over AIME24, AIME25, and AIME26.

Gate construction	Gate Jaccard	Top-1K coverage min./med.	Selection overlap	AIME sim.
Qwen3 + round-robin	1.00	893 / 949	1.00	0.519
Student hidden state + round-robin	0.30	214 / 506	0.47	0.499
Qwen3 + max similarity	0.68	96 / 913	0.76	0.522
Qwen3 + mean similarity	0.40	40 / 702	0.56	0.515

Appendix

Table 25: Effect of the embedding representation and retrieval aggregation rule. Top-1K coverage reports, across target references, the minimum and median number of each reference’s 1,000 nearest candidates retained in the 30K gate. Selection overlap is measured against the final subset obtained with the default gate.

Set	AIME24	AIME25	$\cos>0.7$
Random 10K	0.648	0.639	1.0
Length 10K	0.658	0.642	1.4
Relevance gate 30K	0.691	0.681	7.0
DeGMatch 10K	0.668	0.657	2.9

Appendix

Table 26: Post-hoc semantic proximity to AIME24/25 evaluation problems. The first two columns report the mean cosine similarity of the ten nearest training problems per evaluation problem. The final column reports the mean number of neighbors with cosine similarity above 0.7. Evaluation problems are used only for this diagnostic.

Set	Length (tokens)	Student NLL	Grad. norm	AIME sim.	NT+Comb.	Olympiad
Pool	4,935 [2,959, 8,107]	1.14 [0.90, 1.46]	31.2	0.47	0.29	0.73
Random	4,974 [3,019, 8,116]	1.15 [0.91, 1.47]	31.5	0.47	0.29	0.73
Length	13,940 [12,669, 15,574]	1.77 [1.58, 2.01]	42.3	0.49	0.38	0.67
MaxPPL	11,786 [9,066, 14,460]	1.95 [1.84, 2.13]	42.2	0.48	0.56	0.70
RDS+	10,104 [7,403, 13,364]	1.50 [1.28, 1.74]	41.2	0.50	0.47	0.63
DeGMatch	9,196 [6,823, 12,600]	1.56 [1.33, 1.82]	41.0	0.52	0.48	0.64

Appendix

Table 27: Properties of mathematical subsets selected for AIME. Length and student NLL are reported as median [IQR] when available. NT+Comb. is the combined fraction of examples whose dataset-provided problem_type is Number Theory or Combinatorics; Olympiad is the fraction whose dataset-provided source is olympiads .

Method	Tokens	Med. len.	A24	A25	A26	Avg.
Rand-LC	35.3M	12,102	38.4	30.6	32.1	33.70
DeGMatch	36.1M	12,274	37.9	31.0	35.6	34.83

Appendix

Table 28: Length-controlled downstream comparison under a fixed 3K selection budget. Rand-LC (length-controlled random sampling) draws from the same 30K AIME relevance gate as DeGMatch, excluding DeGMatch-selected examples, and approximately matches its complete-sequence length distribution using 20 quantile strata.

	AIME	AMC	HMMT	MATH	GK
AIME	1.00	0.21	0.27	0.12	0.11
AMC		1.00	0.20	0.13	0.12
HMMT			1.00	0.11	0.12
MATH500				1.00	0.14
Gaokao					1.00

Appendix

Table 29: Pairwise Jaccard overlap among target-specific 10K DeGMatch subsets. Two independently sampled uniform 10K subsets from the 93,733-example mathematical pool have expected Jaccard overlap of approximately 0.056.

Rank shift	#	Med. len.	Med. NLL	Alg.	Comb.	Geom.	NT
Raw top 1K $\rightarrow$ below 10K	258	9,374	1.57	66	42	104	39
Raw below 10K $\rightarrow$ DeGMatch top 1K	276	14,572	1.85	28	102	47	86

Appendix

Table 30: Candidate groups undergoing large rank changes after deconfounding within the AIME relevance gate. The final four columns report counts of algebra, combinatorics, geometry, and number-theory problems; less frequent categories are omitted.

Ex.	Selected training problem	Shared reasoning structure	Deconf. rank	Raw rank	Semantic rank
A	Six people are seated in a row, where three designated people must not sit next to one another.	Forbidden-adjacency arrangement problem closely matching the reference at both the problem and solution level.	1	6,879	1
B	Three red, two white, and two blue balls are arranged in a row; compute the probability that the two white balls or the two blue balls are adjacent.	Both solutions define adjacency events, merge adjacent objects into blocks, count the resulting arrangements, and combine overlapping cases using inclusion–exclusion.	13	14,215	4,862

Appendix

Table 31: Illustrative target-conditioned matches for the AIME 2020-II-9 reference. All ranks are computed among the 30K examples in the relevance gate for this reference, with rank 1 denoting the closest match. Semantic rank uses Qwen3 problem-embedding cosine similarity; the corresponding cosine similarities are 0.80 for Example A and 0.43 for Example B.

Variant	Rel.	Transform	Dir.	Mag.
Random	–	–	–	–
Relevance only	✓	–	–	–
Raw-gradient dir.	✓	None	✓	–
Centered dir.	✓	Center	✓	–
Deconfounded dir.	✓	Center + PC50	✓	–
Deconfounded mag.	✓	Center + PC50	–	✓

Appendix

Table 32: Component-validation variants. Rel. denotes target-relevance filtering; Center denotes candidate-population mean subtraction; PC50 denotes removal of the top 50 principal components; and Dir. and Mag. denote directional alignment and target-subspace magnitude, respectively.

Method	AIME24	AIME25	AIME26	AMC23	HMMT	MATH500	GK24	GK25	Avg.	Tokens
Length	45.0	33.3	40.4	82.3	21.7	91.8	74.3	81.0	58.73	142.3M
Gate30K+Length	45.5	35.0	40.0	83.1	23.0	92.0	74.3	80.0	59.11	114.6M
MaxPPL	41.7	34.6	38.5	82.0	20.3	92.4	73.1	80.8	57.93	117.2M
Gate30K+MaxPPL	44.8	31.7	41.9	82.7	22.0	92.8	73.6	80.8	58.79	105.5M
DeGMatch	47.7	33.8	42.6	85.6	23.7	92.8	76.9	82.7	60.73	86.9M

Appendix

Table 33: Gate-controlled heuristic baselines under a fixed 10K-example budget. Gate30K applies the corresponding heuristic only within the same 30K AIME relevance gate used by DeGMatch. Avg. is the unweighted mean over the eight mathematical benchmarks.

Variant	Score	A24	A25	A26	Avg.	Tokens
Direction only ( $\gamma=0$ )	$d_{i}$	34.6	31.9	36.7	34.40	30.8M
Weak magnitude ( $\gamma=\frac{1}{2}$ )	$d_{i}\sqrt{m_{i}}$	36.0	30.0	34.6	33.53	34.6M
DeGMatch ( $\gamma=1$ )	$\mathbf{d_{i}m_{i}}$	37.9	31.0	35.6	34.83	36.1M
Strong magnitude ( $\gamma=2$ )	$d_{i}m_{i}^{2}$	37.1	32.1	34.6	34.60	37.1M
Length-normalized	$d_{i}m_{i}/\sqrt{T_{i}}$	34.2	30.0	32.1	32.10	28.1M
Length-residualized	$d_{i}\exp(m_{i}^{\mathrm{res}})$	36.0	32.9	35.0	34.63	27.1M

Appendix

Table 34: Sensitivity to alternative magnitude scalings under a fixed 3K selection budget.

T_{i}

denotes trajectory length and

m_{i}^{\mathrm{res}}

denotes the length-residualized magnitude. All variants use the same deconfounded gradient representation and 30K AIME relevance gate.

Method	AIME	AMC	Gaokao	MATH500	HMMT	Avg.
Heuristic
Random	60.3 (6.0K)					60.3
Length	142.3 (14.2K)					142.3
MaxPPL	117.2 (11.7K)					117.2
Target-aware
RDS+	104.9 (10.5K)	91.0 (9.1K)	63.3 (6.3K)	56.6 (5.7K)	110.0 (11.0K)	85.2

Appendix

Table 35: Training tokens for each 10K-example mathematical subset, in millions, with mean tokens per example in parentheses. Target-agnostic methods use the same subset for every target, whereas target-aware methods select a separate subset for each target. No selected sequence exceeds the 32,768-token training limit, so no example is truncated.

Method	# Ex.	AIME24	AIME25	AIME26	Avg.	Tokens
Random	1K	27.7	22.9	29.2	26.60	6.1M
	3K	30.4	27.3	32.3	30.00	18.0M
	5K	39.4	30.0	30.8	33.40	30.1M
	10K	39.2	30.8	34.4	34.80	60.3M
	15K	41.2	30.8	39.2	37.07	90.0M
Length	1K	30.8	24.8	29.4	28.33	17.7M

Appendix

Table 36: Detailed results across selection budgets. We report performance on AIME24, AIME25, and AIME26, their arithmetic mean, and the total number of training tokens in each selected subset.

Data Selection for Language Model Fine-Tuning

See all 25 papers

GRACE: Gradient-aligned Reasoning Data Curation for Efficient Post-training
May 13, 2026Junjie Li, Ziao Wang, NingXuan Ma +2Training Data Selection LLM Post-Training
Harbin Institute of Technology, Shenzhen, China · Hong Kong Baptist University, China
Mitigating Shortcut Reasoning in Language Models: A Gradient-Aware Training Approach
Mar 21, 2026Hongyu Cao, Kunpeng Liu, Dongjie Wang +1Shortcut Learning in Language Models LLM Reasoning
Arizona State University · Clemson University · University of Kansas
Capacity-Dependent Effects of Data Selection for Reasoning
Aug 13, 2026Cuong Dang, Hoang Anh Just, Ruoxi JiaInstruction-Tuning Data Selection Supervised Fine-Tuning
Department of Electrical and Computer Engineering Virginia Tech Blacksburg, VA, USA

May 13, 2026J/K move · Enter open · S save

GRACE: Gradient-aligned Reasoning Data Curation for Efficient Post-training

Junjie Li, Ziao Wang, NingXuan Ma, Jianghong Ma+1

Harbin Institute of Technology, Shenzhen, China · Hong Kong Baptist University, China

Training Data Selection LLM Post-Training Gradient Alignment Reasoning Trace Analysis

Open paper AskPDF arXiv

Existing reasoning data curation pipelines score whole samples, treating every intermediate step as equally valuable. In reality, steps within a trace contribute very unevenly, and selecting reasoning data well requires assessing them individually. We present GRACE, a gradient-aligned curation method that views each reasoning trace as a sequence of optimization events and scores every step by two complementary signals: its alignment with the answer-oriented gradient direction, and its consistency with the preceding reasoning trajectory. Step-level scores are aggregated into a sample-level value for subset selection, using only the model's internal optimization signals and no external reward models or step annotations. To make this scalable, GRACE introduces a representation-level gradient proxy that estimates step-level alignment from token-level upstream signals in a single forward pass. Post-training Qwen3-VL-2B-Instruct on MMathCoT-1M, GRACE reaches 108.8% of the full-data performance with 20% of the data and retains 100.2% with only 5%, with subsets that transfer effectively across model backbones.

Abstract

Data Selection for Language Model Fine-Tuning

Explore similar work

GRACE: Gradient-aligned Reasoning Data Curation for Efficient Post-training

Mitigating Shortcut Reasoning in Language Models: A Gradient-Aware Training Approach

Capacity-Dependent Effects of Data Selection for Reasoning