Domando LLMs através da Escalonamento de Taxas de Aprendizado com Agrupamento 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

Resumo

O treinamento de grandes modelos de linguagem (LLMs) apresenta desafios devido à sua escala massiva e arquiteturas heterogêneas. Embora otimizadores adaptativos como o AdamW ajudem a lidar com variações de gradiente, eles ainda lutam com a estimativa eficiente e eficaz das taxas de aprendizado por parâmetro, resultando em instabilidade no treinamento, convergência lenta e baixa compatibilidade com técnicas de ajuste fino eficiente em parâmetros (PEFT). Este trabalho introduz o Scaling with Gradient Grouping (SGG), um wrapper de otimizador que melhora a estimativa adaptativa da taxa de aprendizado por meio de agrupamento dinâmico e escalonamento específico por grupo. O SGG primeiro agrupa as estatísticas de gradiente em cada camada em clusters e, em seguida, aplica um escalonamento específico por cluster para calibrar as taxas de aprendizado para cada parâmetro, impondo assim restrições coletivas por grupo enquanto mantém uma adaptação precisa por parâmetro. Experimentos em diversos benchmarks de (M)LLMs mostram que o SGG se integra perfeitamente com otimizadores existentes e oferece ganhos consistentes e convergência mais rápida em relação às linhas de base, com vários tamanhos de modelo. Sua estabilidade em diferentes tamanhos de lote e taxas de aprendizado estabelece o SGG como uma escolha robusta para a otimização 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.