GradAlign: Gradient-Aligned Data Selection for LLM Reinforcement Learning
Organizations: Institute for Interdisciplinary Information Sciences (IIIS), Tsinghua University · Language Technologies Institute (LTI), Carnegie Mellon University
Abstract
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
Figures & tables
| Val | Held-out Test | Ratio of Corrupted ( ) | ||||
|---|---|---|---|---|---|---|
| Method | AMC22 | AMC23 | AIME2425 | Test Avg. | Step 0 | Step 100 |
| GradAlign (Ours) | 49.1 | 68.2 | 20.5 | 44.4 | 17.8% | 29.9% |
| Random | 47.5 | 58.5 | 15.8 | 37.2 | 50.0%* | 50.0%* |
| AccGreedy | 18.1 | 28.5 | 1.9 | 15.2 | 81.8% | 87.1% |
| Align | 45.7 | 61.5 | 16.6 | 39.1 | 33.5% | 47.5% |
| Performance | Ratio of Countdown ( ) | |||
|---|---|---|---|---|
| Method | Countdown(Val) | Countdown(Test) | Step 0 | Step 50 |
| GradAlign (Ours) | 33.8 | 34.0 | 92.9% | 62.1% |
| Random | 14.1 | 10.4 | 11.8%* | 11.8%* |
| AccGreedy | 21.7 | 15.2 | 7.8% | 1.5% |
| Align | 31.1 | 28.2 | 47.2% | 60.6% |
| SuperGPQA | TheoremQA | AIME/AMC | MMLU-Pro | Test | ||||||
|---|---|---|---|---|---|---|---|---|---|---|
| Method | Val | Test | Val | Test | Val | AMC23 | AIME2425 | Val | Test | Avg. |
| GradAlign | 33.5 | 35.7 | 55.7 | 59.3 | 48.8 | 64.2 | 18.3 | 59.8 | 57.9 | 47.1 |
| Random | 29.6 | 35.5 | 53.5 | 55.2 | 47.6 | 60.0 | 17.2 | 57.8 | 57.1 | 45.0 |
| AccGreedy | 31.7 | 31.6 | 53.1 | 54.9 | 45.0 | 62.3 | 17.8 | 57.5 | 55.1 | 44.3 |
| Align | 29.8 | 32.7 | 52.8 | 57.9 | 46.4 | 61.5 | 16.1 | 57.1 | 56.6 | 45.0 |
| GradAlign (val. set size) | Baselines | |||||
| 10 | 30 | 100 | Random | AccGreedy | Align | |
| Test Set Accuracy | 32.7 | 32.1 | 34.0 | 10.4 | 15.2 | 28.2 |
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
| Hyperparameter | MMLU-Pro | SuperGPQA & TheoremQA | Math (AIME/AMC) | Countdown | Noisy Rewards (Distinguishability) |
|---|---|---|---|---|---|
| Model Configuration | |||||
| Base Model | Qwen2.5-1.5B-Math-Instruct | Qwen3-8B-Base | Qwen3-8B-Base | Qwen2.5-1.5B-Math | Qwen3-8B-Base |
| GRPO Training | |||||
| Learning Rate | |||||
| Optimizer | AdamW | AdamW | AdamW | AdamW | AdamW |
| Training Problems per Step ( ) | 32 | 128 | 128 | 128 | 128 |
| Number of Samples | 32 | 128 | 512 |
|---|---|---|---|
| Cosine Similarity | 0.30 | 0.49 | 0.79 |
RL for Language Models
Reinforcement learning for language models uses feedback-driven policy updates to improve reasoning, alignment, and interactive behavior.
See all 433 papers