GradAlign: LLM 강화 학습을 위한 기울기 정렬 데이터 선택 방법

GradAlign: Gradient-Aligned Data Selection for LLM Reinforcement Learning

강화 학습(RL)은 대규모 언어 모델(LLM)의 사후 훈련 패러다임으로 자리 잡았지만, 성능은 훈련 문제의 품질에 매우 민감합니다. 이러한 민감성은 RL의 비정상성에서 비롯됩니다. RL에서는 정책이 계속 발전함에 따라 샘플이 생성되고, 학습은 탐색 및 보상 피드백에 의해 결정되며, 이는 고정된 경로를 가진 지도 미세 조정(SFT)과는 다릅니다. 결과적으로, 기존 연구에서는 종종 수동 큐레이션 또는 간단한 휴리스틱 필터(예: 정확도)에 의존하는데, 이는 부정확하거나 유용성이 낮은 문제를 포함할 수 있습니다. 본 논문에서는 LLM 강화 학습을 위한 기울기 정렬 데이터 선택 방법인 GradAlign을 제안합니다. GradAlign은 신뢰할 수 있는 검증 세트를 사용하여 정책 기울기와 검증 기울기가 일치하는 훈련 문제를 우선적으로 선택하여 적응형 학습 커리큘럼을 제공합니다. 우리는 GradAlign을 세 가지 어려운 데이터 환경(신뢰할 수 없는 보상 신호, 분포 불균형, 낮은 유틸리티의 훈련 코퍼스)에서 평가한 결과, GradAlign이 기존의 방법보다 일관되게 우수한 성능을 보임을 확인했습니다. 이는 비정상적인 정책 최적화를 탐색하고 더 안정적인 학습을 유도하며 최종 성능을 향상시키는 데 방향성 기울기 신호의 중요성을 강조합니다. 구현 코드는 https://github.com/StigLidu/GradAlign 에서 확인할 수 있습니다.

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