Taming LLMs by Scaling Learning Rates with Gradient Grouping (2506.01049v1)

Abstract: Training 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.

Scaling Learning Rates in LLMs with Gradient Grouping

The paper "Taming LLMs by Scaling Learning Rates with Gradient Grouping" addresses a critical challenge in the optimization of LLMs: improving the estimation and application of learning rates to achieve stable, fast convergence across heterogeneous model architectures. Traditional adaptive optimizers, despite offering high adaptability per parameter, often fail to accommodate the learning dynamics of LLMs without incurring significant computational overhead or compromising performance under parameter-efficient fine-tuning (PEFT).

Core Contributions and Methodology

The paper introduces the Scaling with Gradient Grouping (SGG) method, which serves as an optimizer wrapper capable of enhancing learning rate efficiency through dynamic grouping and group-specific scaling of gradients. SGG operates on a dual-level optimization strategy:

1. Dynamic Grouping: It clusters gradient statistics in each layer, accommodating the unique optimization behavior inherent to specific model components such as attention heads and MLP layers. This counters the static grouping or pre-defined group-based approaches that often overlook within-layer gradient variations.
2. Group-Specific Scaling: Post-grouping, SGG applies scaling factors tailored to the aggregate behavior of each cluster, aligning them with the broader layer and model-wide gradient trends. This nuanced scaling aids in harmonizing the diverse optimization pathways within LLMs, allowing for retained parameter-wise adaptability without risking group-level uniformity.

Experimental Outcomes

Experiments conducted across various model scales and benchmarks, including general language pre-training and multimodal tasks, demonstrate SGG's effectiveness. Key results highlighted in the paper are as follows:

• Performance Gains: SGG enables consistent performance improvements in models ranging from 60 million to 1 billion parameters when integrated with optimizers like Adam, Adafactor, and LAMB. For instance, in LLaMA's pre-training on C4, SGG-enhanced models show notable perplexity reductions (e.g., perplexity decreased by 1.26% to 3.75% across different model sizes).
• Faster Convergence: Enhanced convergence speeds were observed, with models achieving target performance metrics more swiftly than baseline optimizers. This stability across diverse learning rates and batch sizes suggests reduced sensitivity to these critical hyperparameters, addressing the well-documented 'surge' phenomenon in LLM training.
• Improved Compatibility with PEFT: SGG's ability to maintain or surpass the performance of full-rank training despite fewer parameters underscores its value in resource-constrained environments, as evidenced by its effectiveness with LoRA and other PEFT strategies.

Implications and Future Directions

The paper paves the way for further exploration into adaptive optimization strategies within LLM contexts. The SGG approach illustrates the potential for more intelligently parameterized learning rates that respect the structural and statistical nuances of model components. In practical terms, optimizing computational resources without sacrificing performance is crucial for implementing LLMs in wider industrial applications, particularly where access to large-scale computation might be limited.

Future work might explore alternative clustering and scaling strategies, as the adaptability of SGG to different model regimes hints at broader applicability across tasks beyond those demonstrated. The insights gained from understanding the intra-layer gradient correlations could guide more comprehensive model architecture designs, potentially influencing the fabric of future LLM developmental practices.

In conclusion, the SGG method signifies a meaningful advancement in optimizing LLM training, promoting efficiency without compromising the thoroughness required in high-stakes linguistic and multimodal tasks. The paper promotes a refined understanding of gradient dynamics and leverages this for tangible improvements in model training régimes.