Taming LLMs by Scaling Learning Rates with Gradient Grouping
Journal:
arXiv
Published Date:
Jun 1, 2025
Abstract
Training large language models (LLMs) poses challenges due to their massive
scale and heterogeneous architectures. While adaptive optimizers like AdamW
help address gradient variations, they still struggle with efficient and
effective parameter-wise learning rate estimation, resulting in training
instability, slow convergence, and poor compatibility with parameter-efficient
fine-tuning (PEFT) techniques. This work introduces Scaling with Gradient
Grouping (SGG), an optimizer wrapper that improves adaptive learning rate
estimation by dynamic grouping and group-specific scaling. SGG first groups
gradient statistics in each layer into clusters and then applies
cluster-specific scaling to calibrate learning rates for each parameter, thus
imposing collective group-wise constraints while maintaining precise
per-parameter adaptation. Experiments on diverse (M)LLM benchmarks show that
SGG integrates seamlessly with existing optimizers, and offers consistent gains
and faster convergence over baselines, with various model sizes. Its stability
across varying batch sizes and learning rates establishes SGG as a robust
choice for LLM optimization.
