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

2025-02-21

AI Summary

To address training instability in post-norm Transformer-based large language models (LLMs) caused by gradient explosion and vanishing, this paper proposes Scale–Distribution Decoupling (SDD): an explicit decomposition of the weight matrix in fully connected layers into a learnable per-dimension scale vector and a normalized distribution component. This constitutes the first structural decoupling of weight parameters with theoretical interpretability and lightweight parameterization. SDD jointly optimizes gradient condition number and activation stability. Experiments demonstrate that SDD significantly improves training stability across diverse LLM architectures and normalization configurations, consistently outperforming baselines including LayerNorm and RMSNorm. Crucially, it incurs zero inference overhead and maintains full compatibility with mainstream training frameworks.

Technology Category

Machine Learning: Large Multimodal Models (LMMs)Natural Language Processing: (Large) Language ModelsComputer Vision: Large Vision Models

Application Category

Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphsUser Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendationSearch and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information

Abstract

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 $ extbf{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.

Problem

• stabilize training of large language models
• prevent gradient explosion and dissipation
• improve optimization efficiency in deep networks

Innovation

• Scale-Distribution Decoupling technique
• Normalizes activations for stability
• Uses learnable scaling vector