Domestiquer les LLM en ajustant les taux d'apprentissage par regroupement de gradients

Taming LLMs by Scaling Learning Rates with Gradient Grouping

June 1, 2025

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

cs.AI

Résumé

L'entraînement de grands modèles de langage (LLM) présente des défis en raison de leur échelle massive et de leurs architectures hétérogènes. Bien que les optimiseurs adaptatifs comme AdamW aident à gérer les variations de gradient, ils peinent encore à estimer efficacement et précisément les taux d'apprentissage paramètre par paramètre, ce qui entraîne une instabilité lors de l'entraînement, une convergence lente et une faible compatibilité avec les techniques de réglage fin à paramètres efficaces (PEFT). Ce travail introduit Scaling with Gradient Grouping (SGG), un wrapper d'optimiseur qui améliore l'estimation des taux d'apprentissage adaptatifs grâce à un regroupement dynamique et à une mise à l'échelle spécifique à chaque groupe. SGG regroupe d'abord les statistiques de gradient dans chaque couche en clusters, puis applique une mise à l'échelle spécifique à chaque cluster pour calibrer les taux d'apprentissage pour chaque paramètre, imposant ainsi des contraintes collectives au niveau des groupes tout en maintenant une adaptation précise paramètre par paramètre. Les expériences sur divers benchmarks de (M)LLM montrent que SGG s'intègre parfaitement aux optimiseurs existants et offre des gains constants ainsi qu'une convergence plus rapide par rapport aux méthodes de référence, pour différentes tailles de modèles. Sa stabilité face à des tailles de lots et des taux d'apprentissage variables établit SGG comme un choix robuste pour l'optimisation des LLM.

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.