StableAdamW: When Stability Hurts Language Model Optimization

View PDF

Abstract:This paper presents a detailed investigation of StableAdamW, a modified version of the AdamW optimizer designed to improve training stability in language models through variance stabilization and enhanced learning rate warmup. Despite promising theoretical foundations and positive results in initial ablation studies, StableAdamW failed to outperform the standard AdamW baseline on the FineWeb benchmark (5.088 vs 4.927 validation loss). Through extensive experiments and analysis, we identify several key factors contributing to this underperformance, including over-constrained optimization dynamics and reduced adaptability to sharp minima. Our findings challenge the common assumption that increased training stability necessarily leads to better model performance. The paper provides valuable insights into the delicate balance between stability and convergence in language model optimization, offering guidance for future optimizer development. We conclude that while stability improvements can be beneficial in certain contexts, they may come at the cost of reduced model performance in language model training.