Desacoplamiento de Escala-Distribución: Habilitando el Entrenamiento Estable y Efectivo de Modelos de Lenguaje a Gran Escala

Scale-Distribution Decoupling: Enabling Stable and Effective Training of Large Language Models

February 21, 2025

Autores: Ya Wang, Zhijian Zhuo, Yutao Zeng, Xun Zhou, Jian Yang, Xiaoqing Li

cs.AI

Resumen

La estabilidad del entrenamiento es un desafío persistente en el preentrenamiento de modelos de lenguaje grandes (LLMs, por sus siglas en inglés), particularmente para arquitecturas como los Transformers Post-Norm, que son propensos a la explosión y disipación de gradientes. En este artículo, proponemos Desacoplamiento de Escala y Distribución (SDD, por sus siglas en inglés), un enfoque novedoso que estabiliza el entrenamiento al desacoplar explícitamente la escala y la distribución de la matriz de pesos en capas completamente conectadas. SDD aplica un mecanismo de normalización para regular las activaciones y un vector de escala aprendible para mantener gradientes bien condicionados, previniendo efectivamente la explosión y disipación de gradientes. Esta separación mejora la eficiencia de la optimización, especialmente en redes profundas, al garantizar una propagación estable de los gradientes. Los resultados experimentales demuestran que nuestro método estabiliza el entrenamiento en diversas arquitecturas de LLMs y supera a las técnicas existentes en diferentes configuraciones de normalización. Además, el método propuesto es ligero y compatible con frameworks existentes, lo que lo convierte en una solución práctica para estabilizar el entrenamiento de LLMs. El código está disponible en https://github.com/kaihemo/SDD.

English

Training stability is a persistent challenge in the pre-training of large language models (LLMs), particularly for architectures such as Post-Norm Transformers, which are prone to gradient explosion and dissipation. In this paper, we propose Scale-Distribution Decoupling (SDD), a novel approach that stabilizes training by explicitly decoupling the scale and distribution of the weight matrix in fully-connected layers. SDD applies a normalization mechanism to regulate activations and a learnable scaling vector to maintain well-conditioned gradients, effectively preventing gradient explosion and dissipation. This separation improves optimization efficiency, particularly in deep networks, by ensuring stable gradient propagation. Experimental results demonstrate that our method stabilizes training across various LLM architectures and outperforms existing techniques in different normalization configurations. Furthermore, the proposed method is lightweight and compatible with existing frameworks, making it a practical solution for stabilizing LLM training. Code is available at https://github.com/kaihemo/SDD.