PaperlayerView exact revision on arXiv

Source-linked AI summary

Post-LayerNorm Is Back: Stable, ExpressivE, and Deep

Chen Chen, Lai Wei

arXiv:2601.19895v2cs.LGcs.CL

TL;DR

Conventional LLM scaling is reaching diminishing returns, while depth scaling is theoretically attractive but difficult to optimize reliably at extreme depths. The paper identifies the ResNet-style residual path as the source of Post-LN gradient vanishing and introduces Keel, a Highway-connected Post-LN Transformer. Keel trains robustly beyond 1000 layers and improves depth-scaling behavior relative to Pre-LN, with important gains depending on regime and data availability.

  • Problem

    Conventional scaling axes are plateauing, while depth scaling remains under-exploited because Transformer architectures become unstable or underuse deeper layers at extreme depths.

  • Method

    The paper analyzes Post-LN gradient dynamics and introduces Keel, which replaces the ResNet-style residual path with a Highway-style gated connection.

  • Results

    Keel trains robustly at depths exceeding 1000 layers and delivers consistent improvements and more favorable depth-scaling behavior than Pre-LN baselines.

  • Takeaways & Limitations

    A Highway-style connection can make Post-LN a practical foundation for extreme-depth Transformer scaling without specialized initialization or elaborate training heuristics.

  • Takeaways & Limitations

    The study focuses on depth scaling; wider settings may require stronger stabilization, gains may be smaller at high width-to-depth ratios, and Keel is not recommended in low-data regimes.

Abstract

from arXiv · show

Large language model (LLM) scaling is hitting a wall. Widening models yields diminishing returns, and extending context length does not improve fundamental expressivity. In contrast, depth scaling offers theoretically superior expressivity, yet current Transformer architectures struggle to train reliably at extreme depths. We revisit the Post-LayerNorm (Post-LN) formulation, whose instability at scale caused its replacement by Pre-LN in modern LLMs. We show that the central failure mode of Post-LN arises from the ResNet-style residual pathway, which introduces gradient vanishing in deep networks. We present Keel, a Post-LN Transformer that replaces this residual path with a Highway-style connection. This modification preserves the gradient flow through the residual branch, preventing signal vanishing from the top layers to the bottom. Unlike prior methods, Keel enables stable training at extreme depths without requiring specialized initialization or complex optimization tricks. Keel trains robustly at depths exceeding 1000 layers and consistently improves perplexity and depth-scaling characteristics over Pre-LN. These findings indicate that Post-LN, when paired with a Highway-style connection, provides a simple and effective foundation for building deeply scalable LLMs, opening the possibility for future infinite-depth architectures.

1 Introduction

Conventional LLM scaling axes are yielding diminishing returns, while depth scaling remains difficult because current architectures become unstable and underuse deeper layers. The paper analyzes Post-LN’s failure mode and introduces Keel, a Highway-connected Post-LN Transformer that trains stably at extreme depths.

  • Motivation: Width, context, and parameter scaling are showing diminishing returns, motivating architectural approaches that deliver more expressivity per parameter.Width saturates quickly, context scaling becomes increasingly expensive, and parameter growth alone does not unlock qualitatively new behaviors.
  • Motivation: Depth scaling could provide richer functions and hierarchical reasoning, but current architectures become unstable at extreme depths and yield worse returns than width scaling.
  • Post-LN and Pre-LN: Pre-LN stabilizes early training but weakens gradient propagation and deeper-layer contributions, limiting depth scaling and representational expressivity.
  • Post-LN and Pre-LN: Post-LN preserves large gradients and expressive inter-layer coupling, yet normalization of summed residual and transformed features causes deep-training instability.
  • Keel: Keel replaces the ResNet-style residual branch with a lightweight Highway-style gated connection that regulates forward and backward information flow.The gate dynamically balances carry and transform signals without requiring special initialization or customized residual scaling.
  • Results: Keel trains robustly beyond 1000 layers and maintains smooth convergence under aggressive learning rates, while improving expressiveness across 64–1024 layers.The reported gains are particularly pronounced in Math and Code, where performance increases by +16.5% over Pre-LN baselines.

2 Preliminary

