소규모 언어 모델 에이전트를 위한 안정적인 강화 학습 연구

Towards Robust Reinforcement Learning for Small-Scale Language Model Agents

70~500만 파라미터 범위의 소규모 언어 모델(SLM)을 강화 학습으로 조정하는 것은 종종 불안정하다고 여겨지지만, 그 근본적인 실패 원인은 체계적으로 조사되지 않았습니다. 본 연구에서는 최첨단(SOTA) 연구에서 Proximal Policy Optimization (PPO)를 사용하여 15개의 (모델, 데이터셋) 조합을 학습했습니다. 실험에는 Pythia-70M, 160M, 410M 및 SmolLM2-135M, 360M 모델이 TinyStories, CNN/DailyMail 및 Wikitext-103 데이터셋으로 구성되었습니다. 소규모 언어 모델에서 세 가지 재현 가능한 실패 모드가 확인되었으며, 이는 표준 PEFT/TRL 파이프라인에서의 LoRA 매개변수 동결, bfloat16 사용 시 중요도 비율에서의 숫자 오버플로우 및 보상 모델 오류로 인한 정책 붕괴입니다. 이러한 문제는 어댑터 병합 및 재초기화 기술, PPO 업데이트 중 float32 정밀도 사용 및 보상 조정, 중요도 비율 보호 및 가중치 복구를 포함하는 세 계층 안전 메커니즘을 사용하여 해결되었습니다. 본 논문에서는 용량 여유 가설(capacity-headroom hypothesis)을 제안하며, 이는 SLM 규모에서의 PPO 성능이 모델 파라미터 수보다는 유창한 지도 학습 모델(PPL < 20)과 구별력 있는 보상 신호에 의해 결정된다는 것입니다. 제안된 시스템은 모든 실험에서 안정적으로 수렴했으며, 유창한 사전 지식과 유용한 보상 신호를 가진 조합에서 SFT 기준보다 더 높은 선호도 승률을 달성했습니다. 또한, 훨씬 적은 양의 학습 데이터를 사용하면서 instruction-tuning 기반 모델보다 우수한 성능을 보여주었습니다. 모든 체크포인트, 선호도 데이터셋 및 학습 스크립트는 공개적으로 제공됩니다.$^§$

The alignment of Small Language Models (SLMs) in the 70--500M parameter range using reinforcement learning is often considered unstable, though the underlying failure mechanisms have not been systematically investigated. In the State-of-the-Art (SOTA) research, fifteen (model, corpus) configurations were trained using Proximal Policy Optimization (PPO). The experiments included Pythia-70M, 160M, 410M and SmolLM2-135M, 360M on the TinyStories, CNN/DailyMail, and Wikitext-103 corpora. Three reproducible failure modes were identified in small-scale language models: silent LoRA parameter freezing in standard PEFT/TRL pipelines, numerical overflow in importance ratios when using bfloat16, and catastrophic policy collapse due to reward-model error. These issues were addressed using a merge-and-reinitialize adapter technique, float32 precision during PPO updates, and a three-layer safety mechanism comprising reward whitening, importance-ratio guarding, and weight rollback. In this paper, a capacity-headroom hypothesis is proposed, which states that PPO performance at the SLM scale depends on both a fluent supervised model ($\text{PPL}<20$) and a discriminative reward signal, rather than on the number of model parameters. The proposed system converged stably in all experiments and improved preference win rate over the SFT baseline in configurations with a fluent prior and an informative reward signal. Furthermore, it outperformed instruction-tuned baselines while requiring significantly less training data. All checkpoints, preference datasets, and training scripts are publicly released$^§$.