Do We Need Adam? Surprisingly Strong and Sparse Reinforcement Learning with SGD in LLMs
Abstract
Reinforcement learning (RL), particularly RL from verifiable reward (RLVR), has become a crucial phase of training large language models (LLMs) and a key focus of current scaling efforts. However, optimization practices in RL largely follow those of next-token-prediction stages (e.g., pretraining and supervised fine-tuning), despite the fundamental differences between RL and these stages emphasized by recent work. One such practice is the use of the AdamW optimizer, which is widely adopted for training large-scale transformers despite its high memory overhead. Our analysis shows that both momentum and adaptive learning rate of AdamW are less influential in RL than in SFT, leading us to hypothesize that RL benefits less from Adam’s per-parameter adaptive learning rates and momentum. Confirming our hypothesis, our experiments demonstrate that the substantially more memory-efficient SGD, which is known to perform poorly in supervised learning of large-scale transformers, matches or even outperforms AdamW in RL for LLMs. Remarkably, full fine-tuning with SGD updates fewer than 0.02% of model without any sparsity-promoting regularization, more than 1,000 times fewer than AdamW. Our analysis offers potential reasons for this update sparsity. Our findings provide fresh insights into the optimization dynamics of RL in LLMs and demonstrate that RL can be substantially more parameter-efficient than previously recognized.
Lay Summary
Reinforcement learning is now a major part of post-training language models. This stage commonly uses AdamW, an optimizer that works well in standard language model training but requires much more memory than simpler alternatives such as SGD. We ask whether AdamW is necessary for reinforcement-learning-based post-training. Contrary to conventional wisdom, we find that SGD can match or outperform AdamW across verifiable tasks. SGD also changes only a tiny fraction of the model’s parameters, suggesting that RL improves models through highly localized updates. These results show that SGD can make RL-based post-training cheaper and more memory-efficient, while also suggesting that optimization principles from standard language model training may not directly apply to reinforcement learning.