This section reviews Highway and Residual Networks alongside Post-LN, Pre-LN, and hybrid normalization. It emphasizes how residual-path design and normalization placement shape gradient propagation, stability, and depth utilization.

  • Highway Networks: Highway layers use trainable gates to regulate how much input is transformed versus directly carried forward.This bypass can prevent gradient attenuation when transformations would otherwise weaken the signal.
  • Residual Networks: Residual Networks replace the Highway gate with a fixed identity path, but its interaction with normalization determines gradient propagation and activation accumulation across depth.
  • Post-LN: Post-LN normalizes after residual addition, coupling residual and transformed signals and making gradients sensitive to activation statistics and deep-network attenuation.
  • Post-LN: DeepNorm mitigates Post-LN instability with depth-dependent residual scaling and weight-initialization down-scaling, using α=L^0.25 and β=L^-0.25 for decoder-only architectures.
  • Pre-LN: Pre-LN provides a clean identity gradient path and stable optimization without special initialization, but deeper layers often contribute less to updates, harming depth scaling.
  • Hybrid Variants: HybridNorm and Mix-LN combine Post-LN and Pre-LN placements, improving robustness in some regimes without fundamentally resolving gradient degeneration at very large depths.

3 KEEL

Keel is a Post-LN Transformer that stabilizes deep training by scaling and normalizing the shortcut pathway while preserving Post-LN structure.

  • Architecture: Keel introduces a Highway-style scalar α on the shortcut connection and an additional Layer Normalization before the residual function.These are the two critical structural modifications relative to vanilla Post-LN.
  • Architecture: α = L, where L is the total number of Attention and FFN sub-layers, is the recommended setting for very deep models.For smaller models, α > 1 can instead be tuned as a hyperparameter.
  • Architecture: The additional normalization stabilizes gradient flow through the residual branch, preventing attenuation while retaining the global Post-LN topology.Keel is classified as Post-LN because its shortcut carries a normalized signal, even though the transformation input is also normalized.
  • Implementation: Keel removes selected normalization and shortcut scaling operations in the first attention and FFN layers to stabilize signal initialization.The implementation also uses larger learning rates than standard Pre-LN baselines and omits additive LayerNorm bias.
  • Gradient analysis: Standard Post-LN gradients exponentially decay toward lower layers, whereas Keel’s formulation prevents gradient vanishing in the deep limit.The Keel derivation introduces α and an additional normalization step to alter the cumulative gradient behavior.

4 Discussions

The discussion contrasts Keel’s gradient-centered design with DeepNorm and connects depth-wise propagation to sequence-wise recurrence. It argues that Keel’s topology and normalization provide a stable signal for very deep LLMs.

  • Comparison with DeepNorm: Keel targets vanishing lower-layer gradients, whereas DeepNorm primarily bounds forward output magnitude as depth increases.Keel’s analysis reports a gradient norm of approximately 1 regardless of depth L.
  • Comparison with DeepNorm: DeepNorm depends on specialized L−0.25 initialization, while Keel’s design is motivated by gradient dynamics rather than initialization constraints.The discussion notes that initialization benefits can diminish as massive pre-training moves weights away from their initial state.
  • Design evolution: Naive shortcut scaling improved stability over vanilla DeepNorm but remained below standard Pre-LN in training stability and final convergence.The authors attribute this result to high variance from the unnormalized input to the residual function.
  • Design evolution: Learnable input scaling improved forward stability but attenuated gradients as β decreased, leaving stability inferior to Pre-LN.The later normalization step was introduced to decouple variance control from gradient scaling.
  • Design evolution: Adding normalization before β preserves the gradient magnitude through F, with affine scaling terms cancelling the gradient-vanishing effect observed earlier.The final Keel formulation absorbs the external β into LayerNorm’s internal affine weights.
  • Depth-wise recurrence: The paper identifies a structural isomorphism between sequence-wise recurrence and depth-wise propagation in residual networks.It interprets Transformer layer updates as optimization steps along the depth dimension.
  • Depth-wise recurrence: Keel’s shortcut scaling and residual normalization are presented as counterparts to gating and state normalization used to stabilize recurrent test-time-training objectives.The discussion connects these mechanisms to sequence-modeling systems such as Titans and LaCT.

5 Experiments

