STAPO: LLM 강화 학습의 안정화를 위한 희귀한 오해석 토큰 제거 방법

STAPO: Stabilizing Reinforcement Learning for LLMs by Silencing Rare Spurious Tokens

강화 학습(RL)은 대규모 언어 모델(LLM)의 추론 능력을 크게 향상시켰지만, 기존의 RL 미세 조정 방법은 안정성을 유지하기 위해 엔트로피 정규화 및 가중치 재조정과 같은 휴리스틱 기술에 크게 의존합니다. 실제로 이러한 방법들은 종종 훈련 후반 단계에서 성능 저하를 겪으며, 이는 추론 품질 저하 및 불안정한 훈련으로 이어집니다. 우리의 분석에 따르면, RL에서 토큰별 정책 기울기의 크기는 토큰 확률 및 로컬 정책 엔트로피와 음의 상관 관계를 가집니다. 우리는 훈련 불안정성이 전체 토큰의 약 0.01%에 불과한, '오해석 토큰'에 의해 발생할 수 있다는 것을 발견했습니다. 이러한 토큰이 올바른 응답에 나타날 때, 이는 추론 결과에 큰 영향을 미치지 않지만 전체 시퀀스 레벨의 보상을 그대로 받기 때문에, 비정상적으로 증폭된 기울기 업데이트를 초래합니다. 이러한 불안정성을 완화하기 위해, 우리는 낮은 확률, 낮은 엔트로피, 그리고 양의 이점을 갖는 특징 신호를 통해 오해석 토큰을 효율적으로 식별하는 S2T(silencing spurious tokens) 메커니즘을 설계하고, 최적화 과정에서 이러한 토큰의 기울기 변화를 억제합니다. 이 메커니즘을 그룹 기반 목표에 통합하여, 우리는 안정적이고 효과적인 대규모 모델 개선을 위한 '오해석 토큰 인식 정책 최적화(STAPO)'를 제안합니다. Qwen 1.7B, 8B 및 14B 기본 모델을 사용한 6개의 수학적 추론 벤치마크에서, STAPO는 GRPO, 20-Entropy 및 JustRL과 비교하여 엔트로피 안정성 측면에서 일관되게 우수한 성능을 보이며, 평균적으로 7.13% (ρT=1.0, top-p=1.0) 및 3.69% (ρT=0.7, top-p=0.9)의 성능 향상을 달성했습니다.

Reinforcement Learning (RL) has significantly improved large language model reasoning, but existing RL fine-tuning methods rely heavily on heuristic techniques such as entropy regularization and reweighting to maintain stability. In practice, they often suffer from late-stage performance collapse, leading to degraded reasoning quality and unstable training. Our analysis shows that the magnitude of token-wise policy gradients in RL is negatively correlated with token probability and local policy entropy. We find that training instability can be caused by a tiny fraction of tokens, approximately 0.01%, which we term spurious tokens. When such tokens appear in correct responses, they contribute little to the reasoning outcome but inherit the full sequence-level reward, leading to abnormally amplified gradient updates. To mitigate this instability, we design an S2T (silencing spurious tokens) mechanism to efficiently identify spurious tokens through characteristic signals with low probability, low entropy, and positive advantage, and then suppress their gradient perturbations during optimization. Incorporating this mechanism into a group-based objective, we propose Spurious-Token-Aware Policy Optimization (STAPO), which promotes stable and effective large-scale model refinement. Across six mathematical reasoning benchmarks using Qwen 1.7B, 8B, and 14B base models, STAPO consistently demonstrates superior entropy stability and achieves an average performance improvement of 7.13% ($ρ_{\mathrm{T}}$=1.0, top-p=1.0) and 3.69% ($ρ_{\mathrm{T}}$=0.7, top-p=0.9) over GRPO, 20-Entropy, and JustRL.