Domando los LLMs mediante la Escalación de Tasas de Aprendizaje con Agrupación de Gradientes

Taming LLMs by Scaling Learning Rates with Gradient Grouping

June 1, 2025

Autores: Siyuan Li, Juanxi Tian, Zedong Wang, Xin Jin, Zicheng Liu, Wentao Zhang, Dan Xu

cs.AI

Resumen

El entrenamiento de modelos de lenguaje a gran escala (LLMs, por sus siglas en inglés) presenta desafíos debido a su escala masiva y arquitecturas heterogéneas. Aunque optimizadores adaptativos como AdamW ayudan a abordar las variaciones en los gradientes, aún luchan con la estimación eficiente y efectiva de las tasas de aprendizaje por parámetro, lo que resulta en inestabilidad durante el entrenamiento, convergencia lenta y poca compatibilidad con técnicas de ajuste fino eficiente en parámetros (PEFT, por sus siglas en inglés). Este trabajo introduce Scaling with Gradient Grouping (SGG), un envoltorio de optimización que mejora la estimación de tasas de aprendizaje adaptativas mediante agrupación dinámica y escalado específico por grupo. SGG primero agrupa las estadísticas de gradiente en cada capa en clústeres y luego aplica un escalado específico por clúster para calibrar las tasas de aprendizaje para cada parámetro, imponiendo así restricciones colectivas a nivel de grupo mientras mantiene una adaptación precisa por parámetro. Los experimentos en diversos puntos de referencia de (M)LLM muestran que SGG se integra perfectamente con optimizadores existentes y ofrece ganancias consistentes y una convergencia más rápida en comparación con las líneas base, independientemente del tamaño del modelo. Su estabilidad frente a diferentes tamaños de lotes y tasas de aprendizaje establece a SGG como una opción robusta para la optimización de LLMs.

English

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.