SpanNorm: Reconciling Training Stability and Performance in Deep Transformers

TL;DR

SpanNorm addresses the entrenched instability–performance trade-off in deep Transformer normalization by fusing PreNorm and PostNorm ideas into a block-spanning residual design with block-output normalization. The approach maintains a stable signal propagation through principled variance control and a Scale Init for depth scalability, enabling stable training of ultra-deep models and offering improvements in both dense and MoE settings. Theoretical analyses show SpanNorm bounds gradient dynamics and prevents both vanishing gradients and representation collapse, while extensive experiments demonstrate consistent gains in downstream tasks and resilience to extreme depths. Practically, SpanNorm is a drop-in replacement that preserves computational efficiency and scales to industrial-scale models, potentially enabling more powerful and stable LLMs.

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. Structurally, SpanNorm establishes a clean residual connection that spans the entire transformer block to stabilize signal propagation, while employing a PostNorm-style computation that normalizes the aggregated output to enhance 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 also 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.

Figures (14)

• Figure 1: Comparisons of our proposed SpanNorm with PostNorm, PreNorm, and other advanced variants. Here, we take the dense model as an instance, and MHA denotes the multi-head attention, FFN denotes the feedforward network. Note that MHA can also be replaced by GQA, MLA and other attention variants. When switching to the MoE models, FFN could be replaced by MoE modules.
• Figure 2: Early-stage training stability analysis. We train dense models with a fixed hidden dimension $d=1536$ across increasing depths (12, 24, 48 layers) for 5000 steps to evaluate the impact of initialization on stability.
• Figure 3: Gradient norm dynamics (Analysis of the 24-layer model in Figure \ref{['fig:scale_init']}). We visualize layer-wise gradient norms for the exact 24-layer configuration that exhibited instability. (See Figure \ref{['fig:grad_seed']} for the full visualization of all 24 layers).
• Figure 4: Performance and scalability across scales. (Left) SpanNorm consistently outperforms baselines on a 64-layer 5B dense model (200B tokens). (Middle & Right) On an ultra-deep 128-layer 6.5B model (400B tokens), SpanNorm achieves superior performance with robust convergence, whereas PostNorm and HybridNorm suffer from catastrophic divergence in the early phase (Right).
• Figure 5: Cosine similarity between the outputs of each pair of layers for MoE-A2.4B-16B model.

Theorems & Definitions (2)

• Theorem 4.1: Condition for Stable Depth Scaling
• proof