ArXiv: 2601.19895
🎯 Pitch
A 512-layer Post-LN Transformer trained for 1T tokens scores 60.9 on GSM-8K—a full 10 points higher than its Pre-LN twin—simply by swapping the residual connection for a gated highway. The same design trains stably past 1000 layers and tolerates learning rates that would crash standard architectures, finally unlocking depth as a viable scaling axis.
1. Executive Summary
This paper introduces Keel, a Post-LayerNorm Transformer architecture that replaces the standard ResNet-style residual connection with a Highway-style gated pathway to eliminate the gradient vanishing that has made Post-LN unstable at scale. The authors train deep decoder-only LLMs—ranging from 64 to 1024 layers at up to 3B parameters—on both private data (up to 1T tokens) and the FineWeb-EDU dataset, comparing Keel against Pre-LN, vanilla Post-LN, DeepNorm, HybridNorm, and Mix-LN baselines. Keel enables stable training at depths exceeding 1000 layers while tolerating substantially higher learning rates (e.g., a maximum tolerable LR of 1.01 × 10⁻² vs. 7.65 × 10⁻³ for Pre-LN at 64 layers, and 6.31 × 10⁻³ vs. 4.67 × 10⁻³ at 512 layers), yielding consistent downstream gains that grow with depth—most prominently a +16.5% relative improvement on Math & Code benchmarks and a +10-point absolute gain on GSM-8K (60.9 vs. 51.0) at 512 layers after 1T-token pretraining. The paper establishes that Post-LN architectures can be rehabilitated for extreme-depth LLM training through structural gradient-flow preservation, but only when the residual branch is explicitly designed to prevent signal attenuation rather than relying on initialization tricks or hybrid normalization placement.
2. Context and Motivation
The Core Problem: Depth Scaling Is Theoretically Promising but Practically Broken
The paper addresses a fundamental tension in LLM architecture design: depth scaling offers theoretically superior expressivity compared to width scaling, yet current Transformer architectures cannot train reliably at extreme depths. This is not merely an engineering inconvenience — it represents a hard ceiling on one of the most promising axes for improving model capability per parameter.
The theoretical case for depth is well-established in the deep learning literature. Deeper networks can represent exponentially richer function classes and support more hierarchical reasoning, because each additional layer can compose with all previous layers to form progressively more abstract representations. In the context of Transformers, this means that a deeper model with the same parameter count as a wider model should, in principle, be able to learn more complex patterns — particularly those requiring multi-step logical deduction, algorithmic reasoning, or deep chains of dependencies.
However, the paper observes that current LLMs fail to capitalize on this theoretical advantage. The opening paragraph of Section 1 states this directly:
"Width scaling saturates quickly, context scaling grows increasingly expensive, and parameter growth alone does not unlock qualitatively new behaviors. As a result, LLM scaling is hitting a wall, and there is increasing interest in architectural directions that can deliver more expressivity per parameter."
The critical empirical observation is that "depth scaling delivers substantially worse returns than width scaling under current architectures." This is the gap the paper seeks to close. The authors are not proposing depth scaling as an alternative to width scaling — they are arguing that the poor returns from depth are an artifact of architectural instability, not a fundamental limitation of depth itself. If the instability can be fixed, depth becomes a viable axis for improving model quality without increasing parameter count.
Why This Problem Matters: The End of Easy Scaling
The importance of this problem is anchored in the broader context of LLM scaling hitting diminishing returns across conventional axes:
- Width scaling saturation: The paper notes that "width scaling saturates quickly." This aligns with empirical observations across the field — making models wider (larger hidden dimensions, more attention heads) produces logarithmic returns in loss at exponential cost in parameters and FLOPs. Once a model reaches a certain width, additional parameters in the same layer contribute little to representational capacity.
- Context scaling expense: Extending context windows improves the model's ability to process long-range dependencies but does not change the model's fundamental expressivity per token. It also incurs quadratic (or near-quadratic) computational costs in standard attention mechanisms, making it an increasingly expensive path to marginal capability improvements.
- Parameter growth without new behaviors: The paper asserts that "parameter growth alone does not unlock qualitatively new behaviors." This is a stronger claim — that beyond a certain scale, adding parameters to existing architectures produces incremental improvements on existing benchmarks but does not enable emergent capabilities in reasoning, planning, or abstraction that require fundamentally different architectural support.
Together, these observations paint a picture of a field approaching the limits of what current architectures can deliver through brute-force scaling. Depth represents one of the few remaining degrees of freedom that could break through this ceiling — but only if the optimization challenges can be overcome.
The practical implications extend beyond academic research. If depth scaling can be made reliable, organizations can build models with superior reasoning capabilities without increasing the total parameter count (and thus inference cost). This is the "expressivity per parameter" argument: deeper, narrower networks can potentially match or exceed the performance of wider networks at lower total computational cost per forward pass, since inference FLOPs scale roughly linearly with parameters but depth adds only sequential computation rather than more computation per token.
The Layer Normalization Placement Debate
The paper frames its contribution within a specific lineage of architectural choices: the placement of Layer Normalization relative to residual connections. This seemingly minor detail has an outsized impact on training dynamics.
The original Post-LN formulation. The Transformer, as introduced by Vaswani et al. (2017), placed LayerNorm after the residual addition in each sublayer. In this Post-LN configuration, the output of each sublayer is:
This means that both the residual (identity) signal and the transformed signal pass through normalization together. This coupling has a crucial consequence: gradients flowing backward must pass through the Jacobian of LayerNorm applied to the sum of two potentially misaligned signals. As the paper's analysis shows (Section 3.2, Equation 15), this causes the gradient magnitude to decay exponentially with depth — scaling as —which means lower layers receive vanishing signals during backpropagation in deep networks.
This gradient vanishing is not a minor effect. At 512 layers, the cumulative gradient attenuation means that the earliest layers receive effectively no training signal, regardless of how large the loss is at the output. The model cannot learn representations in its bottom layers, functionally reducing its effective depth to something much smaller than its architectural depth.
The Pre-LN solution and its hidden cost. The Pre-LN formulation (Xiong et al., 2020) addressed this instability by inverting the normalization placement:
Here, the residual path bypasses normalization entirely. The gradient through the residual connection is simply — no attenuation from LayerNorm's Jacobian. This enables stable training at scale without specialized initialization, which is why Pre-LN became the de facto standard in modern LLMs (GPT-3, PaLM, LLaMA, etc.).
However, the paper argues that Pre-LN introduces its own structural limitation: reduced depth utilization. Because gradients flow predominantly through the identity connection, deeper layers contribute diminishingly to the update signal. The paper cites prior work (sun2025curse) for this claim:
"gradients flow predominantly through the identity connection, deeper layers often contribute diminishingly to the update signal, reducing depth utilization and harming scaling behavior"
The consequence is that Pre-LN models do not fully leverage their architectural depth. Shallow layers are effectively "redundant" — they can be removed with minimal impact on model performance — because the gradient signal preferentially updates layers closer to the output. This is empirically demonstrated in the paper's layer redundancy analysis (Section 8). When individual layers are removed from Pre-LN models, shallow-layer removal causes minimal perplexity increase, while deep-layer removal causes larger degradation. This asymmetry indicates that Pre-LN models behave as if they were effectively shallower than their architectural depth.
The paper's key insight is that this tradeoff — between stability (Pre-LN) and depth utilization (Post-LN) — is not inherent but is an artifact of the ResNet-style residual connection. By changing the residual pathway itself, both properties can be achieved simultaneously.
Where Prior Approaches Fall Short
The paper identifies specific shortcomings in existing methods that have attempted to address deep Transformer training:
DeepNorm (Wang et al., 2022). DeepNorm introduced two modifications to Post-LN: (1) a depth-dependent scaling factor on the residual branch, and (2) a corresponding down-scaling of weight initialization by . The forward pass becomes . The theoretical motivation was to bound the forward output magnitude to prevent explosion as depth increases. However, the paper identifies two critical limitations:
First, the theoretical framing is wrong for the actual failure mode. DeepNorm targets forward-pass variance control, but the paper's gradient analysis (Section 3.2) shows that the primary instability mechanism in Post-LN is backward gradient vanishing, not forward variance explosion. Bounding the forward pass does not inherently guarantee a healthy backward gradient flow. Equation 15 shows that even with bounded forward activations, the gradient magnitude through standard Post-LN still decays as .
Second, reliance on initialization creates fragility at scale. DeepNorm's weight initialization scaling () creates a strong dependency on the initial state of the network. In the context of LLM pretraining at trillion-token scale, the model weights drift substantially from their initialization. The paper argues that "as the training dynamics evolve, the benefits of a specific initialization strategy often diminish, leading to instability in later stages of training." The empirical evidence supports this: the paper reports that "while DeepNorm is effective for standard encoder-decoder architectures, its stability and performance degrade when applied to large-scale decoder-only LLMs, often underperforming standard Pre-LN baselines" (Section 4.1).
Admin (Liu et al., 2020). Admin proposed initializing Post-LN models with a "profiling" phase followed by specialized parameter initialization. Like DeepNorm, it treats the instability as an initialization problem rather than a structural one, and suffers from the same limitation: initialization-dependent benefits that degrade over long training trajectories.
Hybrid normalization strategies (HybridNorm, Mix-LN). These approaches interleave Post-LN and Pre-LN blocks (HybridNorm) or transition from Post-LN in early layers to Pre-LN in later layers (Mix-LN). The idea is to get the representational benefits of Post-LN's coupling where it matters (lower layers, which in standard Pre-LN receive the weakest gradients) while retaining Pre-LN's stability elsewhere.
The paper acknowledges that these hybrid designs "provide improved robustness over pure Post-LN and can outperform pure Pre-LN in certain regimes." However, they do not fundamentally resolve the gradient degeneration problem because they still use the same ResNet-style residual connection in their Post-LN blocks. The underlying mechanism of gradient vanishing through the LayerNorm Jacobian (Equation 15) remains active wherever Post-LN blocks appear, limiting how deep those blocks can be stacked before the attenuation becomes prohibitive.
The paper concludes that prior work has addressed symptoms (output magnitude explosion through initialization, placement placement through hybridization) rather than the root cause (the ResNet-style residual pathway's interaction with LayerNorm in the backward pass). This sets up the paper's central claim: by changing the residual connection itself to a Highway-style pathway, the gradient vanishing can be eliminated structurally, without depending on initialization or hybrid placement.
How This Paper Positions Itself
The paper frames its contribution as a return to Post-LN with a structural fix, not an intermediate hybrid solution or an initialization-dependent workaround. The positioning is explicit and confident — the title itself declares "Post-LayerNorm Is Back."
The key distinction from prior work is that Keel addresses the gradient problem at its source: the residual branch. By introducing (1) a scalar scaling factor on the residual connection (a Highway-style gate without the usual constraint, relying on the final LayerNorm for output magnitude) and (2) an additional LayerNorm on the input to the transformation branch, the gradient magnitude through each layer becomes:
This is the paper's central theoretical result (Equation 19, Section 3.3). With , the cumulative gradient product does not decay at all as — it converges to 1. This means the training signal reaches the earliest layers at full strength regardless of architectural depth. The gradient vanishing that plagues Post-LN (Equation 15, ) is eliminated not by avoiding LayerNorm in the residual path (Pre-LN's approach) but by restructuring the residual path itself so that LayerNorm's Jacobian no longer causes attenuation.
The paper carefully argues that Keel is still a Post-LN architecture despite the additional normalization in the transformation branch. The classification depends on the shortcut branch:
"the distinction between Pre-LN and Post-LN depends on the shortcut branch, not the input to the transformation function."
In Keel, the shortcut carries , where is the output of the previous layer's LayerNorm — a normalized signal. This matches the Post-LN structural definition (the shortcut is normalized), distinguishing it from Pre-LN (where the shortcut carries the unnormalized signal from the previous layer).
The paper positions Keel not as a compromise between Post-LN and Pre-LN but as Post-LN done correctly — a rehabilitation of the original Transformer's normalization placement by fixing the one component (the residual connection) that caused it to fail at scale.
This positioning has important implications for the field. If the paper is correct that Post-LN's expressivity advantages (large gradients in deeper layers, strong inter-layer coupling) can be preserved while eliminating its instability, then the architectural landscape shifts: Pre-LN's dominance was not because normalization-before-residual is intrinsically better, but because it was a practical workaround for a fixable structural flaw in Post-LN. Keel suggests that future architectures can return to Post-LN's more expressive coupling while training at scales that were previously only achievable with Pre-LN.
3. Technical Approach
3.1 Reader Orientation
The paper proposes Keel, a modified Transformer architecture that changes how information flows through the residual connections between layers. Keel solves the problem that standard Post-LayerNorm Transformers cannot train at extreme depths (hundreds or thousands of layers) because gradients vanish as they propagate backward through the network—replacing the standard residual connection with a Highway-style gated pathway eliminates this vanishing while preserving Post-LN's stronger inter-layer coupling that makes deep layers more expressive.
3.2 Big-Picture Architecture (Diagram in Words)
The Keel architecture modifies a single Transformer sublayer (attention or feed-forward). A standard Post-LN sublayer takes an input $\mathbf{x}_l$, applies a transformation $\mathcal{F}_l$ (attention or FFN), sums the input and transformed output, and normalizes the result. Keel changes this in three ways:
-
Highway-style scalar gate (
$\alpha$): The input$\mathbf{x}_l$is multiplied by a scalar$\alpha = L$(where$L$is the total number of sublayers) before being added to the transformed branch. This weights the residual pathway so that the gradient through it does not decay with depth. -
Input normalization to the transformation: An additional LayerNorm is placed on the input before it enters
$\mathcal{F}_l$. This stabilizes the variance of the transformed branch's input, preventing the transformation function from receiving signals at uncontrolled scales. -
Output normalization (Post-LN topology): The sum
$\alpha \mathbf{x}_l + \mathcal{F}_l(\text{LN}(\mathbf{x}_l))$is passed through a final LayerNorm, producing$\mathbf{x}_{l+1}$. This makes Keel structurally Post-LN—the shortcut branch carries a normalized signal (since$\mathbf{x}_l$itself was the output of the previous layer's LayerNorm).
The overall forward pass for layer $l$ is:
Information flows as follows: input $\mathbf{x}_l$ (already normalized from the previous layer) splits into two paths—one is scaled by $\alpha$, the other passes through a fresh LayerNorm and then through the attention or FFN transformation. The two branches are summed, then normalized to produce $\mathbf{x}_{l+1}$. The very first attention and FFN layers omit the final LayerNorm and the $\alpha$ scaling, degrading to standard Pre-LN blocks to ensure stable signal initialization from the embedding layer.
3.3 Roadmap for the Deep Dive
- First, the forward pass of Keel in full detail—the two structural modifications (Highway scaling and input normalization), how they differ from vanilla Post-LN, and why the
$\alpha = L$setting is critical for gradient preservation. - Second, the gradient analysis that motivates the architecture—why standard Post-LN causes exponential gradient vanishing (
$\mathcal{O}(2^{-L/2})$), how Keel's modifications change the gradient dynamics to prevent this vanishing (converging to 1 as$L \to \infty$), and what assumptions are embedded in this analysis. - Third, the argument that Keel is genuinely a Post-LN architecture despite the input normalization—the structural criterion (shortcut branch normalization), the non-redundancy of successive LayerNorms, and why this classification matters for understanding the architecture's properties.
- Fourth, the relationship to DeepNorm—how Keel's gradient-flow framing differs from DeepNorm's forward-magnitude framing, and why structural stabilization outperforms initialization-dependent stabilization in large-scale LLM training.
- Fifth, the design evolution from naive residual scaling to the final Keel formulation—the three intermediate attempts, what each revealed about the gradient dynamics, and how redundant parameters were systematically eliminated to arrive at the minimal formulation.
3.4 Detailed, Sentence-Based Technical Breakdown
This is primarily an architectural innovation paper whose core idea is that the instability of Post-LN Transformers at scale is caused specifically by gradient vanishing through the ResNet-style residual branch interacting with LayerNorm's Jacobian, and that replacing this residual branch with a Highway-style gated connection eliminates the vanishing while preserving Post-LN's expressivity advantages.
The Keel Forward Pass
The Keel architecture modifies a single Transformer sublayer (which can be either an attention sublayer or a feed-forward sublayer—both use the same structural modification). The forward propagation for the $l$-th sublayer is defined by Equation 8:
where $\mathbf{x}_l \in \mathbb{R}^d$ is the input activation vector from the previous sublayer, $\alpha$ is a scalar weight on the residual (shortcut) connection, $\mathcal{F}_l$ is the transformation function for this sublayer (either multi-head attention or a feed-forward network), and $\text{LN}(\cdot)$ denotes RMS-based Layer Normalization.
What it computes: the sublayer takes the incoming activation $\mathbf{x}_l$, sends it through two parallel pathways, and recombines them. The left pathway scales $\mathbf{x}_l$ by $\alpha$ and passes it straight through (the "carry" or "skip" path). The right pathway applies a fresh Layer Normalization to $\mathbf{x}_l$, then feeds the normalized result through the sublayer's main transformation $\mathcal{F}_l$ (attention or FFN), producing a transformed output. The two pathway outputs are added together element-wise, and the sum is passed through a final Layer Normalization to produce $\mathbf{x}_{l+1}$. The result $\mathbf{x}_{l+1}$ is a normalized activation vector of the same dimension as the input, ready to be fed into the next sublayer.
Why this form: there are two deliberate deviations from vanilla Post-LN ($\mathbf{x}_{l+1} = \text{LN}(\mathbf{x}_l + \mathcal{F}_l(\mathbf{x}_l))$). First, the scalar $\alpha$ weights the residual pathway more heavily. The paper sets $\alpha = L$ where $L$ is the total number of sublayers (both attention and FFN layers counted together—so a 512-layer model with alternating attention and FFN means $L = 1024$ sublayers). This $\alpha$ acts as a Highway-style gate, but with a critical difference from standard Highway networks: the paper notes explicitly that "unlike standard gating mechanisms that require coefficients to sum to 1 (e.g., $(1-\lambda)\mathbf{x} + \lambda\mathcal{F}(\mathbf{x})$), we rely on the final Post-LN to normalize the output magnitude, rendering explicit variance constraints on the summation unnecessary." In other words, because the final LayerNorm will rescale the sum to unit norm (up to the learned affine weights), the absolute magnitudes of the two branches before summation do not need to be constrained—they just need to be in the right ratio for gradient propagation. The $\alpha = L$ choice ensures that the residual signal dominates the sum at initialization, which is what the gradient analysis shows is necessary to prevent vanishing.
Second, the additional LayerNorm on the input to $\mathcal{F}_l$ normalizes the signal before it enters the transformation. In standard Post-LN, $\mathcal{F}_l$ receives $\mathbf{x}_l$ directly, which is already normalized (since it is the output of the previous layer's LayerNorm). So why add another LayerNorm? The paper's design evolution (Section 4.2) reveals the answer: between the previous layer's output and the current layer's transformation, the signal passes through the scaling by $\alpha$ and the residual addition with the previous transformation's output. This means the input to $\mathcal{F}_l$ is not $\mathbf{x}_l$ in isolation—it is $\mathbf{x}_l$ as it enters the current sublayer, which may have accumulated variance from the previous sublayer's operations. The additional LayerNorm explicitly controls this variance, preventing the transformation function from receiving signals at uncontrolled scales. This is what the design evolution calls "decoupling scale and variance"—the normalization separates the problem of controlling input magnitude from the problem of controlling gradient magnitude.
The paper makes a subtle but important implementation note about the very first attention and FFN layers. As shown in Figure 2, "for the very first attention and FFN layers, we remove the final Post-LN and the scaling factor $\alpha$." This means these initial layers effectively degrade to standard Pre-LN blocks. The purpose is to ensure "stable signal initialization from the embedding layer"—the embedding outputs are not normalized, so applying Post-LN immediately would create an unstable initial forward pass. Starting with Pre-LN gives the signal a chance to stabilize before the Highway-style Post-LN mechanism takes over.
All Layer Normalization operations in Keel use RMS-based normalization (the standard in modern LLMs) with learnable affine weights $\boldsymbol{\gamma}$ but with the additive bias term set to zero ($\boldsymbol{\beta} = 0$). The normalization is:
where $\|\mathbf{x}\|_2$ is the L2 norm of the vector, and $\boldsymbol{\gamma} \in \mathbb{R}^d$ is the learnable element-wise scaling parameter. The paper defines the scalar magnitude of this parameter as $\gamma = \|\boldsymbol{\gamma}\|_\infty$ (the maximum absolute value among the elements of $\boldsymbol{\gamma}$), which is used in the gradient analysis to bound the Jacobian magnitude.
The Gradient Analysis: Why Post-LN Fails and How Keel Fixes It
The core theoretical contribution of the paper is a gradient-flow analysis that identifies the precise mechanism of failure in Post-LN and shows how Keel's modifications eliminate it. This analysis operates in two parts: first, deriving the gradient decay in standard Post-LN, and second, showing that Keel's $\alpha = L$ scaling prevents this decay.
Standard Post-LN gradient decay. The forward pass of the $l$-th sublayer in standard Post-LN is:
where the subscript on $\text{LN}_l$ indicates that each sublayer has its own LayerNorm with its own learnable affine parameters. The paper defines $\mathbf{z}_l = \mathbf{x}_l + \mathcal{F}_l(\mathbf{x}_l)$ as the pre-normalization state—the sum of the residual and transformed signals before normalization.
During backpropagation, the gradient of the loss $\mathcal{L}$ with respect to $\mathbf{x}_l$ is computed via the chain rule:
The Jacobian of the layer output with respect to its input expands as:
The $\mathbf{I}$ term comes from the residual connection (the derivative of $\mathbf{x}_l$ with respect to itself), and the $\partial\mathcal{F}_l / \partial\mathbf{x}_l$ term comes from the transformation branch.
The critical insight is that the Jacobian of the LayerNorm itself controls the gradient flow through both terms. The paper derives a bound on the spectral norm (or L2 operator norm) of this Jacobian:
where $\gamma_l$ is the magnitude (max absolute value) of the affine weight vector of the $l$-th layer's LayerNorm, and $\gamma_{l-1}$ is the magnitude of the previous layer's affine weights. The $\sqrt{2}$ factor arises from the geometry of summing the residual and transformed signals—the norm of the sum $\mathbf{z}_l = \mathbf{x}_l + \mathcal{F}_l(\mathbf{x}_l)$ is approximately $\sqrt{2}$ times the norm of either term individually if the two vectors are orthogonal (a reasonable approximation in high dimensions for independently parameterized branches).
What this bound means: each LayerNorm's backward Jacobian has magnitude approximately proportional to the ratio of its own affine weights to the previous layer's weights. At initialization, when all $\gamma$ values are approximately 1, each layer's Jacobian magnitude is approximately $1 / \sqrt{2}$—meaning each layer attenuates the gradient by roughly 30%.
Why this form (and why it causes vanishing): the $1/\sqrt{2}$ factor comes from the fact that LayerNorm divides by the norm of its input $\mathbf{z}_l$. When $\mathbf{x}_l$ and $\mathcal{F}_l(\mathbf{x}_l)$ are roughly equal in norm and approximately orthogonal, the norm of their sum is $\sqrt{2}$ times the norm of either component. The LayerNorm Jacobian includes a term that scales inversely with this norm, so it shrinks by $1/\sqrt{2}$. This attenuation compounds across layers.
The cumulative gradient magnitude across all $L$ layers (the product of per-layer Jacobian magnitudes) therefore scales as:
What this means operationally: if a gradient signal starts at the output layer with magnitude 1, by the time it reaches the input layer it will have been multiplied by $(1/\sqrt{2})^L = 2^{-L/2}$. For a model with $L = 100$ sublayers, this is approximately $2^{-50} \approx 8.9 \times 10^{-16}$—effectively zero in floating-point arithmetic. The bottom layers receive no training signal regardless of how large the loss is at the output. This is the root cause of Post-LN's inability to train at depth.
Keel's gradient-preserving modification. Keel changes the forward pass to include the $\alpha$ scaling and the additional input normalization, as discussed above. Reformulating with explicit notation for the two LayerNorms in each sublayer (the "inner" LN on the transformation input and the "outer" LN on the summed output):
where $\text{LN}_{l,1}$ is the output (outer) normalization and $\text{LN}_{l,2}$ is the input (inner) normalization for the transformation branch.
The revised gradient magnitude through the $l$-th sublayer's residual connection is derived as:
where $\gamma_{l,1}$ is the affine weight magnitude of the outer LN in layer $l$, $\gamma_{l-1,1}$ is the affine weight magnitude of the outer LN in the previous layer, and $\gamma_{l,2}$ is the affine weight magnitude of the inner LN in layer $l$.
What this means: the denominator now has two terms inside the square root. The first term $\gamma_{l-1,1}^2 \alpha^2$ represents the contribution from the scaled residual pathway—it grows with $\alpha$. The second term $\gamma_{l,2}^2$ represents the contribution from the transformation branch. When $\alpha$ is large (specifically $\alpha = L$), the first term dominates the denominator, and the fraction approximates $\gamma_{l,1} \cdot \alpha / (\gamma_{l-1,1} \cdot \alpha) = \gamma_{l,1} / \gamma_{l-1,1}$. At initialization, when all $\gamma$ values are approximately 1, this ratio is approximately 1—meaning each layer passes the gradient through without attenuation.
The paper sets $\alpha = L$ (the total number of sublayers) and analyzes the asymptotic behavior as $L \to \infty$. Assuming all $\gamma$ values are initialized to 1 (which they are, since the affine weights start at 1), the cumulative gradient product across all $L$ layers becomes:
The algebraic manipulation uses $\alpha = L$ and the approximation $\gamma_{l,2}^2 \approx 1$ and $\gamma_{l-1,1}^2 \approx 1$ to simplify the denominator to $\sqrt{L^2 + 1}$. The fraction $L / \sqrt{L^2 + 1}$ approaches 1 from below as $L$ grows. The limit $\lim_{L \to \infty} (L/\sqrt{L^2 + 1})^L = 1$ is a well-known limit (it is the limit of $(1 + 1/L^2)^{-L/2}$, which converges to $\exp(0) = 1$).
Why this form works: the key is that $\alpha = L$ grows with depth. In standard Post-LN, the residual connection has weight 1 regardless of depth, so the denominator's norm is dominated by the contribution from the additive branch $\mathcal{F}_l(\mathbf{x}_l)$, producing the $1/\sqrt{2}$ attenuation. In Keel, the residual connection's weight grows linearly with depth, so the residual signal dominates the sum increasingly at deeper layers. The outer LayerNorm divides by a norm that is dominated by $\alpha$ times the norm of $\mathbf{x}_l$, which cancels the $\alpha$ in the numerator, leaving the gradient magnitude near 1.
This analysis assumes the $\gamma$ values remain near their initialization values. In practice, the $\gamma$ values are learned and can drift during training. However, as long as they do not diverge catastrophically (which the training stability results suggest they do not), the gradient preservation property holds to first order.
The paper also sets $\alpha$ to a fixed value equal to $L$ rather than making it learnable. The justification is that "setting $\alpha = L$ is critical for maintaining training stability in very large-scale or deep models." For smaller architectures where vanishing or exploding gradients are less pronounced, the paper notes that $\alpha$ can be treated as a tunable hyperparameter with $\alpha > 1$ to potentially accelerate convergence. But for the extreme depths studied in this paper (up to 1024 layers, meaning $L = 2048$ sublayers), the fixed $\alpha = L$ is essential.
Why Keel Is a Post-LN Architecture
The paper goes to considerable length to argue that Keel should be classified as Post-LN despite the presence of an additional LayerNorm inside the transformation branch. This classification matters because it determines what properties the architecture inherits from the Post-LN family—specifically, strong inter-layer coupling and large gradients in deeper layers.
The paper's argument rests on a structural criterion: the distinction between Pre-LN and Post-LN depends on the shortcut branch, not the input to the transformation function. The paper demonstrates this by reformulating the standard Post-LN forward pass to make the shortcut explicit:
In standard Post-LN, the input $\mathbf{x}_l$ is the output of the previous layer's LayerNorm. Therefore, the forward pass can be written as:
(where the outer LN is omitted for clarity of the structural argument). This shows that in Post-LN, the shortcut branch carries $\text{LN}_{l-1}(\mathbf{x}_l)$—a normalized signal.
In Pre-LN, the forward pass is:
Here, the shortcut branch carries $\mathbf{x}_l$ directly—the unnormalized signal from the previous layer.
For Keel, the equivalent reformulation shows that the shortcut carries $\alpha \text{LN}_{l-1}(\mathbf{x}_l)$—a scaled but normalized signal:
Since the shortcut carries a normalized signal (like Post-LN) rather than an unnormalized signal (like Pre-LN), Keel is structurally Post-LN. The additional inner LayerNorm $\text{LN}_l$ normalizes the input to the transformation, but this is an internal detail of the transformation branch—it does not change the structural role of the shortcut.
Non-redundancy of successive LayerNorms. A natural question arises: if Keel applies $\text{LN}_l(\text{LN}_{l-1}(\mathbf{x}_l))$—two successive LayerNorms—can these be merged into a single operation? The paper argues they cannot, and the reasoning reveals an important design detail.
A LayerNorm operation is:
The first normalization $\text{LN}_{l-1}$ divides by the norm of its input and multiplies by its own affine weights $\boldsymbol{\gamma}_{l-1}$. The second normalization $\text{LN}_l$ takes this already-scaled output and re-normalizes it—dividing by its norm (which now reflects the $\boldsymbol{\gamma}_{l-1}$ scaling) and multiplying by its own affine weights $\boldsymbol{\gamma}_l$.
If the two $\boldsymbol{\gamma}$ vectors were identical (or if one were the inverse of the other), the operations could theoretically collapse. But they are independently learned, meaning the product $\boldsymbol{\gamma}_{l-1} \odot \boldsymbol{\gamma}_l$ (element-wise) is not representable as a single $\boldsymbol{\gamma}$ vector applied after a single normalization, because the norm computation in the second LayerNorm depends on the already-scaled values from the first LayerNorm. The two-step normalization allows the network to learn a more flexible transformation than a single normalization step would permit.
The paper notes that the motivation for this two-step design is discussed in Section 4.2—it emerged from the design evolution, where a learnable scaling parameter $\boldsymbol{\beta}$ was introduced before the transformation input (Attempt 2), and then explicit normalization was added before $\boldsymbol{\beta}$ to decouple gradient magnitude from input variance (Attempt 3). The final formulation absorbed the now-redundant $\boldsymbol{\beta}$ into the inner LayerNorm's $\boldsymbol{\gamma}$, leaving two successive LayerNorms as the minimal representation of the desired computation.
Keel vs. DeepNorm: Gradient Flow vs. Output Magnitude
The paper positions Keel in explicit contrast to DeepNorm (Wang et al., 2022), which was the most prominent prior attempt to stabilize Post-LN for deep architectures. The comparison reveals a fundamental difference in how the two approaches understand the instability problem.
DeepNorm's theoretical premise is that instability arises from the forward output magnitude growing without bound as depth increases. Each sublayer adds a transformation output to the residual stream, and without normalization, the signal norm would grow roughly linearly with depth. DeepNorm addresses this by choosing $\alpha$ to maintain constant variance in the forward pass:
with $\alpha = L^{0.25}$ for decoder-only architectures. The weight initialization is also scaled down by $\beta = L^{-0.25}$ to complement the residual scaling. The idea is that the smaller initialization reduces the magnitude of $\mathcal{F}_l(\mathbf{x}_l)$, and the larger $\alpha$ compensates on the residual path, so the sum has controlled variance.
Why this framing is insufficient, according to Keel's authors: bounding the forward pass does not guarantee a healthy backward gradient flow. The gradient through a Post-LN layer passes through the LayerNorm Jacobian, which depends on the norm of the summed signal $\mathbf{z}_l = \alpha \mathbf{x}_l + \mathcal{F}_l(\mathbf{x}_l)$. Even if this norm is controlled (not exploding), the Jacobian can still cause systematic attenuation if the ratio of the transformation branch's contribution to the residual branch's contribution is not right. Specifically, with DeepNorm's $\alpha = L^{0.25}$, the residual branch grows slowly with depth, but the transformation branch (even with small initialization) still contributes non-trivially to the denominator of the Jacobian, causing cumulative attenuation—though slower than standard Post-LN's $\mathcal{O}(2^{-L/2})$.
Keel's $\alpha = L$ grows much faster ($L$ vs. $L^{0.25}$), meaning the residual signal dominates the sum more aggressively, which is what the gradient analysis shows is necessary to achieve a cumulative product of approximately 1.
The initialization dependency problem. DeepNorm relies on a specialized weight initialization ($L^{-0.25}$ down-scaling) to keep the transformation outputs small. The paper argues that in the context of LLM-scale pretraining on trillions of tokens, "the model weights drift significantly from their initialization." The benefits of a specific initialization strategy diminish as training progresses because the weights move away from their initial values, and the carefully balanced forward-pass variance can degrade. As the training dynamics evolve, the DeepNorm scheme may lose its stabilizing effect, "leading to instability in later stages of training."
Keel eliminates this dependency by building the stabilization into the architecture itself—the $\alpha = L$ scaling and the additional input normalization are structural features of the forward and backward passes, not properties of the initial weight values. This "structural approach ensures stability throughout the entire training trajectory, rather than just the initial phase." The empirical evidence supports this: the paper reports that DeepNorm "often underperform[s] standard Pre-LN baselines" when applied to large-scale decoder-only LLMs, suggesting that its initialization-dependent stabilization does not hold up over the course of full pretraining.
The paper also notes that DeepNorm was primarily demonstrated on encoder-decoder architectures, not the decoder-only architectures that dominate modern LLM training. The different architectural topology (cross-attention, different depth ratios between encoder and decoder) may affect how the forward-pass variance and backward gradient dynamics interact with DeepNorm's scaling scheme, making the transfer to decoder-only LLMs less straightforward than the original paper might suggest.
Design Evolution: From Naive Scaling to Keel
Section 4.2 presents a step-by-step ablation of how the Keel formulation was developed, starting from a naive application of residual scaling and incrementally adding components based on observed failure modes. This narrative is valuable because it makes the design choices concrete—each modification is a direct response to a specific problem observed in the previous attempt.
Goal of the evolution: to find a formulation that both (a) prevents gradient vanishing through the residual pathway and (b) does not introduce new failure modes (like gradient attenuation through the transformation branch or training instability). The endpoint is the minimal set of modifications that achieve these goals simultaneously.
Attempt 1: Naive Residual Scaling. The initial attempt was a straightforward adaptation of DeepNorm's idea to scale the residual connection, but without DeepNorm's specialized initialization:
with $\alpha = L$. This is essentially DeepNorm with a different scaling exponent ($L$ instead of $L^{0.25}$) and standard initialization.
Observation: this formulation exhibited better stability than vanilla DeepNorm on LLMs, but "both its training stability and final convergence performance lagged behind the standard Pre-LN baseline." The hypothesized cause was that the unnormalized input to $\mathcal{F}_l$ resulted in high variance within the residual branch. Since $\mathbf{x}_l$ comes from the previous layer's LayerNorm, it is already normalized, but the addition of the scaled residual from the previous step means the signal entering $\mathcal{F}_l$ has an effective scale that depends on the accumulation of previous transformations—the variance is not explicitly controlled.
What this reveals: scaling the residual pathway alone is not sufficient. While it may help the gradient through the shortcut, the transformation branch's input distribution is uncontrolled, which complicates the optimization landscape for the attention and FFN parameters.
Attempt 2: Learnable Input Scaling. To control the variance entering the transformation, the next attempt introduced a learnable vector $\boldsymbol{\beta}$ (initialized to a vector of ones) that scales the input before it enters $\mathcal{F}_l$:
The idea is that $\boldsymbol{\beta}$ can learn to scale down the input to $\mathcal{F}_l$, reducing the variance of the transformed output and stabilizing the forward pass.
Observation: this improved forward pass stability—as $\boldsymbol{\beta}$ learned to take on smaller values, the variance entering $\mathcal{F}_l$ was reduced, and the overall training became more stable. However, it introduced a new problem: "gradient attenuation." The gradient flowing backward through $\mathcal{F}_l$ is multiplied by $\boldsymbol{\beta}$ (by the chain rule, $\partial(\boldsymbol{\beta} \odot \mathbf{x}_l) / \partial\mathbf{x}_l = \text{diag}(\boldsymbol{\beta})$). When $\boldsymbol{\beta}$ shrinks to control forward variance, it proportionally scales down the gradients flowing back through the transformation—"effectively choking the learning signal for the attention and FFN blocks." The overall stability remained inferior to Pre-LN.
What this reveals: there is a fundamental tension between controlling forward variance (small $\boldsymbol{\beta}$) and maintaining backward gradient magnitude (large $\boldsymbol{\beta}$). A simple multiplicative gate cannot resolve this tension because forward and backward effects of the gate are coupled—whatever the gate does to the forward signal, it does inversely to the backward gradient.
Attempt 3: Decoupling Scale and Variance. To resolve this tension, the next attempt separated the problem of input variance control from the problem of gradient propagation by introducing an explicit LayerNorm before the learnable scaling:
The inner LayerNorm first normalizes $\mathbf{x}_l$ to unit norm, which controls the input variance regardless of what has accumulated in the residual stream. Then $\boldsymbol{\beta}$ scales this normalized signal. The gradient analysis for this configuration (the paper's Equation 27, which is presented in the text but not numbered sequentially in the main flow) shows:
The crucial observation is that the $\gamma_{l,2}$ term (from the inner LayerNorm's affine weights) appears in both the numerator and denominator—in the numerator from the inner LN's backward Jacobian and in the denominator indirectly through the outer LN's Jacobian. This means "the scaling factor $\gamma_{l,2}$ (associated with the inner LN) is effectively canceled by the affine weights $\gamma_{l-1,1}$ of the previous layer, preventing the gradient vanishing observed in Attempt 2."
Why this works: the normalization before $\boldsymbol{\beta}$ breaks the coupling between forward variance control and backward gradient attenuation. The LayerNorm sets the input to $\mathcal{F}_l$ to unit norm regardless of what $\boldsymbol{\beta}$ does, so $\boldsymbol{\beta}$ can be small to control variance without also shrinking the gradient (since the gradient no longer passes through $\boldsymbol{\beta}$ alone—it passes through the LayerNorm Jacobian first, which has its own scaling properties that compensate).
Final Formulation (Keel). The final step recognizes that $\boldsymbol{\beta}$ is now redundant. The inner LayerNorm already has learnable affine weights $\boldsymbol{\gamma}$ that can perform the same element-wise scaling as $\boldsymbol{\beta}$. By merging the external $\boldsymbol{\beta}$ into the internal $\boldsymbol{\gamma}$ of the LayerNorm, the formulation simplifies to:
This is the minimal representation: two modifications to standard Post-LN—the $\alpha = L$ scaling on the residual, and the explicit LayerNorm on the input to the transformation. Both are structural, neither depends on initialization, and together they prevent gradient vanishing while maintaining forward stability.
Depth-wise Test-Time Training Interpretation
Section 4.3 offers a conceptual framing that connects Keel's depth-wise stabilization to recent work in sequence modeling, particularly linear attention and Test-Time Training (TTT) approaches. This framing is speculative (it does not drive any of the paper's empirical results) but provides an interesting lens for understanding why the specific modifications in Keel are effective.
The connection works in two steps. First, residual network layers can be interpreted as gradient descent steps. Consider the standard Pre-LN residual update:
where $\mathcal{G}(\mathbf{x}_l; \mathbf{W})$ is the transformation block (attention or FFN). This can be viewed as a gradient step on a "depth-wise objective":
where $\langle\cdot, \cdot\rangle$ denotes an inner product. The update rule $\mathbf{x}_{l+1} = \mathbf{x}_l + \mathcal{G}(\mathbf{x}_l; \mathbf{W})\mathbf{W}_o^\top$ is the gradient of $\mathcal{L}_{\text{depth}}^{(l)}$ with respect to $\mathbf{x}_l \mathbf{W}_o$, evaluated at the current activation.
Second, this depth-wise optimization is structurally analogous to sequence-wise recurrence in linear attention models. In linear attention, the hidden state updates as:
which can be interpreted as a gradient step on a sequence-wise objective $\mathcal{L}_{\text{seq}}^{(t)} = -\langle \mathbf{S}_t \mathbf{k}_t, \mathbf{v}_t \rangle$—maximizing alignment between the retrieved state and the current value.
The analogy: depth-wise propagation in residual networks is to sequence-wise recurrence in linear attention models as TTT along the depth dimension is to TTT along the sequence dimension. Both are optimization processes that refine a hidden state through incremental updates, and both face stability challenges when the number of steps grows large.
From this perspective, Keel's $\alpha$ scaling is analogous to gating mechanisms in sequence models (like Titans or LaCT) that control how much of the previous state is carried forward versus how much is updated. The additional LayerNorm is analogous to state normalization techniques that prevent the hidden state's norm from growing or shrinking uncontrollably over many update steps. Both serve the same role in their respective domains: ensuring that the optimization process (whether along sequence length or network depth) remains stable over many steps.
The paper notes that "deep modeling exhibits phenomena similar to those in sequence modeling." Specifically: (1) the first few layers are especially critical (analogous to attention sinks in sequence modeling, where initial tokens absorb disproportionate attention), and (2) deeper layers become increasingly significant (analogous to recency bias in sequence modeling, where recent tokens receive more attention). The layer redundancy analysis (Section 8) supports the first claim: removing the first layer of a deep model destroys performance (PPL spikes above $10^4$), while removing shallow intermediate layers causes minimal degradation.
This conceptual framing is not empirically validated in the paper—it is presented as a direction for future research: "techniques developed to improve long-context sequence recurrence can likely be adapted to develop infinite-depth model propagation, and vice versa." The value is in providing an intuitive bridge between the architectural choices in Keel and the broader landscape of sequence modeling research.
4. Key Insights and Innovations
Innovation 1: The Root Cause of Post-LN Instability Is Backward Gradient Vanishing, Not Forward Variance Explosion
The paper's most fundamental conceptual contribution is a diagnostic reframing of why Post-LN fails at depth. Prior work—most prominently DeepNorm (Wang et al., 2022)—operated under the assumption that the instability of deep Post-LN networks arises from forward-pass variance explosion: each sublayer adds a transformed output to the residual stream, and without proper scaling, activation magnitudes grow without bound as depth increases, causing optimization to diverge. This forward-variance framing motivated a family of solutions that bound output magnitudes through initialization-dependent scaling (DeepNorm, Admin) or that avoid the problem entirely by moving normalization before the residual addition (Pre-LN).
Keel's authors perform a gradient-flow analysis (Section 3.2) that identifies a different failure mechanism entirely: exponential backward gradient vanishing through the LayerNorm Jacobian. Even if forward activations are perfectly controlled (which they are, since LayerNorm already constrains output magnitudes to unit norm), the backward signal decays as because each LayerNorm's Jacobian attenuates the gradient by roughly per sublayer. This is not a forward variance problem—it is a backward signal propagation problem.
The significance of this reframing extends beyond the specific architectural fix that Keel proposes. It explains a puzzle that the forward-variance view could not: why DeepNorm, which successfully controls forward output magnitude in encoder-decoder architectures, fails to stabilize large-scale decoder-only LLMs. DeepNorm's scaling factors are derived to bound forward variance, but bounding the forward pass does not guarantee a healthy backward gradient flow. The gradient still passes through the LayerNorm Jacobian, and if the balance between the residual and transformation branches in the pre-normalization sum is not right, attenuation still compounds across layers—just more slowly than in vanilla Post-LN. The paper reports that DeepNorm "often underperforms standard Pre-LN baselines" when applied to decoder-only LLMs (Section 4.1), which is exactly what a backward-gradient-failure account would predict: DeepNorm partially addresses the wrong problem and therefore does not fully resolve the true failure mode.
This shift in diagnostic framework also explains why Pre-LN works despite being less expressive. Pre-LN's gradient path through the residual connection is , with no LayerNorm Jacobian to cause attenuation. The backward signal propagates cleanly, enabling stable training at scale. But this comes at a cost: because gradients flow predominantly through the identity connection, deeper layers contribute diminishingly to parameter updates—the effective depth of a Pre-LN model is smaller than its architectural depth, as confirmed by the layer redundancy analysis in Section 8 showing that shallow layers in Pre-LN models can be removed with minimal perplexity increase. The field had accepted this depth-utilization penalty as the price of stability. Keel's gradient analysis shows that this tradeoff is not inherent—it is a consequence of how the ResNet-style residual connection interacts with LayerNorm's Jacobian, and a different residual pathway can eliminate the tradeoff entirely.
Evidence anchor: the cumulative gradient magnitude product for standard Post-LN (Equation 15, Section 3.2) decays as , while Keel's modified pathway (Equation 19, Section 3.3) converges to 1 as . The stability experiments (Table 1, Section 5.1) confirm the practical consequence: vanilla Post-LN diverges at learning rates of at 64 layers, while Keel tolerates —nearly two orders of magnitude higher.
Innovation 2: Structural Stabilization Trumps Initialization-Dependent Stabilization for Long-Horizon LLM Training
A second conceptual contribution is the argument that architectural stabilization is fundamentally more robust than initialization-dependent stabilization for large-scale LLM training, and that this distinction matters more than the specific scaling constants chosen. This is not a claim about which method achieves better final accuracy (though Keel does outperform DeepNorm empirically)—it is a claim about how stabilization should be achieved given the realities of trillion-token pretraining.
DeepNorm (Wang et al., 2022) and Admin (Liu et al., 2020) treat Post-LN instability as solvable through specialized initialization: choose the right initial weight scales and residual multipliers, and the network starts in a regime where gradients flow properly. The implicit assumption is that if the network starts in a well-conditioned region of parameter space, training will remain stable. But LLM pretraining at trillion-token scale fundamentally challenges this assumption: the model weights drift substantially from their initialization over the course of training. The carefully constructed initial balance between residual and transformation branches degrades as the weights evolve, and the stabilization that depended on that balance can erode—"leading to instability in later stages of training" (Section 4.1).
Keel takes the opposite approach: bake the stabilization into the architecture itself, so it holds regardless of where the weights are. The scaling and the additional input normalization are structural features of the forward and backward passes that operate on the activations, not properties of the initial weight values. The gradient analysis shows that the cumulative gradient product converges to 1 at any point in training as long as the (LayerNorm affine weight) values have not diverged catastrophically—it does not depend on weights staying near their initialization.
This distinction has implications beyond the specific methods compared in the paper. It suggests a design principle for deep architectures more broadly: any stabilization mechanism that depends on the network being near its initial state is fragile under long training horizons. The more an architecture relies on initialization for stability, the more likely it is to encounter training instability mid-run as weights drift—which is precisely the pattern of "irrecoverable instability" that the paper identifies as one of three divergence modes (Section 5.1, Figure 3b): loss spikes from which the model never recovers. These mid-training instabilities are particularly damaging because they waste compute that has already been invested and often require restarting from an earlier checkpoint with modified hyperparameters.
The HybridNorm and Mix-LN approaches (Section 2.4) occupy an intermediate position: they are structurally hybrid (interleaving Post-LN and Pre-LN blocks) rather than initialization-dependent, but they still use the standard ResNet-style residual connection in their Post-LN blocks. As a result, the gradient vanishing mechanism (Equation 15) remains active in those blocks, limiting how deep the Post-LN segments can be before attenuation becomes prohibitive. Keel's structural fix addresses the residual pathway itself, making the stabilization independent of both initialization and block placement.
Evidence anchor: Table 1 (Section 5.1) shows that DeepNorm achieves a maximum tolerable learning rate of only at both 64 and 512 layers—barely better than vanilla Post-LN's at 64 layers and at 512 layers. This suggests that DeepNorm's initialization-dependent stabilization provides minimal robustness beyond what standard Post-LN offers, especially as depth increases. Keel achieves maximum tolerable LRs an order of magnitude higher ( at 64 layers, at 512 layers), consistent with a structural mechanism that remains effective regardless of weight drift.
Innovation 3: Depth Scaling Can Be Restored as a Viable Axis for Improving Expressivity Per Parameter
The paper's third conceptual contribution is an empirical demonstration that depth scaling can deliver meaningful improvements in model quality when the optimization barrier is removed, challenging the widespread practical assumption that width scaling is inherently more effective than depth scaling for Transformers. This is not a theoretical claim about representational capacity—that deeper networks can represent richer functions has been known since the early days of deep learning. Rather, it is an empirical finding that the poor returns from depth scaling in current LLMs are an artifact of architectural instability, not a fundamental property of depth itself.
The field's default has been to scale width over depth. Models like GPT-3, PaLM, and LLaMA are wide relative to their depth—the width-to-depth ratio is typically large, and scaling efforts historically add parameters by increasing hidden dimension rather than adding layers. This preference is rational given the evidence: deeper Pre-LN models show diminishing returns, with additional layers contributing less and less to model quality because Pre-LN's gradient dynamics cause shallow layers to be undertrained (the "effective depth" problem documented in Section 2.4 and Appendix 8). Given a fixed parameter budget, the practical choice has been to go wider rather than deeper because wider models train more reliably and deliver better downstream performance.
Keel's experiments directly challenge this conventional wisdom. The "Deeper vs. Wider" comparison (Section 5.5, Table 5) holds the parameter count fixed at 3B and compares three configurations: a wide Pre-LN model (128 layers, 2048 hidden dim), a deep Pre-LN model (512 layers, 1024 hidden dim), and a deep Keel model (same 512-layer topology). The deep Pre-LN model matches the wide Pre-LN model on average downstream score (52.2 vs. 52.3) despite having a worse training loss, suggesting that depth provides some benefit even under Pre-LN but that optimization difficulties prevent it from pulling ahead. The deep Keel model, however, achieves an average score of 55.5—a +3.2 point improvement over the deep Pre-LN baseline and +3.3 points over the wide baseline. On reasoning-intensive tasks, the gap is larger: GSM-8K improves from 38.1 (deep Pre-LN) or 35.3 (wide Pre-LN) to 43.8 with Keel.
The significance of this result is not the absolute numbers (3B parameters is modest by current standards) but the reversal of the depth-vs-width ordering. When depth scaling is limited by optimization instability, wider models win by default because they train more reliably—not because width is intrinsically more expressive. When the optimization barrier is removed (by Keel's gradient-preserving residual pathway), deeper models can outperform wider ones at the same parameter count because they can learn more hierarchical representations. This validates the theoretical motivation for depth scaling and reframes the problem: the bottleneck was never depth itself, but the architectures we were using to achieve it.
The depth scaling results (Section 5.3, Table 3) reinforce this finding across a broader range of depths. From 64 to 1024 layers, Keel consistently outperforms Pre-LN, but the gain grows with depth: +1.7 points at 64 layers, +1.2 at 128 layers, +3.8 at 512 layers, +3.0 at 1024 layers. This is exactly the pattern one would expect if Pre-LN's effective-depth problem worsens with architectural depth—the deeper the Pre-LN model, the more layers are essentially redundant, so the gap between Pre-LN and a properly trained architecture widens. At 1024 layers, Keel achieves a GSM-8K score of 58.6 versus Pre-LN's 49.8, and HumanEval of 32.9 versus 29.9. Pre-LN shows signs of stagnation between 512 and 1024 layers on reasoning tasks, whereas Keel continues to improve—consistent with the interpretation that Pre-LN's effective depth saturates while Keel's continues to grow with architectural depth.
Evidence anchor: Table 3 (Section 5.3) and Table 5 (Section 5.5) provide the primary evidence. The depth scaling trend in Table 3 shows Keel's average advantage growing from +1.7 at 64L to +3.0 at 1024L. Table 5 shows the reversal of the depth-vs-width ordering at 3B parameters.
Innovation 4: The Training Loss vs. Downstream Performance Discrepancy Reveals a Hidden Cost of Poor Depth Utilization
A more subtle but potentially far-reaching finding emerges from the observation that training loss and downstream task performance are not always positively correlated in deep LLMs (Section 9). The paper documents two instances of this discrepancy, and together they suggest that standard pretraining metrics may systematically underestimate the benefits of better depth utilization.
The first instance compares the wide Pre-LN model (128 layers, 2048 hidden dim) with the deep Pre-LN model (512 layers, 1024 hidden dim) at the same 3B parameter budget. The wide model achieves a significantly lower training loss (Figure 8a), yet the deep model matches or slightly exceeds it on downstream tasks (Table 5: 52.2 vs. 52.3 average). The second instance compares Pre-LN and Keel at a low learning rate of . Pre-LN achieves a slightly lower training loss than Keel (Figure 8b), yet Keel demonstrates better downstream performance (Table 2: 50.5 vs. 49.5 average at this learning rate).
The standard assumption in LLM pretraining is that lower perplexity (or training loss) translates to better downstream performance—this is the basis for using validation loss as a proxy for model quality during training and for scaling laws that predict downstream capability from pretraining loss. The paper's finding challenges this assumption specifically in the context of depth scaling: a model with poorer depth utilization (wide or Pre-LN) can achieve lower training loss by overfitting to the statistical patterns in the pretraining data through its wide layers, while a model with better depth utilization (deep Keel) learns more structured, hierarchical representations that generalize better to downstream reasoning tasks even if its raw next-token prediction loss is slightly higher.
This is not presented as a fully explained phenomenon—the paper states "we aim to further investigate the mechanisms behind this phenomenon in future work" (Section 9). But even as an empirical observation, it has practical implications: if training loss systematically underestimates the downstream value of depth, then the field's historical preference for width scaling may have been reinforced by a misleading metric. Wide models look better during pretraining because they optimize the training objective more effectively, but deep models with proper gradient flow may actually learn more useful representations. This would explain why deep Pre-LN models have historically disappointed—not only are they harder to train, but the standard metric for evaluating training progress (loss curves) may obscure whatever benefits their depth does provide.
The paper's recommended response is practical: "evaluating end-task performance during pre-training is essential" rather than relying solely on training loss as a proxy for model quality. This is a methodological contribution that could affect how future depth-scaling experiments are designed and evaluated.
Evidence anchor: Figure 8 and Section 9 provide the loss curve comparisons. Table 5 (deep vs. wide Pre-LN) and Table 2 (low-LR Pre-LN vs. Keel) provide the downstream performance data that contradicts the loss ordering.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. The paper uses two distinct training data sources across its experiments. The primary experiments (Sections 5.1–5.3, 5.5–5.7) use an internal private dataset of up to 1T tokens for pre-training, plus continued pre-training (CPT) on an additional 250B tokens to enhance reasoning and coding capabilities. Data scaling experiments (Section 5.4) use the FineWeb-EDU dataset (Penedo et al., 2024) at 10B and 40B token budgets. Layer redundancy analysis (Section 8) uses the C4 validation dataset for perplexity evaluation.
-
Base model(s). All experiments use decoder-only Transformer models with RMS-based Layer Normalization, RoPE positional embeddings, and standard multi-head attention. Hidden dimensions are 1024 or 2048 (Table 8, Section 10), with model sizes ranging from 3B parameters. The maximum depth tested is 1024 layers (meaning 2048 total sublayers counting both attention and FFN). The authors deliberately choose a topology with an "extreme depth-to-width ratio of 0.5 (512 layers vs. 1024 width)" to stress-test optimization stability where gradient preservation is hardest (Section 5.6). All models are trained from scratch—there is no fine-tuning from existing pretrained checkpoints.
-
Metrics. Three categories of metrics are used across experiments. Training stability is measured via Maximum Tolerable Learning Rate (Max LR)—"the highest learning rate a model can sustain during warm-up stage without diverging" under an aggressive peak of (Section 5.1). Downstream performance is measured via zero-shot and few-shot accuracy on a suite of 14 benchmarks spanning general knowledge (MMLU 5-shot, ARC-Easy 25-shot, ARC-Challenge 25-shot, HellaSwag 0-shot, LAMBADA 0-shot, PIQA 0-shot, AGI-Eval 0-shot, Winogrande 0-shot, CommonsenseQA 0-shot), reasoning (GSM-8K 5-shot), coding (HumanEval 0-shot, MBPP 0-shot), and multilingual understanding (CMMLU 5-shot, C-Eval 5-shot), with evaluation conducted via the lm-evaluation-harness (Gao et al., 2023). Training quality is monitored via loss curves and, in the layer redundancy analysis (Section 8), via perplexity increase on C4 when individual layers are removed.
-
Baselines. The paper compares Keel against five architectural and normalization baselines, all implemented within the same decoder-only Transformer framework to ensure fair comparison. Pre-LN (Xiong et al., 2020) is the modern standard, placing LayerNorm before the residual addition; it serves as the primary competitive baseline throughout. Vanilla Post-LN (Vaswani et al., 2017) places LayerNorm after the residual addition with no modifications. DeepNorm (Wang et al., 2022) scales the residual branch by and down-scales weight initialization by . HybridNorm (Cao et al., 2023) interleaves Post-LN and Pre-LN blocks throughout the network. Mix-LN (Li et al., 2024) applies Post-LN in lower layers and transitions to Pre-LN in upper layers. All baselines use the same RMS-based LayerNorm with learnable affine weights and zero bias, consistent with Keel's configuration.
-
Generation budget / compute accounting. The paper measures compute indirectly through training hyperparameters rather than FLOPs. The primary axes for fair comparison are: (a) matching parameter count (all comparisons at the same depth and width, or at the same total parameter budget in the deeper-vs-wider experiment), (b) matching training data volume (same number of tokens for each compared pair), and (c) optimal learning rate tuning per method (each architecture uses its own best learning rate, determined either by the maximum tolerable LR test in Section 5.1 or by explicit sweep in Section 5.2). Inference-time compute is not studied—this is purely a training-time architectural comparison. The paper argues that the ability to use higher learning rates is itself an advantage of Keel, not a confound, because it reflects a more well-conditioned optimization landscape.
-
Cross-validation / statistical protocol. The paper does not report confidence intervals, standard errors, or statistical significance tests for any of its downstream benchmark results. The depth scaling results (Table 3, Section 5.3) represent single training runs per configuration (no repeated seeds). The learning rate sweep (Table 2, Section 5.2) uses three learning rates per model, also single runs each. The 1T-token pretraining results (Table 6, Section 5.7) include only a single run per architecture. For SFT (Table 7, Section 5.7), the authors perform grid search over learning rates and training epochs and report the best result per architecture—this is a best-of-N selection rather than a statistical protocol that accounts for variance. The Max LR experiments (Section 5.1) use a deterministic protocol: increase learning rate linearly from 0 to divergence, record the LR at the step immediately preceding divergence. This produces a single number per architecture-depth configuration with no variance estimate. Overall, the experimental design prioritizes demonstrating architecture-level effects over quantifying statistical reliability—a reasonable choice for an architectural innovation paper, but one that leaves open the question of how much the reported improvements would vary across random seeds.
Main Quantitative Results
Training Stability: Maximum Tolerable Learning Rate (Section 5.1)
The central stability result appears in Table 1. At 64 layers with a 5000-step warmup and , Keel tolerates a maximum learning rate of , compared to for Pre-LN—a 32% improvement. Vanilla Post-LN diverges at , nearly two orders of magnitude lower than Keel. DeepNorm (), HybridNorm (), and Mix-LN () all remain within a factor of 3 of vanilla Post-LN, indicating that their stabilization mechanisms provide only marginal robustness beyond the standard formulation.
At 512 layers, Keel maintains a Max LR of , versus for Pre-LN—a 35% advantage. The Post-LN variants remain clustered below : DeepNorm and HybridNorm both at , Mix-LN also at , vanilla Post-LN at . The gap between Keel and these methods widens with depth: at 64 layers Keel is roughly 12× more tolerant than the best Post-LN variant (Mix-LN); at 512 layers it is roughly 18× more tolerant. This widening gap is consistent with the theoretical claim that Post-LN's gradient vanishing compounds with depth while Keel's gradient preservation holds independent of depth.
The paper identifies three distinct failure modes for divergence (Figure 3): Loss Stagnation (plateauing early), Irrecoverable Instability (spike with permanent degradation), and Optimization Degradation (anomalously slow convergence without explicit spikes). The Max LR metric captures all three—any run exhibiting any of these behaviors is marked as diverged. This is a more comprehensive failure detection protocol than simply checking for NaN, which would miss the "silent" failure modes (stagnation and degradation) that are common in large-scale LLM training.
Figure 4 provides qualitative evidence of the instability that limits Pre-LN's usable learning rate. At 256 layers with a batch size of 8M tokens and 512 layers with a batch size of 4M tokens, Pre-LN exhibits a "pronounced loss spike" during early training when trained with a higher learning rate—the kind of irrecoverable instability that forces a lower learning rate cap. Keel avoids these spikes at the same learning rates, enabling the higher peak learning rates used in later experiments.
Optimal Learning Rate and Its Effect on Downstream Performance (Section 5.2)
Having established that Keel tolerates higher learning rates, the paper investigates whether this tolerance translates to better downstream performance. Table 2 reports the results of sweeping peak learning rates (1.5, 3.0, and ) for both Pre-LN and Keel at 512 layers with hidden dimension 1024, trained on 250B tokens.
Headline finding: Keel benefits monotonically from higher learning rates, while Pre-LN shows inconsistent and sometimes negative responses. At the highest learning rate common to both (, which is within Keel's Max LR of but exceeds Pre-LN's Max LR of from Table 1 at 512 layers—note that this is trained stably at for Pre-LN since the Max LR test uses a more aggressive peak), Keel achieves an average score of 55.5 across 14 benchmarks, compared to Pre-LN's best configuration (which is at for most tasks despite some degradation) averaging 52.3.
The pattern within Pre-LN reveals instability even at learning rates it can nominally tolerate. Moving from 3.0 to , some tasks improve (MMLU: 48.1 → 52.9; GSM-8K: 31.1 → 38.1; AGI-Eval: 29.7 → 34.7), while others degrade (ARC-Easy: 75.3 → 74.0; ARC-Challenge: 44.4 → 43.6; MBPP: 24.0 → 22.8). This task-dependent instability suggests that even when Pre-LN does not diverge catastrophically, its optimization is poorly conditioned—gradients for different task-relevant subspaces respond differently to the same learning rate, creating a tension where no single learning rate is optimal for all capabilities.
In contrast, Keel shows monotonic or near-monotonic improvement across essentially all benchmarks as learning rate increases from 1.5 to 3.0 to . On GSM-8K, Keel jumps from 30.5 to 36.6 to 43.8. On AGI-Eval, from 30.5 to 34.6 to 39.6. The only exceptions are tiny: ARC-Challenge dips from 45.1 to 44.5 at the middle learning rate before recovering to 48.9 at the highest. This consistency supports the claim that Keel's gradient preservation creates a more uniformly well-conditioned optimization landscape—learning rates that benefit one set of parameters also benefit others, without destructive interference.
The gains from higher learning rates are most pronounced on reasoning-intensive benchmarks. At : GSM-8K shows Keel 43.8 vs. Pre-LN 38.1 (+5.7 points); AGI-Eval shows 39.6 vs. 34.7 (+4.9); CommonsenseQA shows 61.3 vs. 55.7 (+5.6). On knowledge-intensive tasks, the gaps are smaller: MMLU 56.3 vs. 52.9 (+3.4); HellaSwag 67.4 vs. 64.9 (+2.5). This differential benefit—larger on reasoning, smaller on factual recall—is a recurring pattern throughout the paper and suggests that deep architectures stabilized by Keel are particularly valuable for tasks requiring multi-step inference, where hierarchical representations across many layers matter most.
Depth Scaling: Performance at 64–1024 Layers (Section 5.3)
Table 3 is the paper's most comprehensive scaling result, comparing Pre-LN and Keel at 64, 128, 512, and 1024 layers. All models are trained on 190B tokens of general pre-training followed by 60B tokens of continued pre-training (250B total). Learning rates are tuned per configuration: for 64L and 128L models, for 512L models, and for Keel at 1024L versus for Pre-LN at 1024L (because Pre-LN exhibits "substantial instability" at per Figure 4b).
Headline finding: Keel's advantage grows with depth, consistent with the claim that Pre-LN's effective depth problem worsens as architectural depth increases. At 64 layers, Keel leads Pre-LN by +1.7 points on average (39.6 vs. 37.9). At 128 layers, the gap is +1.2 points (46.5 vs. 45.3)—a slight narrowing. At 512 layers, the gap widens to +3.8 points (58.1 vs. 54.3). At 1024 layers, the gap is +3.0 points (60.9 vs. 57.9). The non-monotonicity (narrowing at 128L before widening at 512L) may reflect that at moderate depths, Pre-LN's effective depth is still sufficient to capture most of the training signal, while at extreme depths, the underutilization of shallow layers becomes a binding constraint.
Reasoning benchmarks show the strongest depth-scaling gains for Keel. On GSM-8K: at 64L, Keel leads 12.9 vs. 9.6 (+3.3); at 128L, 28.0 vs. 22.4 (+5.6); at 512L, 49.8 vs. 45.6 (+4.2); at 1024L, 58.6 vs. 49.8 (+8.8). The 1024L gap is particularly striking: Pre-LN's GSM-8K score is identical at 512L and 1024L (49.8 at both depths), suggesting that Pre-LN has saturated—adding 512 more layers provides zero improvement on this reasoning task. Keel, by contrast, jumps from 49.8 at 512L to 58.6 at 1024L, a +8.8 point gain. This is the strongest single piece of evidence that Pre-LN's depth scaling is bottlenecked by gradient dynamics, not by the intrinsic difficulty of leveraging additional layers.
HumanEval shows a similar pattern: Pre-LN improves from 22.6 at 512L to 29.9 at 1024L (+7.3), but Keel improves from 32.3 to 32.9 (+0.6). The absolute gap favors Keel at both depths (32.3 vs. 22.6 at 512L; 32.9 vs. 29.9 at 1024L), but Pre-LN partially catches up at 1024L on this task—unlike GSM-8K, where it completely stagnates. MBPP shows Keel ahead at all depths, with gaps ranging from 0.2 to 9.7 points depending on depth.
Knowledge benchmarks show smaller but consistent Keel advantages. MMLU: gaps range from +1.7 (64L) to +3.1 (512L). HellaSwag: +1.3 (64L) to +2.2 (512L). The multilingual benchmarks (CMMLU, C-Eval) show larger Keel advantages—often +4 points—suggesting that deep hierarchical representations are particularly important for cross-lingual transfer.
A cross-depth comparison within Pre-LN reveals an important pattern. Pre-LN's average score improves from 37.9 (64L) to 45.3 (128L) to 54.3 (512L) to 57.9 (1024L)—gains of +7.4, +9.0, and +3.6 respectively. The diminishing returns are clear: the jump from 512L to 1024L provides less than half the improvement of the jump from 128L to 512L, despite adding the same number of layers. Keel's improvements are: 39.6 → 46.5 → 58.1 → 60.9, with jumps of +6.9, +11.6, and +2.8—also diminishing, but with a larger absolute gain at the 128L→512L transition (+11.6 vs. +9.0) and a larger final score at 1024L (60.9 vs. 57.9).
Data Scaling: Performance at 10B and 40B Tokens (Section 5.4)
Section 5.4 and Table 4 test whether Keel's advantages persist or change as training data volume scales, using the FineWeb-EDU dataset. Both Pre-LN and Keel are configured at 256 layers with hidden dimension 1024 (approximately 3B parameters). Learning rates are tuned independently for each architecture on the 10B-token run, and the best rates (Pre-LN: ; Keel: ) are carried forward to the 40B-token run.
Headline finding: Keel's advantage widens with more data, not less. At 10B tokens, Keel averages 61.5 vs. Pre-LN's 60.3 (+1.2 points). At 40B tokens, Keel averages 68.2 vs. Pre-LN's 66.7 (+1.5 points). The gap on HellaSwag widens from +0.9 to +2.6. Winogrande widens from +3.4 to +2.8 (slight narrowing). SciQ widens from +2.0 to +0.7 (narrowing). The pattern across tasks is not uniform, but the overall trend favors Keel at larger data scales.
Figure 5 provides the training loss curves. At 10B tokens (Figure 5a), Keel exhibits higher loss early in training but "overtakes Pre-LN as training proceeds and achieves lower loss by the end." At 40B tokens (Figure 5b), the same pattern holds but "the performance gap further widens in favor of Keel." The crossover point—where Keel's loss drops below Pre-LN's—appears to occur later in the 40B run than the 10B run, suggesting that Keel requires sufficient training to overcome its initially higher loss but then pulls ahead by a larger margin.
The paper attributes this behavior to "the lower effective depth of Pre-LN relative to Keel, which limits its representational capacity." The interpretation is that Pre-LN's effective depth saturates earlier in training—it hits a ceiling on what its effectively shallower network can learn from the data, while Keel continues to improve because its full architectural depth is actually being utilized. This is presented as evidence that Keel is "particularly well suited for large-scale training, while its advantages are less pronounced in low-data regimes."
Important caveat: The 10B and 40B token budgets are small relative to the 250B–1T token experiments in other sections. This limits how strongly the data scaling results can be extrapolated. The Finding that Keel's advantage "widens with more data" is based on only two data points (10B and 40B), and the widening is modest (+1.2 to +1.5 average). A proper data scaling law experiment would require several orders of magnitude of data scaling (e.g., 10B, 100B, 1T) to establish a clear trend. The paper frames this as suggestive rather than definitive.
Deeper vs. Wider: Fixed Parameter Budget Comparison (Section 5.5)
Table 5 reports the critical test of whether depth scaling can deliver better performance than width scaling at the same parameter count—the experiment that directly challenges the conventional wisdom that wider models are more parameter-efficient. All configurations have 3B parameters and are trained on 250B tokens of private data. Learning rates are tuned per configuration: for both 512-layer models, for the 128-layer wide model.
Headline finding: The deep Keel model (512L, 1024 hidden) outperforms both the deep Pre-LN model and the wide Pre-LN model (128L, 2048 hidden) at the same parameter budget. Keel achieves an average score of 55.5 versus 52.3 for deep Pre-LN (+3.2) and 52.2 for wide Pre-LN (+3.3). The wide Pre-LN model—the architecture that represents the industry standard for stability—is the weakest of the three configurations on average score.
The composition of the gains reveals depth's particular value. On reasoning and math benchmarks: GSM-8K shows Keel 43.8 vs. wide Pre-LN 35.3 (+8.5) and deep Pre-LN 38.1 (+5.7). MBPP shows Keel 26.0 vs. wide 24.4 (+1.6) and deep 22.8 (+3.2). HumanEval shows Keel 19.5 vs. wide 16.5 (+3.0) and deep 17.7 (+1.8). On knowledge benchmarks, the gaps are smaller: MMLU shows Keel 56.3 vs. wide 51.5 (+4.8) and deep 52.9 (+3.4). ARC-Easy shows Keel 77.1 vs. wide 75.2 (+1.9) and deep 74.0 (+3.1).
An interesting sub-finding: the deep Pre-LN model (52.3 average) slightly outperforms the wide Pre-LN model (52.2) despite having a worse training loss (as shown in Figure 8a, discussed in Section 9). This is the training-loss/downstream-performance discrepancy discussed as Innovation 4: the deep model's loss is higher during training, but its downstream performance is comparable or slightly better, particularly on reasoning tasks. The authors interpret this as evidence that "increasing depth has the potential to improve complex reasoning when the model can be trained effectively." Keel realizes this potential by removing the training bottleneck.
Implication for architectural design: If deep-and-narrow models can match or exceed wide-and-shallow models at the same parameter count when properly stabilized, then the parameter-efficiency argument for width scaling weakens. Deep models have lower inference FLOPs per token (since each layer costs the same but there are fewer parameters per layer in a narrow configuration), so the combination of better downstream performance and lower inference cost would make deep Keel strictly preferable to wide Pre-LN at a given parameter budget—assuming the training stability can be maintained, which Keel demonstrates it can.
Large-Scale 1T-Token Pretraining and Supervised Fine-Tuning (Sections 5.6–5.7)
Table 6 reports the largest-scale experiment in the paper: 512-layer, 3B-parameter models trained on 1T tokens (750B general pre-training + 250B continued pre-training). This is a stress test of whether Keel's advantages persist at the scale of practical LLM training. The Pre-LN baseline is capped at a peak learning rate of due to instability (Figure 4a), while Keel trains stably at —a 50% higher learning rate.
Headline finding: Keel achieves a global average of 62.5 vs. Pre-LN's 58.7, a +3.8 point improvement across 14 benchmarks. The largest absolute gains are on reasoning-intensive tasks: GSM-8K: 60.9 vs. 51.0 (+9.9); AGI-Eval: 46.5 vs. 37.9 (+8.6); MBPP: 40.6 vs. 35.0 (+5.6); CommonsenseQA: 69.8 vs. 64.5 (+5.3). Smaller but consistent gains appear on knowledge benchmarks: MMLU: 62.7 vs. 59.5 (+3.2); ARC-Easy: 81.6 vs. 79.7 (+1.9); HellaSwag: 69.8 vs. 68.2 (+1.6).
The Winogrande result is notable: both models score 66.7—the only benchmark where Keel does not outperform Pre-LN. This task involves pronoun resolution, which may depend more on local linguistic pattern matching than on deep hierarchical reasoning, potentially explaining the lack of depth benefit.
Table 7 reports SFT results after both models undergo supervised fine-tuning on a high-quality instruction mix, with grid search over learning rates and epochs to find the best configuration per architecture. Headline finding: The pre-training advantage is preserved and in some cases amplified after SFT. Keel averages 54.8 vs. Pre-LN's 49.3 across the eight benchmarks evaluated (+5.5 points). On the "hard" benchmarks designed to probe capability limits: MMLU-Pro shows Keel 35.6 vs. Pre-LN 26.6 (+9.0); BBH shows 51.7 vs. 46.4 (+5.3); GSM-8K shows 68.8 vs. 58.7 (+10.1).
The amplification on MMLU-Pro (+9.0 post-SFT vs. +3.2 pre-SFT on standard MMLU) is particularly interesting—it suggests that Keel's better pre-training representations provide a stronger foundation for learning complex instruction-following behaviors during fine-tuning, and that the gap widens on tasks requiring nuanced reasoning rather than simple knowledge recall.
Ablation Studies and Robustness Checks
-
Learning rate tolerance across architectures (Table 1, Section 5.1): The Max LR comparison serves as an implicit ablation of the residual pathway design. Vanilla Post-LN ( at 64L), DeepNorm (), HybridNorm (), and Mix-LN () all achieve Max LRs roughly two orders of magnitude below Keel (). This demonstrates that hybrid placement strategies and initialization-dependent scaling (DeepNorm) do not address the root cause of Post-LN instability—the ResNet-style residual pathway's interaction with LayerNorm's Jacobian—while Keel's structural modification does. Pre-LN () comes closest to Keel's stability but still lags, and its depth utilization is worse (as shown in subsequent experiments).
-
Learning rate sweep for Pre-LN vs. Keel (Table 2, Section 5.2): The sweep across three learning rates reveals that Pre-LN's downstream performance responds non-monotonically and task-inconsistently to higher learning rates, while Keel's responds monotonically. This is a robustness check on the claim that Keel provides a more well-conditioned optimization landscape—not just higher stability ceilings but also more predictable scaling with the learning rate hyperparameter. The specific failure cases for Pre-LN at (ARC-Easy drops from 75.3 to 74.0; MBPP drops from 24.0 to 22.8) suggest that even when Pre-LN does not diverge, large learning rates cause destructive interference in specific parameter subspaces.
-
Depth sweep from 64 to 1024 layers (Table 3, Section 5.3): This is a robustness check on the claim that Keel's advantages scale with depth. The non-monotonicity at 128 layers (gap narrows to +1.2 before widening to +3.8 at 512L) suggests that at moderate depths, the benefit of Keel over Pre-LN may be partially offset by other factors (perhaps the higher learning rate used for both at 64L–128L favors Pre-LN, or the continued pre-training phase interacts differently with the two architectures at different depths). The larger point—that Keel's advantage is substantial at extreme depths—is consistently supported from 512L onward.
-
Data scaling from 10B to 40B tokens (Table 4, Section 5.4): This is a robustness check on whether Keel's benefits persist or vanish with increased data. The finding that the gap modestly widens (+1.2 to +1.5 average) suggests that Keel does not merely converge faster to the same asymptote—it converges to a better solution. However, the limited data range (only two data points: 10B and 40B tokens) and small absolute gap changes make this one of the weaker robustness checks in the paper.
-
Deeper vs. wider at fixed parameter budget (Table 5, Section 5.5): This serves as a robustness check on the claim that depth scaling is a viable alternative to width scaling when optimization is stabilized. The deep Keel model outperforming the wide Pre-LN model at the same 3B parameter count demonstrates that the historical advantage of width over depth is not a fundamental property of Transformers but an artifact of training instability. However, only one wide configuration (128L, 2048 hidden) is compared against one deep configuration (512L, 1024 hidden). A stronger check would sweep multiple (depth, width) pairs at the same parameter budget to identify the Pareto frontier.
-
Supervised fine-tuning preserves pre-training advantage (Table 7, Section 5.7): This is a robustness check on whether Keel's pre-training benefits transfer to downstream task performance after instruction tuning. The preservation and occasional amplification (MMLU-Pro: +9.0 post-SFT vs. +3.2 pre-SFT on standard MMLU) suggest that Keel's better pre-training representations create a stronger foundation for fine-tuning, rather than being overwritten by the SFT process.
-
Layer redundancy analysis (Section 8, Figures 6–7): Individual layers are removed from trained Pre-LN and Keel models, and perplexity increase is measured on C4. Removing the first layer destroys performance for both architectures (PPL > ), confirming that the first layer is exceptionally critical—analogous to attention sinks in sequence modeling. Beyond the first layer, Pre-LN shows smaller perplexity increases when shallow layers are removed compared to deep layers, indicating that shallow layers are more redundant in Pre-LN. Keel shows larger perplexity increases across a broader range of layers, indicating more uniform depth utilization. This directly validates the "effective depth" hypothesis: Pre-LN's gradient dynamics cause it to underutilize shallow layers, while Keel's gradient preservation enables more uniform learning across depth.
-
Training loss vs. downstream performance discrepancy (Section 9, Figure 8): At matched learning rates (), Pre-LN achieves slightly lower training loss than Keel (Figure 8b) yet Keel achieves better downstream performance (Table 2). Similarly, a wide Pre-LN model achieves lower training loss than a deep Pre-LN model (Figure 8a) yet scores similarly or worse on downstream tasks (Table 5). This is a robustness check on the use of training loss as a proxy for model quality in deep architectures: it demonstrates that the metric can be misleading, and the authors recommend end-task evaluation during pretraining.
Critical Assessment
The experiments broadly support the paper's central claims, but with important caveats about scope, scale, and what has actually been demonstrated versus what is implied.
Claim: Keel enables stable training at depths exceeding 1000 layers. Supported, but with a learning rate concession. Table 3 demonstrates successful training at 1024 layers for both Keel and Pre-LN, but Pre-LN required a lower peak learning rate ( vs. Keel's ) to avoid instability (Figure 4b). The claim that Keel "trains robustly at depths exceeding 1000 layers" is accurate—it does so at a higher learning rate than Pre-LN and with better downstream performance. However, the paper does not show successful training beyond 1024 layers, and the asymptotic claim in the theoretical analysis (gradient product converges to 1 as ) is validated only up to this finite depth. The gap between Keel and Pre-LN at 1024L (+3.0 average) is smaller than at 512L (+3.8), which could suggest either that the advantage begins to saturate at extreme depths or that the learning rate disadvantage imposed on Pre-LN at 1024L ( vs. Keel's ) accounts for part of the gap. The paper cannot fully disentangle these explanations because Pre-LN cannot be trained at the higher learning rate at 1024L—the instability is the point.
Claim: Keel consistently outperforms Pre-LN on downstream benchmarks, with gains that grow with depth. Supported, but the absolute magnitudes require careful interpretation. Every direct comparison between Keel and Pre-LN at matched parameter counts and token budgets shows Keel ahead (Tables 2, 3, 4, 5, 6, 7). The pattern of gains growing with depth is clear from 512L to 1024L for reasoning tasks, but the average score gap narrows from +3.8 at 512L to +3.0 at 1024L (Table 3)—the gap grows on GSM-8K but shrinks on average. This suggests that the "gains grow with depth" claim is task-dependent and may not hold uniformly across all benchmarks.
Claim: Keel achieves +16.5% improvement in Math & Code over Pre-LN. The specific figure requires context. The +16.5% figure appears in the abstract and Figure 1(b) but refers to a particular comparison that is not identified with a single table. Looking across the results: at 1024 layers (Table 3), Keel leads Pre-LN on GSM-8K by 58.6 vs. 49.8 (+17.7%), on HumanEval by 32.9 vs. 29.9 (+10.0%), and on MBPP by 38.6 vs. 36.2 (+6.6%)—an average of roughly +11.4% across math and code. At 512 layers with 1T tokens (Table 6), Keel leads on GSM-8K by 60.9 vs. 51.0 (+19.4%), HumanEval by 33.5 vs. 29.9 (+12.0%), and MBPP by 40.6 vs. 35.0 (+16.0%)—averaging roughly +15.8%. The +16.5% figure appears to be drawn from one of these groupings (likely the 1T-token 512L results) but is presented as a headline without specifying the exact comparison. This is a minor presentation issue but worth noting for reproducibility.
Claim: Keel revives Post-LN as a viable architecture for LLMs, not just as a hybrid or initialization-dependent fix. Supported in principle, but the experiments are limited to one model scale and family. All experiments use decoder-only Transformers at 3B parameters with hidden dimensions of 1024 or 2048, trained from scratch. The largest experiment is 1T tokens, which is substantial but small relative to production LLM training (multi-trillion tokens). The paper does not test Keel on models at the 7B, 13B, 70B, or larger scales where Pre-LN is most commonly used. The claim that Post-LN "is back" is supported for the scale regime tested but cannot be assumed to generalize to larger models without further evidence.
Significant weaknesses and missing experiments:
-
Single model family, single scale. All experiments use the same decoder-only Transformer architecture with the same hyperparameters (attention heads, RoPE theta, dropout rates—Table 8) at 3B parameters. There is no evidence that Keel's benefits transfer to different model families (e.g., mixture-of-experts architectures, different attention mechanisms beyond standard MHA) or to larger scales where different instability modes might dominate.
-
Private training data for most experiments. The primary experiments (Sections 5.1–5.3, 5.5–5.7) use "internal data" and "private data" that cannot be inspected or replicated. Only the data scaling experiments (Section 5.4) use a public dataset (FineWeb-EDU). This limits independent verification of the core results.
-
No statistical quantification of variance. As noted in the methodology section, all results represent single training runs. The downstream benchmark gaps are often small relative to what might be expected from run-to-run variance (e.g., +1.7 average gain at 64L in Table 3, +1.2 at 128L). Without confidence intervals or multiple seeds, it is impossible to determine whether these small gaps are statistically reliable or within noise. The larger gaps (+3.8 at 512L, +3.0 at 1024L) are more likely to be robust but still lack statistical quantification.
-
No comparison with Keel at different α values. The paper sets based on the theoretical analysis and does not report an ablation over different values (e.g., , , ) to demonstrate that is optimal rather than simply sufficient. The claim that "setting is critical for maintaining training stability in very large-scale or deep models" is asserted but not experimentally validated through an sweep. Such an ablation would substantially strengthen the theoretical-motivation-to-empirical-validation pipeline.
-
The depth-vs-width comparison is a single data point. Table 5 compares one deep configuration (512L, 1024 hidden) against one wide configuration (128L, 2048 hidden) at 3B parameters. A proper scaling law analysis would test multiple (depth, width) pairs at the same parameter count—e.g., (256L, 1408 hidden), (512L, 1024 hidden), (1024L, 724 hidden), etc.—to map out the Pareto frontier of depth vs. width under Keel. The paper does not establish whether 512L is optimal for 3B parameters or whether even deeper configurations would perform better.
-
No combination with other architectural innovations. Keel is tested as a modification to an otherwise standard Transformer. The paper does not explore whether Keel's benefits compound with or are made redundant by other architectural improvements (SwiGLU activations, grouped-query attention, different positional encodings, etc.) that are common in state-of-the-art LLMs.
-
Limited data scaling analysis. The data scaling experiment (Section 5.4) uses only 10B and 40B token budgets—far below the 250B–1T tokens used in other experiments. The finding that Keel's advantage widens with data is based on a single step from 10B to 40B and a small absolute gap change (+1.2 to +1.5 average). A proper data scaling law analysis would require at least three data points across orders of magnitude to establish a trend.
-
The SFT comparison may be confounded by learning rate effects. Table 7 shows Keel outperforming Pre-LN after SFT, but both models were pre-trained at different learning rates (Keel at , Pre-LN at ). The SFT learning rate sweep is performed independently, so the best SFT configuration per architecture is reported. However, the pre-training learning rate difference may affect the quality of the representations accessible during fine-tuning. A cleaner comparison would match pre-training learning rates where possible (even if this means using a suboptimal rate for Keel) to isolate the architectural effect from the learning rate effect.
Where the claims hold conditionally:
-
Keel's stability advantage holds at the learning rates and depths tested. The Max LR results (Table 1) are clear and consistent: Keel tolerates higher learning rates than all Post-LN variants and Pre-LN. The stability advantage at aggressive learning rates is the most robustly demonstrated claim in the paper.
-
Keel's downstream performance advantage holds at the scales tested (3B parameters, up to 1T tokens, up to 1024 layers). Every apples-to-apples comparison favors Keel, with larger margins on reasoning tasks than on knowledge tasks. The advantage is consistent across learning rates, depths, and data volumes, though the magnitude varies.
-
The claim that Post-LN can be rehabilitated holds for Keel's specific formulation at the tested scales. Whether other Post-LN formulations (with different Highway-style gating, different normalization placements, different scaling schemes) would also be rehabilitated is not tested. Keel is a specific architecture, not a proof that all Post-LN variants are now viable.
Experiments that would have strengthened the paper:
- An sweep () at a fixed depth to validate the choice experimentally.
- Multiple random seeds for the key depth scaling and deeper-vs-wider experiments to quantify variance.
- At least one experiment at a larger parameter scale (e.g., 7B or 13B parameters) to test whether Keel's benefits transfer to model sizes where Pre-LN is standard.
- A proper depth-vs-width Pareto frontier at fixed total parameters (sweeping multiple depth-width pairs, not just two configurations).
- A data scaling law experiment with at least three data points (e.g., 10B, 100B, 1T tokens) to establish whether Keel's advantage follows a power-law relationship with data volume.
- Comparison against Pre-LN at matched learning rates for at least one experiment, to isolate the architectural contribution from the learning rate contribution to downstream performance.
6. Limitations and Trade-offs
The α = L Setting Is Theoretically Motivated but Never Empirically Validated Against Alternative α Values
The assumption or constraint. The paper's central theoretical result — that Keel's cumulative gradient product converges to 1 as L → ∞ — depends critically on setting the Highway-style scalar gate to α = L (the total number of sublayers). Equation 19 (Section 3.3) derives this limit assuming α = L and all γ (LayerNorm affine weight) values remain near 1. The paper states:
"setting α = L is critical for maintaining training stability in very large-scale or deep models. For smaller architectures where vanishing or exploding gradients are less pronounced, α can be treated as a tunable hyperparameter (α > 1) to potentially accelerate convergence."
However, this claim — that α = L is critical versus other choices — is asserted from the theoretical derivation but never tested experimentally. No experiment sweeps alternative α values (e.g., α = √L as in DeepNorm, α = L/2, α = 2L, α = L²) at a fixed depth to demonstrate that α = L is optimal rather than simply sufficient.
The consequence. Without an α sweep, it is unknown whether Keel's stability and performance depend sensitively on the exact α = L setting, or whether a broad range of α values would work similarly well. If the former, practitioners must know L (total sublayer count) precisely before training and cannot modify the architecture depth without re-deriving α — a coupling that complicates architecture search. If the latter, the theoretical analysis is less prescriptive than the paper implies, and the specific α = L derivation may not be the active mechanism behind Keel's empirical success. More critically, the possibility that DeepNorm's α = L^0.25 could work equally well if combined with Keel's input normalization (the other structural modification) is never tested — the two modifications (Highway scaling and input normalization) are never ablated independently to determine their separate contributions.
What evidence exists in the paper. The design evolution narrative (Section 4.2) tests only α = L in all four attempts, providing no comparison against other α values. The DeepNorm comparison (Section 4.1) is between architectures with different residual scaling and different initialization schemes and different input normalization (DeepNorm lacks the inner LN), meaning the comparison cannot isolate the effect of α. The stability results in Table 1 compare Keel (α = L) against DeepNorm (α = L^0.25) and other baselines, but these differ in multiple architectural dimensions simultaneously. The paper provides no experiment where α is varied while holding all other Keel components fixed.
Mitigation status. Not addressed. The paper treats α = L as a derived constant rather than a tunable hyperparameter and does not acknowledge the absence of an α sweep as a gap in the empirical validation. The suggestion that α can be tuned for smaller models implies sensitivity, but this sensitivity is never characterized. A practitioner seeking to deploy Keel at a novel depth or scale would need to trust the theoretical derivation without empirical guidance on acceptable α ranges or the performance cost of suboptimal choices.
All Primary Results Use Private Training Data, Rendering the Core Experiments Non-Reproducible
The assumption or constraint. The paper's headline experiments — stability analysis (Section 5.1, Table 1), depth scaling (Section 5.3, Table 3), deeper-vs-wider comparison (Section 5.5, Table 5), and 1T-token pretraining (Section 5.7, Tables 6–7) — all use an internal private dataset described only as "internal data" or "private data." The data mixture, preprocessing pipeline, quality filtering, and domain composition are not disclosed. The only experiment using a public dataset is the data scaling study (Section 5.4, Table 4), which uses FineWeb-EDU at modest 10B and 40B token scales — far below the 250B–1T token budgets of the main experiments.
The consequence. Independent replication of the paper's core claims is impossible without access to the same training data. This is a significant barrier because training data composition strongly affects both training stability (data quality, noise levels, and domain diversity influence gradient variance and loss-spike propensity) and downstream benchmark performance (data overlap with evaluation benchmarks can inflate scores). If the private data contains even modest contamination from the MATH or code benchmarks used for evaluation, the reported improvements on GSM-8K, HumanEval, and MBPP would be overestimated relative to what a public-data replication would achieve.
Furthermore, the reported gap sizes between Keel and Pre-LN cannot be contextualized against the variance that arises from data composition choices. If the private data is unusually favorable to deep architectures (e.g., containing long-form reasoning traces that benefit from hierarchical representations), the Keel advantage may not generalize to other data distributions. The public-data experiment (FineWeb-EDU at 10B–40B tokens) shows a smaller average gap (+1.2 to +1.5 points, Table 4) than the private-data experiments (+3.2 to +3.8 points at comparable model sizes in Tables 5–6), though this could be due to the smaller token budget rather than data composition.
What evidence exists in the paper. The paper provides no characterization of the private data's composition, size distribution, or relationship to evaluation benchmarks. The 1T-token experiments in Section 5.7 are described as "general pre-training stage on the first 750B tokens, followed by continued pre-training (CPT) on the remaining 250B tokens to enhance reasoning and coding capabilities," but the nature of the CPT data and how it was selected is not disclosed. The FineWeb-EDU experiments (Table 4) serve as a partial sanity check that Keel's benefits are not purely an artifact of the private data, but the token scale is too small to verify the headline results.
Mitigation status. Not addressed and not acknowledged as a limitation. The paper does not discuss data transparency, does not commit to releasing the private data or a detailed data sheet, and does not characterize the data in sufficient detail for partial reproducibility (e.g., domain breakdown, deduplication status, contamination analysis against evaluation benchmarks). The reliance on private data is typical for industry LLM papers but limits the scientific verifiability of the claims.
Keel Has Not Been Tested at Parameter Scales Where Pre-LN Is the Dominant Architecture
The assumption or constraint. All experiments in the paper use models with approximately 3B parameters. The largest configuration tested is 512 layers with a hidden dimension of 1024 at 3B parameters (Section 5.7), or 1024 layers with the same hidden dimension (Section 5.3, also ~3B parameters given the depth-to-width relationship). The paper does not test Keel at 7B, 13B, 70B, or larger parameter scales where Pre-LN is the standard architecture in production LLMs. The experiments are restricted to what the paper calls an "extreme depth-to-width ratio" (512 layers vs. 1024 width), which is a topology specifically chosen to stress-test the optimization challenges that Keel addresses — but this topology is unusual for production LLMs, which typically have much larger width-to-depth ratios.
The consequence. The paper's central claim — "Post-LayerNorm Is Back" — implies that Keel enables Post-LN to replace Pre-LN in practical LLM training. But the regime where this substitution would be most impactful is at larger parameter scales (7B+), where Pre-LN's effective-depth limitation most constrains model quality and where training instability would be most costly. At 3B parameters, the results demonstrate a proof of concept but do not establish that Keel scales to the parameter counts used in state-of-the-art LLMs.
There are specific reasons why larger-scale behavior might differ. At larger hidden dimensions, the variance dynamics that Keel's input normalization controls may change — wider layers produce transformations with different variance characteristics, and the gradient analysis in Section 3.3 assumes the γ values remain near 1, an assumption that may break down when attention and FFN layers have very different output magnitudes at larger widths. Additionally, at larger parameter scales, training is often distributed across many devices, and communication patterns could interact with Keel's additional normalization steps (the inner LN before the transformation branch) in ways that affect throughput or introduce numerical issues not present at 3B scale.
What evidence exists in the paper. None. The paper does not include experiments at parameter scales larger than 3B, nor does it discuss how Keel's architectural choices might interact with model width as the parameter budget increases. The width-to-depth ratio tradeoff is explored only in one direction — making the model deeper at fixed parameters (Section 5.5) — but not in the other direction (making the model larger while holding the depth-to-width ratio constant, which is the typical scaling path for production LLMs). The model configuration table (Section 10, Table 8) lists only two hidden dimension sizes (1024 and 2048), both corresponding to 3B parameter configurations.
Mitigation status. Partially acknowledged in Section 7 (Limitation and Future Work), which states:
"As model width increases (e.g., hidden size, number of experts, or FFN dimension), LLM training can also become more unstable. In such wider settings, Post-LN may require a larger α (or stronger stabilization mechanisms) to maintain stable optimization. We leave a thorough investigation of LLM stability under width scaling to future work."
This acknowledgment is general and does not commit to specific scale targets. It also frames width scaling as a separate future investigation rather than as a gap in the current paper's evidence for the headline claim. The paper's title and abstract assert that Post-LN is now viable — without qualification about parameter scale — but the experiments only support this for the narrow, deep topology tested.
Keel's Advantages Are Concentrated on Reasoning Benchmarks; Gains on Factual and Linguistic Tasks Are Modest and Potentially Within Run-to-Run Variance
The assumption or constraint. The paper consistently reports larger absolute improvements on reasoning-intensive benchmarks (GSM-8K, AGI-Eval, CommonsenseQA) than on knowledge-intensive or linguistic pattern-matching benchmarks (MMLU, HellaSwag, ARC-Easy, LAMBADA, Winogrande). This pattern holds across depths and training budgets: at 512 layers after 1T tokens (Table 6), the reasoning gap is +9.9 on GSM-8K, +8.6 on AGI-Eval, and +5.3 on CommonsenseQA, while the knowledge gap is +3.2 on MMLU, +1.9 on ARC-Easy, +1.6 on HellaSwag, and 0.0 on Winogrande. At 64 and 128 layers (Table 3), several knowledge benchmarks show Keel-PreLN differences of less than 2 points.
The consequence. For practitioners whose primary use case is factual QA, language understanding, or pattern recognition rather than multi-step mathematical reasoning or code generation, the case for adopting Keel over Pre-LN is substantially weaker. A +1.6 point improvement on HellaSwag at 512 layers (Table 6) or a +1.7 point average improvement at 64 layers (Table 3) may not justify changing a production architecture, especially given Keel's additional implementation complexity (two modifications per sublayer) and the lack of evidence at larger parameter scales.
More fundamentally, the paper does not investigate whether the reasoning-specific gains arise from Keel's depth utilization (more uniform gradient propagation through all layers) or from the higher learning rates that Keel enables. If the latter — and if a Pre-LN model could achieve similar reasoning performance through a different learning rate schedule or optimization trick that does not require architectural changes — then Keel's value proposition narrows to being a mechanism for enabling higher learning rates without instability, rather than a fundamentally more expressive architecture.
The small absolute gaps on knowledge benchmarks also raise concerns about statistical reliability. The paper reports single training runs for all experiments, with no variance estimates. A gap of +1.6 on HellaSwag or +1.9 on ARC-Easy could plausibly arise from random seed variation, data ordering, or hardware non-determinism. Without multiple seeds or confidence intervals, these small gains cannot be confidently attributed to the architectural change.
What evidence exists in the paper. The pattern of reasoning > knowledge gains is visible across all result tables: Table 2 (learning rate sweep), Table 3 (depth scaling), Table 5 (deeper vs. wider), Table 6 (1T-token pretraining), and Table 7 (SFT). The paper does not provide a systematic breakdown or discuss this differential effect explicitly, but it is consistent enough to constitute a robust empirical regularity within the tested regime. Winogrande's 0.0 gap (Table 6) and the near-zero gaps on several knowledge benchmarks at shallower depths (Table 3) suggest that Keel's benefits are task-family-dependent.
Mitigation status. Not addressed. The paper presents Keel as a general architectural improvement for deep Transformers but does not characterize which task families benefit most or why. The abstract's +16.5% Math & Code figure selectively highlights the domain where gains are largest, while the average scores (which include knowledge benchmarks where gains are smaller) are reported separately without discussing the discrepancy. A practitioner evaluating Keel for a non-reasoning application would need to extrapolate from the reported knowledge-benchmark gaps, which are small and of uncertain statistical significance.
The Paper Provides No Analysis of Training Throughput, Memory Overhead, or Inference Cost
The assumption or constraint. Keel adds two structural modifications to every Transformer sublayer: a scalar multiplication (α) on the residual branch and an additional LayerNorm operation on the input to the transformation function. The paper's experiments measure training stability (Max LR), downstream accuracy, and training loss, but provide no measurements of wall-clock training time, GPU memory consumption, or inference latency relative to Pre-LN. The computational cost of the additional LayerNorm and the α scaling operation is assumed to be negligible, but this assumption is never quantified.
The consequence. For a practitioner deciding whether to adopt Keel, the throughput and memory implications matter. The additional LayerNorm per sublayer requires (a) computing the L2 norm of the input vector (a reduction operation over the hidden dimension), (b) element-wise division by this norm, and (c) element-wise multiplication by the learnable γ vector. For a 512-layer model with 1024 sublayers (counting both attention and FFN), this adds 1024 extra normalization operations per forward pass. While a single LayerNorm is cheap relative to attention or FFN computation, the cumulative overhead across 1024 sublayers may be non-trivial — particularly in inference settings where latency is critical and normalization operations on small vectors can be bound by memory bandwidth rather than compute.
Additionally, the inner LayerNorm's input must be stored for backpropagation (since LayerNorm's Jacobian depends on the input norm and the pre-normalized values), increasing activation memory relative to a Pre-LN model where the transformation branch receives the unnormalized residual-stream value (which is already stored). For very deep models, activation memory is often the binding constraint on batch size and training throughput.
The paper's reliance on "generations" or "tokens" as compute metrics (rather than FLOPs or wall-clock time) obscures these overheads. If Keel requires, say, 5–10% more wall-clock time per training step than Pre-LN at the same parameter count, the reported accuracy gains must be discounted accordingly to assess whether Keel provides a net efficiency improvement.
What evidence exists in the paper. None. The paper does not report training throughput (tokens/second), peak GPU memory usage, or inference latency for any configuration. The computational cost model is entirely absent — the paper focuses exclusively on optimization dynamics (gradient flow, learning rate tolerance) and downstream accuracy. There is no FLOPs analysis comparing Keel and Pre-LN at matched parameter counts.
Mitigation status. Not addressed and not acknowledged as a limitation. The paper treats the additional LayerNorm as costless because normalization is typically a small fraction of total Transformer computation. However, at the extreme depths studied (up to 1024 layers = up to 2048 sublayers), small per-layer overheads compound, and the paper's omission of any throughput analysis leaves a gap in the practical evaluation of Keel's viability.
7. Implications and Future Directions
How This Work Changes the Landscape
This paper shifts the diagnosis of deep Transformer training instability from a forward-pass problem (activation magnitude explosion) to a backward-pass problem (gradient vanishing through the LayerNorm Jacobian), and in doing so reopens an architectural design space—Post-LN—that the field had largely abandoned in favor of Pre-LN. This is best understood as a diagnostic reframing with architectural consequences rather than a paradigm shift: the fundamental mechanisms (LayerNorm, residual connections, gradient backpropagation) are unchanged, but the paper demonstrates that the community's understanding of why Post-LN fails was incomplete, and that fixing the true cause (the ResNet-style residual pathway's interaction with LayerNorm's Jacobian) yields a practical architecture that combines Post-LN's expressivity with Pre-LN's stability.
The magnitude of the shift is moderate but targeted. The paper does not propose a new class of models or a new training paradigm—it modifies two lines of a Transformer sublayer's forward pass. Yet these two lines address what has been, since Xiong et al. (2020), the binding constraint on depth scaling in Transformers. The shift is from viewing Post-LN's instability as an intrinsic property of normalizing-after-residual (and therefore accepting Pre-LN as the necessary compromise) to viewing it as a specific, fixable pathology of how the residual and transformed branches are mixed before normalization.
The paper reconciles a tension that has persisted in the literature since the original Transformer. The 2017 Transformer used Post-LN and demonstrated strong results at 6–12 layers, but attempts to scale this architecture deeper encountered training instability that Pre-LN resolved. The field interpreted this as evidence that Pre-LN is fundamentally more stable, and that Post-LN's stronger inter-layer coupling (which gradients must pass through) comes at an unavoidable cost in optimization difficulty. Keel demonstrates that this tradeoff is not fundamental—the gradient vanishing arises specifically from the equal-weight summation in the ResNet-style residual path, not from the mere fact of normalizing after the addition. By weighting the residual branch more heavily (α = L), the gradient through the LayerNorm Jacobian can be preserved without sacrificing the inter-layer coupling that gives Post-LN its expressivity advantage.
This reframing has several consequences for the research landscape:
The ML community gains a more precise diagnostic framework for Transformer training instability. Prior work diagnosed Post-LN failure through forward variance (DeepNorm) or through gradient norm monitoring without identifying the specific mechanism. The paper's decomposition—showing that the LayerNorm Jacobian attenuates gradients by ~1/√2 per layer due to the geometry of summing two roughly-orthogonal vectors—provides a concrete target for future stabilization efforts. Researchers developing new normalization schemes or residual pathways can evaluate their proposals against this gradient-flow analysis rather than relying solely on empirical stability tests.
The design space for deep Transformers expands beyond Pre-LN and its hybrids. Prior to Keel, the viable options for stable Transformer training at scale were essentially: (1) Pre-LN, with its known effective-depth limitation; (2) hybrid approaches (HybridNorm, Mix-LN) that interleave or transition between Post-LN and Pre-LN blocks but do not fully resolve the gradient issue; or (3) DeepNorm-style initialization-dependent scaling, which the paper shows does not hold up at LLM scale. Keel introduces a fourth option: structurally stabilized Post-LN that does not depend on initialization or block-placement heuristics. This means depth scaling research can now operate from a Post-LN baseline that trains as stably as Pre-LN but with better depth utilization, rather than being forced to accept Pre-LN's expressivity penalty.
Depth scaling becomes empirically viable as a primary axis for improving model quality. The paper's deeper-vs-wider experiment (Table 5) challenges the conventional wisdom that width scaling is more parameter-efficient than depth scaling. The deep Keel model (512L, 1024 hidden) outperforms the wide Pre-LN model (128L, 2048 hidden) at the same 3B parameter count by +3.3 points average, with larger margins on reasoning tasks (+8.5 on GSM-8K). If this result generalizes to larger scales, it suggests that the field's historical preference for width over depth was an artifact of Pre-LN's poor depth utilization, not a fundamental property of Transformer expressivity. This could redirect scaling research toward depth as a first-class axis alongside width and data volume.
The role of learning rate in architecture comparison is elevated. The paper demonstrates that Keel's stability advantage manifests primarily as the ability to use higher learning rates (Table 1: +32% at 64L, +35% at 512L over Pre-LN), and that these higher learning rates translate to better downstream performance (Table 2: Keel at η=6.0×10⁻³ averages 55.5 vs. Pre-LN's best at 52.3). This implies that architecture comparisons that fix learning rates across all methods may systematically underestimate the advantage of more stable architectures. Future architecture research should report per-method optimal learning rates and treat learnable rate tolerance as a legitimate architectural advantage rather than a confound to be controlled away.
The training loss vs. downstream performance discrepancy (Section 9) introduces a methodological caution for depth scaling research. The finding that Pre-LN can achieve lower training loss than Keel at matched learning rates (Figure 8b) yet underperform on downstream tasks (Table 2) challenges the standard practice of using validation perplexity as the primary proxy for model quality during architecture development. If this phenomenon is general—if better depth utilization systematically produces representations that generalize better to reasoning tasks even when next-token prediction loss is slightly higher—then architecture comparisons based solely on loss curves may be systematically biased toward shallower effective-depth architectures. This is a methodological contribution that affects how future depth-scaling experiments should be evaluated.
Follow-Up Research This Work Enables
Ablation of α to determine the sensitivity of Keel's stability to the Highway scaling factor. The paper's central theoretical result (Equation 19) derives gradient preservation under α = L, but this choice is never empirically validated against alternatives. A critical follow-up experiment would sweep α ∈ {1, √L, L/2, L, 2L, L²} at a fixed depth (e.g., 256 layers) and measure both Max LR (Section 5.1 protocol) and downstream benchmarks (Section 5.2 protocol) to determine whether α = L is actually optimal or merely sufficient. If stability is insensitive to α above some threshold (say, α ≥ L/2), then the theoretical derivation overstates the importance of the exact scaling law, and practitioners can tune α as a hyperparameter without risking instability. If stability is sharply peaked at α = L, then the theoretical coupling between depth and residual scaling is genuinely binding, with implications for how Keel-based architectures can be modified. A strong version of this experiment would also test whether DeepNorm's α = L^0.25 can match Keel's stability when combined with Keel's inner LayerNorm (the other structural modification), isolating the contribution of the Highway scaling from the input normalization.
Depth-vs-width Pareto frontier at fixed parameter budgets to establish optimal architecture allocation under Keel. The paper's deeper-vs-wider experiment (Table 5) compares only two configurations (128L/2048 hidden vs. 512L/1024 hidden) at 3B parameters. A proper scaling law study would train multiple (depth, width) pairs at the same total parameter count—e.g., (64L, 2816 hidden), (128L, 2048 hidden), (256L, 1408 hidden), (512L, 1024 hidden), (1024L, 724 hidden)—all using Keel's stabilization, and measure how downstream performance varies with depth-to-width ratio. This would answer whether the 512L configuration tested in the paper is near-optimal for 3B parameters, or whether even deeper configurations (1024L) would perform better at this scale. If the optimal depth-to-width ratio increases with total parameter count, this would suggest that larger models should be systematically deeper than current practice, providing a concrete design guideline for next-generation LLM architectures. The experiment would also reveal whether the reasoning-specific benefits of depth (larger gaps on GSM-8K, AGI-Eval) continue to grow with depth or saturate at some threshold.
Scaling Keel to 7B+ parameters with public data to test whether Post-LN can genuinely replace Pre-LN at production scale. The paper's experiments are restricted to 3B parameters and primarily use private training data. The most important follow-up for establishing Keel's practical relevance is a reproduction at a larger scale—ideally 7B or 13B parameters—trained on a fully public dataset (e.g., FineWeb-Edu, Dolma, or RedPajama) with a public training recipe, and evaluated on standard open benchmarks via lm-evaluation-harness. This would address three gaps simultaneously: (1) whether Keel's stability advantage persists at parameter scales where Pre-LN is the incumbent standard; (2) whether the reported gains can be reproduced on public data, eliminating the concern that private data composition inflates the advantage; and (3) whether the implementation complexity of Keel (two modifications per sublayer) introduces throughput or memory overheads that offset its accuracy gains at scale. A strong version of this experiment would report training throughput (tokens/second), peak GPU memory, and inference latency alongside accuracy, and would include at least three random seeds to quantify variance on the smaller knowledge-benchmark gains.
Combining Keel with other architectural innovations to test for complementarity or redundancy. The paper tests Keel on an otherwise standard Transformer (RMS LayerNorm, standard MHA, standard FFN). Modern LLMs incorporate numerous additional improvements—SwiGLU activations, grouped-query attention (GQA), rotary position embeddings with extended contexts, mixture-of-experts (MoE) layers, and different normalization schemes (e.g., QK-Norm in some architectures). Several of these innovations affect the same gradient pathways that Keel modifies. For example, GQA changes the dimensionality of key-value projections relative to query projections, which could alter the variance characteristics of the attention transformation branch that Keel's inner LayerNorm is designed to control. MoE architectures route tokens to different FFN experts based on the input, introducing stochasticity in the transformation branch that could interact with the Highway-style gating. A systematic study combining Keel with these innovations would determine whether Keel's benefits are additive to, redundant with, or in tension with other architectural improvements that have become standard since the basic Transformer architecture Keel was tested on.
Data scaling laws for Keel vs. Pre-LN across orders of magnitude to test the "effective depth" hypothesis. The paper's data scaling experiment (Section 5.4) uses only two token budgets (10B, 40B) and shows a modest widening of the Keel-PreLN gap from +1.2 to +1.5 average points. The paper interprets this as evidence that Pre-LN's lower effective depth causes it to saturate earlier, while Keel continues to improve. Testing this interpretation properly requires a data scaling law experiment with at least four data points spanning two or more orders of magnitude (e.g., 10B, 100B, 300B, 1T tokens), measuring both training loss and downstream benchmarks for Keel and Pre-LN at matched depth and width. If Pre-LN's performance plateaus earlier than Keel's as data increases—consistent with a model whose effective capacity is saturated—this would provide the strongest evidence that Pre-LN's depth utilization problem is a binding constraint on how much models can learn from large datasets, and that Keel's architectural fix has growing returns with data scale. Conversely, if the gap remains constant or narrows with data, then Pre-LN's effective depth is sufficient to absorb large-scale training data, and Keel's advantage is primarily in training dynamics (faster convergence at higher learning rates) rather than asymptotic representational capacity.
Layer-importance dynamics during training to understand when and how Keel's depth utilization advantage develops. The paper's layer redundancy analysis (Section 8, Figures 6–7) measures post-hoc layer importance by removing individual layers from trained models and measuring perplexity increase. This is a static snapshot of the final trained state. A dynamic version of this experiment would checkpoint models periodically during training (e.g., at 1%, 10%, 50%, 100% of total tokens) and measure layer importance at each checkpoint for both Keel and Pre-LN. This would reveal when in training the effective-depth divergence emerges: does Pre-LN start with uniform layer utilization that degrades over time as gradient vanishing accumulates, or is the shallow-layer redundancy present from early in training? If the former, it suggests that architectural stabilization is most important for long training runs; if the latter, it suggests that even moderate training budgets benefit from Keel's gradient preservation. This experiment would also test whether the "training loss vs. downstream performance discrepancy" (Section 9) emerges gradually during training or is present from the start, which would inform whether validation-loss-based early stopping is safe for deep models.
Practical Applications and Downstream Use Cases
Training deep-and-narrow models for on-device or low-latency deployment where inference cost is the binding constraint. The deeper-vs-wider experiment (Table 5) demonstrates that at 3B parameters, a deep Keel model (512 layers, 1024 hidden) achieves higher accuracy than a wide Pre-LN model (128 layers, 2048 hidden) by +3.3 points on average and +8.5 points on GSM-8K. For deployment scenarios where inference latency or per-token cost is limited—mobile devices, real-time APIs, edge inference—the deep-and-narrow topology offers a direct advantage: each forward pass processes a smaller hidden dimension (1024 vs. 2048), reducing the FLOPs per token in the attention and FFN operations (which scale with hidden dimension squared). The deeper layer count adds sequential computation (more layers to pass through) but not more computation per layer. If inference is compute-bound rather than memory-bandwidth-bound, the deep Keel model will be faster per token than the wide model at the same total parameter count, while achieving better downstream accuracy—a pure Pareto improvement. This use case is most compelling for reasoning-heavy applications (math tutoring, code assistance, logical QA) where Keel's largest accuracy gains are concentrated (+8.5 to +9.9 points on GSM-8K across Tables 5–6).
Large-scale pretraining with higher learning rates to reduce total training time. Keel's stability at elevated learning rates (Table 1: +32% Max LR over Pre-LN at 64L, +35% at 512L) means that practitioners can increase the learning rate without risking the irrecoverable loss spikes that constrain Pre-LN training (Figure 4). If higher learning rates translate to faster convergence—which the learning rate sweep (Table 2) suggests they do, with Keel improving monotonically from η=1.5×10⁻³ to 6.0×10⁻³—then total training tokens (and thus GPU-hours) can be reduced for a target downstream accuracy. For example, if a Pre-LN model requires 1T tokens at η=3.0×10⁻³ to reach a certain accuracy, a Keel model at η=4.5×10⁻³ (50% higher) might reach the same accuracy in fewer tokens, amortizing the per-step overhead of Keel's additional LayerNorm operations. This use case is most relevant for organizations that pretrain LLMs from scratch and for whom training budget (total FLOPs or GPU-hours) is the primary cost, rather than inference cost.
Continued pretraining of deep models on specialized reasoning corpora where depth utilization is critical. The paper uses a two-phase training protocol: general pretraining (750B tokens) followed by continued pretraining (250B tokens) on data designed to "enhance reasoning and coding capabilities" (Section 5.6). The large gains on GSM-8K (+9.9 points at 1T tokens, Table 6) and the amplification under SFT (Table 7, +10.1 on GSM-8K) suggest that Keel's improved depth utilization is particularly valuable during the continued-pretraining phase, where the model must learn to compose its existing knowledge into multi-step reasoning chains. Practitioners training domain-specific models (medical reasoning, legal analysis, scientific problem-solving) could adopt Keel specifically for the continued-pretraining stage—even if the base model is Pre-LN—by fine-tuning an existing checkpoint into a deep Keel configuration (though the paper does not test this transfer scenario, and weight compatibility between Pre-LN and Keel architectures would need to be addressed). The key benefit is that Keel's deeper effective depth allows the model to actually use the additional layers it has for learning hierarchical reasoning patterns, rather than having those layers become redundant as in Pre-LN (Section 8).
Architecture search over depth-to-width ratios with Keel as the stable Post-LN backbone. Before Keel, architecture search over Transformer depth was limited by the knowledge that very deep models would either be unstable (Post-LN) or suffer from poor depth utilization (Pre-LN). Keel removes this constraint, enabling practitioners to search over a wider range of depth-to-width ratios at a given parameter budget without worrying that deep configurations will be inherently disadvantaged by optimization dynamics. A concrete workflow: for a target parameter budget (e.g., 7B parameters), train several Keel configurations at different depths (e.g., 128L, 256L, 512L, 1024L) with adjusted hidden dimensions to hold parameters constant, and select the depth that maximizes downstream performance on the target task distribution. The paper's results suggest that for reasoning-heavy distributions, deeper configurations will be favored; for factual-recall distributions, the optimal depth may be shallower. This architecture search was previously impractical because deep configurations could not be fairly compared against shallow configurations under the same optimization budget (the deep models would train poorly regardless of their representational potential). Keel makes the comparison meaningful by stabilizing the deep models to the same level as shallow ones.
When to Prefer This Method
The paper positions Keel as a replacement for Pre-LN in deep Transformer training, with specific conditions under which the benefits are largest. These tradeoffs are articulated across the results and limitations sections:
-
Prefer Keel over Pre-LN when training decoder-only Transformers at depths exceeding approximately 100 layers (200 sublayers), especially when the target application involves multi-step reasoning, mathematical problem-solving, or code generation—domains where Keel's accuracy gains are largest (+5 to +10 points on GSM-8K, +3 to +6 points on MBPP and HumanEval across experiments). The stability advantage (higher tolerable learning rate, no catastrophic loss spikes) provides operational benefits even for knowledge-intensive applications, but the accuracy benefit on those tasks is modest (+1.5 to +3 points) and may not justify the architectural change alone.
-
Prefer Keel over Pre-LN when the training budget is large (hundreds of billions of tokens or more) and training loss shows signs of saturation. The data scaling experiment (Table 4) and the 1T-token results (Table 6) suggest that Keel's gap over Pre-LN widens with data volume, consistent with Pre-LN hitting an effective-depth ceiling that Keel avoids. In low-data regimes (tens of billions of tokens), Keel's advantage is smaller (Table 4, 10B tokens: +1.2 average) and may not offset the implementation complexity.
-
Prefer Pre-LN over Keel when the width-to-depth ratio is already large (e.g., hidden dimension of 4096+ with fewer than 100 layers). Section 7 notes that Keel "typically requires substantial training data to be effective and is not recommended in low-data regimes" and that "when the width-to-depth ratio is already high, [the effective-depth] issue may be less pronounced, and the gains from Keel may therefore be less substantial." In these regimes, Pre-LN's effective-depth problem is less binding because the model's representational capacity is dominated by its wide layers, and the additional implementation and potential throughput overhead of Keel's dual LayerNorms may not be justified.
-
Prefer Pre-LN or standard practice when inference latency is the absolute binding constraint and the additional per-layer computation of Keel's inner LayerNorm—even if small in absolute terms—is unacceptable. The paper provides no throughput measurements, but the structural addition of an extra normalization per sublayer is nonzero overhead. For applications where microsecond-level latency matters (high-frequency trading, real-time speech systems), this overhead may be disqualifying without further optimization.