STAPO: 희귀 허위 토큰 억제를 통한 대규모 언어 모델 강화학습 안정화

초록

강화학습(RL)은 대규모 언어 모델의 추론 능력을 크게 향상시켰지만, 기존 RL 미세 조정 방법은 안정성을 유지하기 위해 엔트로피 정규화 및 재가중치와 같은 경험적 기법에 크게 의존합니다. 실제로 이러한 방법들은 후기 단계 성능 붕괴를 자주 경험하여 추론 품질 저하와 불안정한 학습을 초래합니다. 우리는 RL에서 토큰 단위 정책 그래디언트의 크기가 토큰 확률 및 지역 정책 엔트로피와 음의 상관관계를 가짐을 도출했습니다. 이 결과를 바탕으로 우리는 훈련 불안정성이 전체 토큰의 약 0.01%에 해당하는 극소수 토큰에 의해 주도된다는 것을 증명하며, 이러한 토큰을 허위 토큰(spurious tokens)으로 명명합니다. 이러한 토큰이 정답 응답에 등장할 경우, 이들은 추론 결과에는 거의 기여하지 않지만 시퀀스 수준의 전체 보상을 상속받아 비정상적으로 증폭된 그래디언트 업데이트를 초래합니다. 이러한 관찰에 기반하여, 우리는 대규모 모델 정제를 위한 허위 토큰 인식 정책 최적화(STAPO)를 제안합니다. STAPO는 이러한 업데이트를 선택적으로 마스킹하고 유효한 토큰에 대한 손실을 재정규화합니다. Qwen 1.7B, 8B, 14B 기본 모델을 사용한 6개의 수학적 추론 벤치마크에서 STAPO는 일관되게 우수한 엔트로피 안정성을 보였으며, GRPO, 20-Entropy, JustRL 대비 평균 7.13%의 성능 향상을 달성했습니다.

English

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 experience late-stage performance collapse, leading to degraded reasoning quality and unstable training. We derive that the magnitude of token-wise policy gradients in RL is negatively correlated with token probability and local policy entropy. Building on this result, we prove that training instability is driven 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. Motivated by this observation, we propose Spurious-Token-Aware Policy Optimization (STAPO) for large-scale model refining, which selectively masks such updates and renormalizes the loss over valid tokens. 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\% over GRPO, 20-Entropy and JustRL.

STAPO: 희귀 허위 토큰 억제를 통한 대규모 언어 모델 강화학습 안정화

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

초록

Support