Parameter-Adaptive AdamW: A Simple Yet Effective Optimization Strategy for Transformer Language Models

View PDF

Abstract:We present a systematic study of parameter-specific adaptation in the AdamW optimizer for transformer language models. While numerous complex optimizer modifications have been proposed, we demonstrate that careful configuration of AdamW with just two parameter groups - weight matrices versus other parameters - achieves a 3.8\% improvement in validation loss (4.741 vs 4.927) over the standard AdamW baseline on the FineWeb benchmark. Our experiments use the Qwen 3 architecture with 134M parameters, trained for 50,000 steps with a batch size of 256. The method's simplicity, requiring no additional computation over AdamW, makes it practical for widespread adoption. We analyze why this grouping strategy works through gradient distribution studies and ablation tests, showing that weight matrices benefit from higher learning rates (0.001 vs 0.0005) and adjusted momentum ($\beta_2=0.98$ vs 0.999). While not surpassing specialized optimizers like StableAdam (3.888 loss), our approach provides reliable improvements with minimal implementation overhead.