Reproduce “Towards Robust Reinforcement Learning for Small-Scale Language Model Agents” on OpenResearch

About this paper

Md Rezwanul Haque, Md. Milon Islam, Fakhri Karray

Researchers at the University of Waterloo and Khulna University of Engineering & Technology identified and resolved key instability issues in applying Proximal Policy Optimization (PPO) to Small Language Models (SLMs) ranging from 70 to 500 million parameters. Their multi-layer safety framework enabled robust PPO alignment, allowing SLMs to achieve competitive performance against larger instruction-tuned baselines with reduced training data, especially when starting with a fluent Supervised Fine-Tuning (SFT) prior.

Key insights

• PPO's instability for SLMs is not inherent but arises from specific, reproducible engineering failure modes within standard PEFT/TRL pipelines.
• The "capacity-headroom hypothesis" proposes that PPO's effectiveness in SLMs depends on a fluent Supervised Fine-Tuning (SFT) prior (Perplexity < 20) and a discriminative reward signal, rather than just model parameter count.
• Each component of the proposed safety framework (merge-and-reinitialize, float32 precision, and a three-layer safety mechanism for distributional collapse) is crucial for addressing distinct failure mechanisms and enabling stable PPO training.

Summary by alphaXiv.