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

TL;DR

This work shows that depth is a viable axis for scaling LLM expressivity if training remains stable. By diagnosing Post-LayerNorm instability as a gradient-vanishing issue tied to the ResNet-style residual path, the authors propose Keel—a Post-LN Transformer augmented with a Highway-style residual gate and residual-branch normalization. The resulting architecture enables stable optimization beyond 1000 layers and yields consistent gains in depth-scalability and reasoning tasks, outperforming Pre-LN baselines across depth and data scales. The findings suggest a practical path toward deeper, more expressive LLMs and motivate exploration of infinite-depth architectures in future work.

Abstract

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.

Figures (8)

• Figure 1: KEEL enables stable, expressive, and deep LLM training. (a) Training Stability: Keel maintains smooth convergence at aggressive learning rates, while Pre-LN exhibits severe instability under the same configuration. (b) Expressiveness: Keel demonstrates superior performance across all capability domains, particularly in Math & Code (+16.5%). (c) Depth Scaling: Keel consistently outperforms Pre-LN across all depths (64-1024 layers). Together, these results demonstrate that Keel's architectural improvements enable stable optimization of ultra-deep networks with enhanced learning efficiency and model expressiveness.
• Figure 2: Illustration of our Keel architecture.
• Figure 3: The illustration of three distinct pathological behaviors of model divergence in the training of LLMs.
• Figure 4: Training loss curves of Pre-LN during the early stage of training. Pre-LN exhibits a pronounced loss spike when trained with a higher learning rate.
• Figure 5: Training loss curves of Pre-LN and Keel on FineWeb-EDU dataset fineweb with varying training tokens. As the training token scales from 10B to 40B tokens, Keel achieves larger gain on training loss compared to Pre-LN baseline.