Don't Let Gains FADE: Breaking Down Policy Gradient Weights in RL
For researchers and practitioners using RL to fine-tune LLMs, this work provides a practical solution to training instability and diversity collapse, with clear improvements over static baselines.
The paper introduces a unifying framework for advantage functions in RL-based LLM post-training, revealing that imbalanced gradient updates cause entropy collapse or weight geometry issues. The proposed FADE method dynamically schedules gradient weights, achieving peak pass@1 20k steps earlier at 7B scale and 2k steps earlier at 32B scale, with best accuracy-diversity trade-off on LiveCodeBench and AIME.
Reinforcement learning post-training dramatically improves LLM reasoning, but suffers from training instability and diversity collapse. Advantage functions offer an appealing fix: they reshape the training objective, reweight which rollouts drive learning, and are trivial to implement. Yet a proliferation of methods makes it unclear which advantage to use and when. We cut through the confusion with a unifying framework that decomposes any advantage into its positive and negative gradient mass along two orthogonal axes. On the sign axis, imbalanced updates collapse either entropy or weight geometry. On the difficulty axis, hard-problem focus sharpens signal but costs sample size. Both trade-offs shift during training: exploration favors balance and hard focus; exploitation favors suppression and medium focus. This motivates FADE (Focal Advantage with Dynamic Entropy), a self-adapting advantage that reads training dynamics to schedule the gradient weight automatically. FADE reaches peak pass@1 20k steps earlier than the best static baseline at the 7B scale and 2k steps earlier at the 32B , while achieving the best accuracy-diversity trade-off across all pass@k on LiveCodeBench and AIME.