# GradAlign: Gradient-Aligned Data Selection for LLM Reinforcement Learning

arXiv: https://arxiv.org/abs/2602.21492 · HTML: https://mlsanity.com/paper/2602.21492/

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 44 of its 50 references that are in the map.
- First posted on arXiv 25 Feb 2026.

## Builds on (most cited first, 30 of 44)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [DAPO: An Open-Source LLM Reinforcement Learning System at Scale](https://mlsanity.com/paper/2503.14476/) (2025)
- [MMLU-Pro: A More Robust and Challenging Multi-Task Language Understanding Benchmark](https://mlsanity.com/paper/2406.01574/) (2024)
- [Understanding Black-box Predictions via Influence Functions](https://mlsanity.com/paper/1703.04730/) (2017)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Scaling Laws for Reward Model Overoptimization](https://mlsanity.com/paper/2210.10760/) (2022)
- [The RefinedWeb Dataset for Falcon LLM: Outperforming Curated Corpora with Web Data, and Web Data Only](https://mlsanity.com/paper/2306.01116/) (2023)
- [LESS: Selecting Influential Data for Targeted Instruction Tuning](https://mlsanity.com/paper/2402.04333/) (2024)
- [Documenting Large Webtext Corpora: A Case Study on the Colossal Clean Crawled Corpus](https://mlsanity.com/paper/2104.08758/) (2021)
- [Open-Reasoner-Zero: An Open Source Approach to Scaling Up Reinforcement Learning on the Base Model](https://mlsanity.com/paper/2503.24290/) (2025)
- [DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining](https://mlsanity.com/paper/2305.10429/) (2023)
- [DataComp-LM: In search of the next generation of training sets for language models](https://mlsanity.com/paper/2406.11794/) (2024)
- [Understanding the Effects of RLHF on LLM Generalisation and Diversity](https://mlsanity.com/paper/2310.06452/) (2023)
- [Data Selection for Language Models via Importance Resampling](https://mlsanity.com/paper/2302.03169/) (2023)
- [Studying Large Language Model Generalization with Influence Functions](https://mlsanity.com/paper/2308.03296/) (2023)
- [SemDeDup: Data-efficient learning at web-scale through semantic deduplication](https://mlsanity.com/paper/2303.09540/) (2023)
- [SWE-smith: Scaling Data for Software Engineering Agents](https://mlsanity.com/paper/2504.21798/) (2025)
- [DeepMath-103K: A Large-Scale, Challenging, Decontaminated, and Verifiable Mathematical Dataset for Advancing Reasoning](https://mlsanity.com/paper/2504.11456/) (2025)
- [TheoremQA: A Theorem-driven Question Answering dataset](https://mlsanity.com/paper/2305.12524/) (2023)
- [Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance](https://mlsanity.com/paper/2403.16952/) (2024)
- [RegMix: Data Mixture as Regression for Language Model Pre-training](https://mlsanity.com/paper/2407.01492/) (2024)
- [On the Accuracy of Influence Functions for Measuring Group Effects](https://mlsanity.com/paper/1905.13289/) (2019)
- [Stream of Search (SoS): Learning to Search in Language](https://mlsanity.com/paper/2404.03683/) (2024)
- [MAmmoTH2: Scaling Instructions from the Web](https://mlsanity.com/paper/2405.03548/) (2024)
- [LIMR: Less is More for RL Scaling](https://mlsanity.com/paper/2502.11886/) (2025)
- [General-Reasoner: Advancing LLM Reasoning Across All Domains](https://mlsanity.com/paper/2505.14652/) (2025)

## Built on by (most cited first, 12 of 12)

- [Socratic-SWE: Self-Evolving Coding Agents via Trace-Derived Agent Skills](https://mlsanity.com/paper/2606.07412/) (2026)
- [Auditing Reward Hackability in Code RL Training Environments](https://mlsanity.com/paper/2606.16062/) (2026)
- [Guiding LLM Post-training Data Engineering with Model Internals from Sparse Autoencoders](https://mlsanity.com/paper/2605.27354/) (2026)
- [Mechanistic Circuit Identification for Controllable Data Generation](https://mlsanity.com/paper/2608.24065/) (2026)
- [PAIR: Pairwise-Aware Inclusion Reweighting for Adaptive Rollout Allocation in RLVR](https://mlsanity.com/paper/2608.11368/) (2026)
- [Gradients Know What Outcomes Don't: Unlocking Reinforcement Learning for LLM Reasoning with Gradient-Aligned Rewards](https://mlsanity.com/paper/2609.03342/) (2026)
- [Let the Target Select for Itself: Data Selection via Target-Aligned Paths](https://mlsanity.com/paper/2605.09404/) (2026)
- [INFUSER: Influence-Guided Self-Evolution Improves Reasoning](https://mlsanity.com/paper/2606.09052/) (2026)
- [Manifold Bandits: Bayesian Curriculum Learning over the Latent Geometry of Large Language Models](https://mlsanity.com/paper/2606.19750/) (2026)
- [CompassPlay: Rewarding the Proposer for Where It Moves the Solver](https://mlsanity.com/paper/2609.32228/) (2026)
- [Not All Rollouts Are Worth Learning: On Trajectory Valuation for Post-Training Reinforcement Learning](https://mlsanity.com/paper/2609.35072/) (2026)
- [LESSER: Post-Training Data Selection with Output-Layer Gradients](https://mlsanity.com/paper/2610.03702/) (2026)

## Abstract (arXiv)

Reinforcement learning (RL) has become a central post-training paradigm for large language models (LLMs), but its performance is highly sensitive to the quality of training problems. This sensitivity stems from the non-stationarity of RL: rollouts are generated by an evolving policy, and learning is shaped by exploration and reward feedback, unlike supervised fine-tuning (SFT) with fixed trajectories. As a result, prior work often relies on manual curation or simple heuristic filters (e.g., accuracy), which can admit incorrect or low-utility problems. We propose GradAlign, a gradient-aligned data selection method for LLM reinforcement learning that uses a small, trusted validation set to prioritize training problems whose policy gradients align with validation gradients, yielding an adaptive curriculum. We evaluate GradAlign across three challenging data regimes: unreliable reward signals, distribution imbalance, and low-utility training corpus, showing that GradAlign consistently outperforms existing baselines, underscoring the importance of directional gradient signals in navigating non-stationary policy optimization and yielding more stable training and improved final performance. We release our implementation at https://github.com/StigLidu/GradAlign