Experiments evaluate Keel against Pre-LN and other normalization baselines for stability, learning-rate robustness, depth scaling, data scaling, and downstream performance. Across these settings, Keel supports more stable optimization and increasingly stronger results at greater depths, especially on reasoning tasks.

  • Stability Analysis: Keel’s stability is measured by Maximum Tolerable Learning Rate, the highest learning rate sustainable during warm-up without divergence.Higher Max LR indicates a more robust optimization landscape and superior training stability.
  • Stability Analysis: 1.01 × 10−2: Keel’s 64-layer Max LR exceeds Pre-LN’s 7.65 × 10−3; at 512 layers, Keel sustains 6.31 × 10−3.Keel significantly improves stability relative to the evaluated baselines, including at extreme depth.
  • Optimal Learning Rate: At η = 6.0 × 10−3, Keel improves monotonically across benchmarks, reaching a global average of 55.5 versus 52.3 for the best Pre-LN configuration.Pre-LN instead shows saturation and task-specific degradation, including ARC-Easy and MBPP.
  • Scalability Analysis: Performance at Depth Scaling: At 1024 layers, Keel scores 58.6 on GSM-8K versus 49.8 for Pre-LN, a gap of 8.8 points; Pre-LN stagnates across several reasoning tasks.Keel’s global-average gains over Pre-LN are approximately +1.7, +1.2, +3.8, and +3.0 points at 64, 128, 512, and 1024 layers.
  • Deeper vs. Wider: Under a fixed 3B-parameter budget, the deep-and-narrow model shows an advantage on complex reasoning benchmarks, while training loss is not necessarily correlated with end-task performance.Keel’s 512-layer model reaches an average score of 55.5, outperforming both deep and wide Pre-LN baselines.
  • Main Results: After 1T-token pre-training, Keel improves the global average by +3.8 points over Pre-LN and gains nearly +10 points on GSM-8K; these reasoning gains persist after SFT.After fine-tuning, the GSM-8K gap remains over 10 points, and MMLU-Pro scores are 35.6 versus 26.6.

6 Conclusion

The paper identifies gradient vanishing from the ResNet-style residual pathway as Post-LN’s primary breakdown mechanism and introduces Keel, which replaces it with a Highway-style connection. Keel enables stable optimization beyond 1000 sub-layers without specialized initialization or elaborate training heuristics, improving depth scaling over Pre-LN.

  • Keel replaces the ResNet-style residual path with a Highway-style connection that preserves gradient flow through the residual branch.The change directly targets depth-induced gradient vanishing while retaining Post-LN’s inter-layer coupling.
  • Keel enables stable optimization at depths exceeding 1000 sub-layers without specialized initialization schemes or elaborate training heuristics.
  • Keel delivers consistent improvements and more favorable depth-scaling behavior compared with widely used Pre-LN baselines.

7 Limitation and Future Work

The work focuses on stable Post-LN training under depth scaling, while leaving width-scaling stability for future investigation. Keel’s benefits may be smaller when width already dominates depth, and it is not recommended in low-data regimes.

  • The study focuses on stabilizing Post-LN LLMs as depth scales, leaving stability under width scaling for future work.Wider settings may require a larger α or stronger stabilization mechanisms.
  • Keel’s gains may be less substantial when the width-to-depth ratio is already high, and the method typically requires substantial training data.The paper therefore does not recommend Keel for low-data regimes.

A Layer Redundancy in Deep LLMs

Layer-removal experiments show that the first few layers are exceptionally critical, while degradation generally increases with depth. Compared with Pre-LN, Keel reduces shallow-layer redundancy and gives deeper layers greater effective importance.

  • The experiments train Pre-LN and Keel on the same 250B tokens and quantify layer redundancy through performance degradation after removing each layer.The metric is ΔPPL_i = PPL_remove i-th layer − PPL_full model.
  • Removing one of the first few layers raises perplexity above 10^4 and destroys model performance.
  • Compared with Pre-LN, removing shallow Keel layers produces a larger PPL increase, indicating reduced shallow-layer redundancy.

B Discrepancy Between Training Loss and Downstream Evaluation

Training loss is not always a reliable proxy for downstream performance in deep LLMs. Both depth-structure comparisons and a Pre-LN–Keel comparison show that lower training loss can coincide with worse end-task results.

  • A shallow-wide Pre-LN model trained on the same 250 billion tokens has lower training loss but performs slightly worse downstream than a deep-narrow model.The models use hidden size 2048 with 64 layers versus hidden size 1024 with 256 layers.
  • At a peak learning rate of 1.5 × 10^-3, Pre-LN has slightly lower training loss than Keel, yet Keel performs better on downstream tasks.
  • In deep LLM training, training loss and end-task performance are not always positively correlated.

C Model Configuration

Table 8 presents the model configurations of Keel and the Pre-LN baselines.

  • Table 8 lists the model configurations for Keel and the Pre-LN baselines.
Loading 2601.19895v2…