GradAlign: Gradient-Aligned Data Selection for LLM Reinforcement Learning
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
Code (0)
등록된 구현이 없습니다.
Tasks
Reinforcement LearningSimilar Papers 제목 키워드 기반
Fast Adversarial Training with Noise Augmentation: A Unified Perspective on RandStart and GradAlign
PGD-based and FGSM-based are two popular adversarial training (AT) approaches for obtaining adversarially robust models. Compared with PGD-based AT, FGSM-based one is significantly faster but fails with catastrophic over…
Data AugmentationGradAlign for Training-free Model Performance Inference
Architecture plays an important role in deciding the performance of deep neural networks. However, the search for the optimal architecture is often hindered by the vast search space, making it a time-intensive process. R…
modelNeural Architecture SearchPGD-2 can be better than FGSM + GradAlign
One major issue of adversarial training (AT) with the fast gradient sign method (FGSM AT) is the phenomenon of catastrophic overfitting, meaning that the trained model suddenly loses its robustness over a single epoch. I…
Understanding and Improving Fast Adversarial Training
A recent line of work focused on making adversarial training computationally efficient for deep learning models. In particular, Wong et al. (2020) showed that $\ell_\infty$-adversarial training with fast gradient sign me…
Learning What Helps: Task-Aligned Context Selection for Vision Tasks
Humans often resolve visual uncertainty by comparing an image with relevant examples, but ViTs lack the ability to identify which examples would improve their predictions. We present Task-Aligned Context Selection (TACS)…
Medical Image ClassificationMedical Image SegmentationReinforcement Learning