SpanNorm: Reconciling Training Stability and Performance in Deep Transformers
Abstract
The success of Large Language Models (LLMs) hinges on the stable training of deep Transformer architectures. A critical design choice is the placement of normalization layers, leading to a fundamental trade-off: the ''PreNorm'' architecture ensures training stability at the cost of potential performance degradation in deep models, while the ''PostNorm'' architecture offers strong performance but suffers from severe training instability. In this work, we propose SpanNorm, a novel technique designed to resolve this dilemma by integrating the strengths of both paradigms. SpanNorm adopts the clean residual path of PreNorm to stabilize signal propagation while employing a PostNorm-style computation that normalizes the output of the residual connection, thereby enhancing model performance. We provide a theoretical analysis demonstrating that SpanNorm, combined with a principled scaling strategy, maintains bounded signal variance throughout the network, preventing the gradient issues that plague PostNorm models, and alleviating the representation collapse of PreNorm. Empirically, SpanNorm consistently outperforms standard normalization schemes in both dense and Mixture-of-Experts (MoE) scenarios, paving the way for more powerful and stable Transformer architectures.
Lay Summary
Training deep Large Language Models involves a fundamental trade-off. The standard "PreNorm" architecture ensures stable training but limits performance as networks grow deeper. Conversely, the "PostNorm" approach offers superior performance but suffers from severe instability, making deep models nearly impossible to train. To resolve this dilemma, we introduce SpanNorm. By breaking the architecture down to its core mechanics, we found that simply changing a single connection line in the network block is enough to capture the best of both worlds. This elegant adjustment maintains a clean pathway across the block for training stability, while normalizing the aggregated output at the very end to maximize learning capacity. This hybrid design successfully prevents signals from vanishing or exploding. SpanNorm consistently outperforms existing methods without adding any computational overhead. We successfully trained a massive 128-layer model, proving its robust stability at extreme scales. Ultimately, it provides a reliable, plug-and-play foundation for building more powerful and efficient AI systems.