ArXiv: 2412.13795
🎯 Pitch
Deep layers in modern LLMs are largely dead weight, contributing almost nothing to performance—but this paper shows it’s not an architectural flaw, it’s a training failure caused by Pre-LN killing gradient flow. By simply applying Post-LN to early layers and Pre-LN to deeper ones, Mix-LN resurrects these dormant layers, yielding consistent perplexity gains up to 7B scale and boosting downstream accuracy by over 4%.
1. Executive Summary
This paper introduces Mix-LN, a hybrid normalization technique that combines Post-Layer Normalization in earlier layers with Pre-Layer Normalization in deeper layers to address the widespread underutilization of deep layers in Large Language Models. The authors demonstrate that Pre-LN—the dominant choice in models like GPT and LLaMA—causes diminished gradient norms in deeper layers, reducing their effectiveness, while Post-LN preserves deeper gradients but suffers from vanishing gradients in early layers. By applying Post-LN to the first 25% of layers and Pre-LN to the remainder, Mix-LN achieves more uniform gradient norms across the network, yielding consistent perplexity improvements over Pre-LN across model sizes from 70M to 7B parameters (e.g., 21.39 vs. 21.92 for LLaMA-250M, 18.18 vs. 18.65 for LLaMA-1B) and substantial gains during supervised fine-tuning (e.g., +4.26% average accuracy on LLaMA-250M across eight downstream tasks), establishing that deep layers can contribute meaningfully to model capacity when the normalization strategy preserves gradient flow through all depths.
2. Context and Motivation
The Core Problem: Deep Layers in Modern LLMs Are Underutilized
The fundamental observation motivating this paper is both counterintuitive and practically significant: in modern LLMs trained with Pre-LN, the deeper layers—those furthest from the input—contribute remarkably little to the model's output. Recent work has shown that these layers can be pruned aggressively (Yin et al., 2023), removed entirely (Gromov et al., 2024; Men et al., 2024), or even swapped without substantially degrading model performance (Lad et al., 2024). For instance, Gromov et al. (2024) demonstrated that removing up to 20–30% of the deepest layers from models like LLaMA-2 yields negligible impact on downstream task accuracy.
This observation has been received primarily as an opportunity for model compression—a way to reduce inference costs by discarding redundant computation (Siddiqui et al., 2024; Zhong et al., 2024; Sreenivas et al., 2024). The authors argue that this framing is fundamentally mistaken. As they state in the introduction:
"Unfortunately, this finding has been largely overlooked by the research community, where many see it primarily as an opportunity for model compression, rather than recognizing it as a potential shortfall in the training process."
The paper's central claim is that deep-layer inefficacy is not an architectural inevitability but a training artifact—specifically, a consequence of the layer normalization choice. This reframes the problem from "how can we compress away useless layers?" to "how can we train so that all layers contribute meaningfully?"
Why This Matters: The Economic and Scientific Stakes
The underutilization of deep layers carries substantial costs along multiple dimensions:
Economic waste in pretraining. Training frontier LLMs requires enormous computational resources—thousands of GPUs running for months, with costs reaching millions of dollars (the paper cites GPT-3's reported costs as representative). If a substantial fraction of model parameters (those in deep layers) are contributing minimally to the learned function, then a significant portion of that investment is wasted. The paper frames this crisply:
"Ideally, all layers in a model should be well-trained, with sufficient diversity in features from layer to layer, to maximize the utility of the network's parameters."
This is not merely an aesthetic preference for "all parameters doing something"—it represents a direct efficiency argument. If deep layers could be made effective, the same parameter count would yield strictly higher performance, equivalent to getting more model capacity for the same training budget.
Downstream transfer consequences. The paper highlights work by Li et al. (2024) showing that deeper layers contribute minimally to performance even during fine-tuning. This means the problem compounds: not only do under-trained deep layers fail during pretraining, but they also cannot be "resuscitated" by supervised fine-tuning or RLHF. A model with poorly trained deep layers has permanently reduced capacity for adapting to downstream tasks, since those layers have not developed rich, transferable representations.
A missed opportunity for capability improvement. The authors position deep-layer inefficacy as a capacity ceiling: current LLMs are essentially operating with fewer effective layers than their architecture provides. If this ceiling can be lifted—if the deepest layers can be trained to contribute meaningful transformations rather than near-identity mappings (as the gradient analysis in Section 2.1 suggests they become under Pre-LN)—then model quality should improve without requiring larger architectures or more training data.
Prior Work and Its Limitations
The paper identifies three strands of prior work that intersect with this problem, each with specific gaps:
Layer normalization analysis (Xiong et al., 2020; Takase et al., 2022, 2023). The theoretical understanding of how different LN placements affect gradient flow has existed for several years. Xiong et al. (2020) provided the key insight: in Post-LN, gradients are multiplied by the Jacobian of each layer normalization, which has a spectral norm less than 1 when the input variance exceeds 1. This causes gradient vanishing in early layers. In Pre-LN, the residual connection is additive and unbounded, so early layers receive healthy gradients, but as outputs accumulate through successive residual connections, the Pre-LN operation effectively becomes an identity mapping in deeper layers—the layers stop learning. This analysis explains why deep layers under Pre-LN are ineffective, but Xiong et al. did not propose a hybrid approach that applies different normalizations at different depths.
Stabilization-focused LN variants (Admin, DeepNorm, Sandwich-LN, B2T). Several methods have attempted to make Post-LN more stable or to improve deep-layer training. Admin (Liu et al., 2020) adds extra parameters to control residual dependencies in Post-LN. DeepNorm (Wang et al., 2024) upscales the residual connection before LN to reduce model updates, enabling deeper Post-LN architectures. Sandwich-LN (Ding et al., 2021) normalizes both the input and output of each sub-layer. B2T (Takase et al., 2022) introduces a residual connection that bypasses all LN except the final one to address early-layer gradient vanishing in Post-LN. The critical limitation of these approaches is that they all treat the normalization strategy uniformly across depth—applying the same modified scheme to every layer. None exploit the depth-dependent complementarity of Pre-LN and Post-LN.
Layer similarity and pruning studies (Yin et al., 2023; Gromov et al., 2024; Men et al., 2024). This line of work documents the empirical fact of deep-layer ineffectiveness and uses it to guide pruning decisions. Gromov et al. (2024) introduced angular distance as a metric for measuring layer similarity, which this paper adopts (Section 2.2). However, these studies treat the phenomenon as an inherent property to be exploited (for compression) rather than as a training pathology to be fixed. The paper explicitly positions itself against this framing, arguing that the research community has been "band-aiding" the symptom rather than investigating the cause.
The Missing Piece: A Depth-Aware Normalization Strategy
The paper's key insight—and what distinguishes it from all prior work—is that the optimal LN strategy is not uniform across depth. Pre-LN and Post-LN have complementary failure modes that are depth-dependent:
- Pre-LN fails deep: gradients decay with depth, deep layers become near-identity mappings, representations become highly similar to neighboring layers.
- Post-LN fails early: gradients vanish in early layers due to accumulated LN Jacobian scaling, leading to training instability and divergence in larger models.
This complementarity directly suggests the hybrid approach: use Post-LN where it works well (early layers, where the LN Jacobian doesn't accumulate too many multiplicative factors) and Pre-LN where Post-LN would fail (later layers, where the accumulated attenuation would be catastrophic). The authors state this motivation explicitly:
"Post-LN has been shown to improve the effectiveness of deeper layers, while Pre-LN is more effective for earlier layers. Therefore, we propose to apply Post-LN to the initial layers and Pre-LN to the later layers, ensuring that the middle and deeper layers benefit from the advantages of both methods."
The paper also connects this to a more fundamental principle: the goal is not merely to avoid collapse or divergence, but to promote representation diversity across layers. Quoting Yang et al. (2023):
"Representation should change substantially from layer to layer in order to most effectively make use of the parameters of a network."
Mix-LN is positioned as the natural normalization strategy for achieving this ideal, by ensuring gradient flow is sufficient at all depths for each layer to learn distinct transformations.
How the Paper Validates and Motivates the Problem
Before proposing Mix-LN, the paper conducts a systematic hypothesis evaluation (Section 2) that serves as both validation of the diagnosis and motivation for the solution. The approach has two parts:
Controlled comparison using in-house models. The authors train LLaMA-130M models from scratch where the only difference is the normalization type (Pre-LN vs. Post-LN), using identical architectures, training data (C4), and hyperparameters. This isolates the normalization effect from confounds present in open-weight model comparisons.
Two-pronged evaluation metrics. They measure layer effectiveness through both representational and functional lenses:
-
Angular distance (Equation 8): quantifies how much the representation changes from one layer to the next. Layers whose outputs are nearly identical to their neighbors are redundant—they can be removed without consequence. The metric is averaged over 256K tokens from C4 to eliminate sampling noise.
-
Performance drop upon removal (Equation 9): measures how much the model's accuracy degrades when a specific layer is excised. This is the gold-standard importance measure, but is more expensive to compute than angular distance.
The results (Figure 3) confirm the hypothesis with striking clarity. In the Post-LN model, angular distance is lowest (yellow) in the first 3–4 layers and grows larger (purple) with depth—meaning early layers are the most redundant. Performance drop follows the same pattern: removing early layers causes negligible accuracy loss, while removing layers 9–11 is catastrophic. In the Pre-LN model, the pattern is inverted: angular distance decreases with depth, with layers 7–11 showing the highest similarity, and removing deep layers causes almost no performance change.
The gradient norm plot (Figure 3e) provides the mechanistic explanation: Post-LN shows large gradients in deep layers (layer 11 has ~0.6) but vanishing gradients in early layers (layer 0 has ~0.05), while Pre-LN shows the reverse (layer 0 has ~0.4, layer 11 has ~0.1). This directly validates the theoretical analysis from Section 2.1.
The open-weight large-scale comparison (LLaMA2-7B vs. BERT-large) in Figure 2 corroborates these findings at scale: LLaMA2-7B (Pre-LN) shows increasingly redundant deep layers (layers 20–30 highly similar) with negligible pruning impact, while BERT-large (Post-LN) shows the opposite pattern with early layers being the most redundant.
The Gap This Paper Fills
In summary, the paper identifies a clear, previously unaddressed gap: no existing normalization strategy accounts for the depth-dependent complementarity of Pre-LN and Post-LN. Prior work either (a) analyzed the failure modes of each variant theoretically without proposing a hybrid, (b) developed stabilization techniques that apply uniformly across depth, or (c) documented deep-layer ineffectiveness as a pruning opportunity without treating it as a training problem. Mix-LN is the first method to explicitly exploit the fact that Pre-LN and Post-LN have opposite depth-dependent behaviors, combining them at different depths to achieve balanced gradient norms throughout the entire network. The paper positions this not as an incremental normalization tweak but as a principled fix to a training pathology that has been implicitly accepted as architectural fact by the broader community.
3. Technical Approach
3.1 Reader Orientation
Mix-LN is a hybrid layer normalization strategy for Transformer models that applies different normalization placements at different depths — Post-LN in the early layers, Pre-LN in the later layers — to ensure healthy gradient flow throughout the entire network. The system solves the problem that Pre-LN (used in virtually all modern LLMs) causes deep layers to receive such small gradients that they learn near-identity transformations and contribute minimally to model output, while Post-LN preserves deep-layer gradients but causes gradient vanishing in early layers. The "shape" of the solution is a single hyperparameter $\alpha$ that controls what fraction of the earliest layers use Post-LN, with all deeper layers using standard Pre-LN — a remarkably simple intervention that rebalances gradient norms across depth without changing layer architecture, training objectives, or data pipelines.
3.2 Big-Picture Architecture (Diagram in Words)
The Mix-LN system has two conceptual components, though in practice it is a single architectural modification:
-
Layer normalization placement assignment function — Given a model with
$L$total Transformer layers and a hyperparameter$\alpha \in [0,1]$, this component assigns Post-LN to layers 0 through$\lfloor \alpha L \rfloor - 1$(the first$\alpha L$layers, floored) and Pre-LN to layers$\lfloor \alpha L \rfloor$through$L-1$(all remaining layers). The assignment is static and applied at initialization — there is no dynamic switching during training. -
Modified Transformer layer implementations — Each layer's internal normalization placement follows its assignment. Post-LN layers compute
$\text{LN}(x + \text{Attention}(x))$for the attention sub-layer and$\text{LN}(x + \text{FFN}(x))$for the feed-forward sub-layer (normalization after the residual addition). Pre-LN layers compute$x + \text{Attention}(\text{LN}(x))$and$x + \text{FFN}(\text{LN}(x))$(normalization before the sub-layer, then residual addition). The same RMSNorm implementation is used in both cases — only the position of the normalization changes.
Information flows through the network as follows: input tokens → embedding layer → Post-LN Transformer layers for the first $\lfloor \alpha L \rfloor$ blocks → Pre-LN Transformer layers for the remaining blocks → final layer norm (if present) → output head. During backpropagation, gradients flow through both types of layers, but the Post-LN early layers experience the multiplicative Jacobian effect described in Equation 5 while the Pre-LN later layers preserve residual-path gradient flow. The key is that the Post-LN Jacobian attenuation does not have enough layers to accumulate to catastrophic levels (since only $\alpha L$ layers apply it), while the Pre-LN identity-mapping tendency does not have enough layers to fully manifest in the deepest layers since $(1-\alpha)L$ is smaller than $L$.
3.3 Roadmap for the Deep Dive
-
First, the formal definition of Post-LN and Pre-LN, their derivatives (Equations 1–4), and the critical approximation that reveals why they have opposite depth-dependent behaviors. This is the mechanistic foundation for everything that follows.
-
Second, the Jacobian analysis of layer normalization (Equations 5–7), which shows that LN acts as a gradient-downscaler with factor
$1/\sigma$where$\sigma$is the input standard deviation. Understanding this is essential to see how Post-LN causes gradient vanishing and Pre-LN prevents it — and why combining them works. -
Third, the Mix-LN formulation itself — the layer assignment rule, the hyperparameter
$\alpha$, and the intuition for why Post-LN-early + Pre-LN-late is the correct ordering (rather than the reverse). -
Fourth, the design rationale — why this particular hybrid was chosen over alternatives like interleaving, per-layer learnable mixing, or using modified normalization variants (DeepNorm, Sandwich-LN, etc.) in a depth-dependent way.
-
Fifth, the practical implementation considerations — compatibility with RMSNorm, interaction with training hyperparameters (warmup, learning rate), and the scaling behavior of
$\alpha$at different model sizes.
3.4 Detailed, Sentence-Based Technical Breakdown
This is a theoretically-motivated empirical methods paper whose core idea is that the depth-dependent failure modes of Pre-LN and Post-LN are complementary, and a simple layer-wise assignment strategy can combine their strengths while mitigating their weaknesses.
Post-LN and Pre-LN: Formal Definitions and Derivatives
The paper begins by precisely defining the two normalization strategies in terms of how they compose layer normalization with the residual connection. Understanding these equations is essential because their derivatives — computed in Equations 3 and 4 — directly reveal the gradient flow properties that determine layer effectiveness at different depths.
Post-Layer Normalization. In Post-LN, the normalization operation is applied after the residual addition. For a sub-layer function $F$ (either multi-head attention or feed-forward network) and input $x$:
where $x$ is the input to the sub-layer (the output of the previous sub-layer or the embedding), $F(x)$ is the output of the attention or FFN computation applied to $x$, and $\text{LN}(\cdot)$ is layer normalization.
What this computes: the sub-layer computes its transformation $F(x)$, the result is added to the input via the residual connection $x + F(x)$, and then layer normalization is applied to the sum. The output is a normalized version of the residual sum, with mean shifted to (approximately) zero and variance scaled to (approximately) one.
Why this form: placing LN after the residual addition means that the normalization operation sees the combined contribution of the skip connection and the sub-layer. This constrains the magnitude of the residual path output, preventing it from growing unboundedly. However, as the derivative analysis below reveals, this constraint comes at the cost of attenuating gradients in early layers.
Pre-Layer Normalization. In Pre-LN, the normalization is applied before the sub-layer computation, and the residual connection is a pure addition:
where the symbols are as defined above, but now $\text{LN}(x)$ normalizes the input before it enters the sub-layer.
What this computes: the input $x$ is first normalized, then the sub-layer $F$ operates on the normalized representation, and the result is added back to the original (unnormalized) input $x$. The residual path is additive without any subsequent normalization constraint.
Why this form: by normalizing before the sub-layer but not constraining the residual output, the residual stream can grow in magnitude as representations from successive layers accumulate. This stabilizes early-layer gradients (the residual path provides an unattenuated highway) but causes deep layers to contribute diminishing transformations — the additive accumulation means that the relative contribution of each new layer to the residual stream becomes progressively smaller compared to the accumulated sum of all previous layers.
Derivative analysis — Post-LN gradient. Applying the chain rule to Equation 1:
where $I$ is the identity matrix (from the derivative of the residual connection $x$ with respect to itself), $\frac{\partial F(x)}{\partial x}$ is the Jacobian of the sub-layer computation, and $\frac{\partial \text{LN}(x + F(x))}{\partial (x + F(x))}$ is the Jacobian of the layer normalization evaluated at the residual sum — denoted $J_{\text{LN}}(x')$ where $x' = x + F(x)$.
What this computes: the gradient of the Post-LN output with respect to the input $x$ has two multiplicative factors. The right factor $(I + \partial F / \partial x)$ represents the gradient contribution from the residual connection (identity) plus the sub-layer. The left factor $J_{\text{LN}}$ is the LN Jacobian, which scales this combined gradient. Critically, the LN Jacobian multiplies the entire gradient — meaning any attenuation from LN applies to the residual path as well.
Why this matters for depth: in a Post-LN network with $L$ layers, the gradient from the loss at layer $L$ back to layer $0$ must pass through $L$ such LN Jacobians. Each of these Jacobians has a spectral norm that is (as shown below) approximately $1/\sigma_{x'}$, where $\sigma_{x'}$ is the standard deviation of the LN input. Since $\sigma_{x'}$ is typically greater than 1 during training, each LN Jacobian shrinks the gradient magnitude. After $L$ layers, the gradient has been attenuated by a factor of approximately $\prod_{l=1}^{L} 1/\sigma_l$, causing gradient vanishing in early layers of deep Post-LN models.
Derivative analysis — Pre-LN gradient. Applying the chain rule to Equation 2:
where the symbols are as before, but note the structural difference: the identity matrix $I$ appears additively and is not multiplied by any LN Jacobian.
What this computes: the gradient of the Pre-LN output has two additive terms. The first term $I$ is the gradient from the residual connection — it passes the gradient back to the input unchanged. The second term accounts for the sub-layer's contribution, but crucially, it is multiplied by the LN Jacobian $\partial \text{LN}(x) / \partial x$ only once (for that layer's LN), and this product is added to (not multiplied with) the identity term.
Why this matters for depth: the residual path $I$ provides a gradient highway that is completely unaffected by LN attenuation. Even if the sub-layer's gradient contribution $\frac{\partial F}{\partial \text{LN}} \cdot J_{\text{LN}}$ is small, the identity term ensures that gradients can flow back through arbitrarily many Pre-LN layers without multiplicative attenuation. This is why Pre-LN solves the gradient vanishing problem. However, because the residual stream accumulates $x + F_1(\text{LN}_1(x)) + F_2(\text{LN}_2(\cdot)) + \dots$, the output of layer $\ell$ is dominated by the accumulated sum of all previous layers. The derivative of each individual layer's contribution relative to the total becomes increasingly small with depth, meaning deep Pre-LN layers contribute near-identity transformations — they stop learning effectively.
The Critical Approximation: LN Jacobian as a Scalar Down-Scaler
The paper adapts a result from Xiong et al. (2020) and Takase et al. (2023) that provides an interpretable approximation for the layer normalization Jacobian. This approximation is what makes the qualitative difference between Pre-LN and Post-LN mathematically tractable.
The derivation proceeds as follows. Let $x'$ be the input to a layer normalization operation (in Post-LN, $x' = x + F(x)$; in Pre-LN, $x' = x$). The Jacobian of LN with respect to its input is:
where $d$ is the hidden dimension (e.g., 768 for LLaMA-130M), $\|x'\|_2$ is the L2-norm of the input vector, $I$ is the $d \times d$ identity matrix, and $\frac{x' x'^{\top}}{\|x'\|_2^2}$ is the outer product of the normalized input direction with itself — a rank-1 matrix representing the projection onto the input direction.
What this computes: the exact Jacobian of layer normalization with respect to its vector input. It is a $d \times d$ matrix that describes how a small change in each component of $x'$ affects each component of the normalized output. The structure is $\frac{\sqrt{d}}{\|x'\|_2}$ times a matrix that subtracts the projection onto the input direction from the identity — meaning changes parallel to $x'$ are suppressed while changes orthogonal to $x'$ are passed through at scale $\sqrt{d}/\|x'\|_2$.
The paper then follows the assumption from prior work that $x'$ approximately follows a normal distribution with mean zero. Under this assumption, we can rewrite $x' = \sigma_{x'} \cdot z$ where $z$ is standard normal ($\mathcal{N}(0, 1)$ component-wise) and $\sigma_{x'}$ is the standard deviation. Substituting:
Now the critical simplification: since $d \gg 1$ in LLMs (e.g., $d = 4096$ in LLaMA-7B), the term $z z^{\top} / d$ involves dividing the outer product of a $d$-dimensional standard normal vector by $d$. As $d$ grows, each element of this matrix becomes small (its entries are $z_i z_j / d$, which scale as $1/d$). For large $d$, we approximate this term as negligible, yielding:
where $\sigma_{x'}$ is the standard deviation of the LN input.
What this approximation means operationally: the layer normalization Jacobian behaves approximately like a scalar multiplier $1/\sigma_{x'}$ times the identity matrix. Each LN operation simply scales all gradient components by $1/\sigma_{x'}$ — it shrinks gradients when $\sigma_{x'} > 1$ and amplifies them when $\sigma_{x'} < 1$. The paper notes that "in practice, we observe that $\sigma_{x'}$ gradually grows larger than one during training," meaning LN acts as a gradient attenuator during typical training regimes.
Why this approximation matters: it converts the complicated matrix Jacobian into a single interpretable scalar. For Post-LN, the gradient from layer $L$ back to layer $0$ is approximately multiplied by $\prod_{l=1}^{L} 1/\sigma_l$, which is a product of $L$ numbers each less than 1 — this decays exponentially with depth, explaining the vanishing gradient problem. For Pre-LN, the residual connection bypasses this multiplicative chain entirely — the $I$ term in Equation 4 provides an unattenuated path, explaining why Pre-LN avoids gradient vanishing. But the reason Pre-LN fails in deep layers is subtler: as the residual stream accumulates, the relative magnitude of each new layer's contribution compared to the running sum diminishes, so the derivative through the learned transformation $\frac{\partial F(\text{LN}(x))}{\partial x}$ becomes small relative to the identity path — the layer effectively learns to output near-zero increments, making its weights less effective.
Mix-LN: The Hybrid Formulation
The core insight of Mix-LN is that we can exploit the depth-dependent complementarity of these failure modes by applying the two normalization strategies to different depth ranges. The formulation is:
For a Transformer with $L$ total layers (indexed $0, 1, \dots, L-1$):
- Layers
$0, 1, \dots, \lfloor \alpha L \rfloor - 1$use Post-LN (Equation 1) - Layers
$\lfloor \alpha L \rfloor, \lfloor \alpha L \rfloor + 1, \dots, L-1$use Pre-LN (Equation 2)
where $\alpha \in [0, 1]$ is a hyperparameter controlling the fraction of layers that use Post-LN. The floor operation $\lfloor \cdot \rfloor$ and ceiling $\lceil \cdot \rceil$ handle non-integer boundaries — $\lfloor \alpha L \rfloor$ Post-LN layers, $\lceil (1-\alpha)L \rceil$ Pre-LN layers.
What this computes operationally: no dynamic computation — this is a static architectural choice made at model initialization. For each layer index $\ell$, the implementation checks $\ell < \lfloor \alpha L \rfloor$ and branches to either the Post-LN or Pre-LN forward pass and gradient computation. The forward pass, loss computation, and backpropagation proceed as usual, just with different normalization placements at different depths.
Why Post-LN on early layers and Pre-LN on deeper layers (and not the reverse): this ordering is critical and follows directly from the gradient analysis. Post-LN causes gradient attenuation through its multiplicative LN Jacobian factor $1/\sigma$. If Post-LN were applied to later layers, the gradient from the loss would pass through fewer such attenuating factors to reach early layers — but the early layers would still receive gradients through Pre-LN's unattenuated additive path. The damage would be limited but the benefit would also be limited: the deep layers (under Post-LN) would receive large gradients, but early layers wouldn't benefit from Post-LN's strong gradient signal in the forward direction.
The chosen ordering (Post-LN early, Pre-LN late) exploits the fact that the gradient from deep layers back to early layers is what suffers from vanishing in pure Post-LN. By limiting Post-LN to only $\lfloor \alpha L \rfloor$ early layers, the gradient from these early Post-LN layers back to the input passes through only $\lfloor \alpha L \rfloor$ attenuating LN Jacobians — not $L$ of them. For $\alpha = 0.25$ and $L = 12$, this means only 3 layers of attenuation rather than 12, a factor of roughly $(1/\sigma)^3$ instead of $(1/\sigma)^{12}$ in accumulated scaling. The deeper Pre-LN layers then have healthy gradient flow through their residual identity paths, but with a crucial improvement: because the Post-LN early layers have constrained the residual stream magnitude (Post-LN normalizes after each residual addition), the Pre-LN layers operate on inputs whose scale is controlled, preventing the unbounded accumulation that causes Pre-LN layers to contribute diminishing returns.
The gradient norm balancing mechanism. The paper includes a plot (Figure 6a) demonstrating the effect. Under pure Pre-LN (Figure 3e), gradient norms start around 0.4 in layer 0 and decay to about 0.1 by layer 11 — a 4× reduction. Under pure Post-LN, gradient norms start near 0.05 in layer 0 (vanishing) but grow to 0.6 in layer 11 — a 12× increase from bottom to top. Mix-LN with $\alpha = 0.25$ produces a flatter profile: the early layers have higher gradients than pure Pre-LN (they benefit from Post-LN's stronger backward signal), while the deeper layers have higher gradients than pure Pre-LN (the reduced number of Pre-LN layers means less accumulation and more distinct per-layer contributions). The result is "more consistent gradient norms across all layers" that enables "all parts of the network—both shallow and deep layers—to contribute effectively to model training."
Why the last Pre-LN layers still work. The paper acknowledges a subtle point: even with Mix-LN, the very deepest layers use Pre-LN and could theoretically suffer from the same identity-mapping problem. However, the number of Pre-LN layers is reduced from $L$ to $\lceil (1-\alpha)L \rceil$ — for $\alpha = 0.25$ and $L = 12$, this is 9 layers instead of 12. The diminished gradient problem in Pre-LN accumulates with the number of consecutive Pre-LN layers, since the residual stream magnitude grows with each additive step. By shortening the Pre-LN segment, the deepest layer's relative contribution to the residual stream is larger (fewer preceding layers have accumulated), so its learned transformation $F(\text{LN}(x))$ adds meaningful signal rather than being swamped by the running sum.
Hyperparameter α: The Post-LN Ratio
The hyperparameter $\alpha$ controls the transition point between normalization strategies. The paper conducts a sweep over $\alpha \in \{0, 0.167, 0.25, 0.33, 0.417, 0.5, 1.0\}$ using LLaMA-1B to determine the optimal value, with results shown in Table 7.
The sweep results. At $\alpha = 0$ (pure Pre-LN), perplexity is 18.65. As $\alpha$ increases (more early layers use Post-LN), perplexity improves: 18.34 at $\alpha = 0.167$, reaching the minimum of 18.18 at $\alpha = 0.25$. Beyond this optimum, perplexity degrades: 18.41 at $\alpha = 0.33$, 18.55 at $\alpha = 0.417$, 18.86 at $\alpha = 0.5$. At $\alpha = 1.0$ (pure Post-LN), training collapses to perplexity 1434 — consistent with Post-LN's known instability at scale.
What this sweep reveals about the mechanism. The U-shaped relationship between $\alpha$ and perplexity confirms that the benefits of Post-LN (higher gradients in deeper layers) and its costs (gradient attenuation in early layers) must be balanced. Too few Post-LN layers ($\alpha$ small) and the gradient benefit to deep layers is insufficient — the Pre-LN segment is too long, and deep layers still suffer from diminished contributions. Too many Post-LN layers ($\alpha$ large) and the accumulated attenuation $(1/\sigma)^{\alpha L}$ becomes problematic, starving early layers of gradient signal. The optimum at $\alpha = 0.25$ suggests that for the LLaMA architecture at this scale, the "right" amount of Post-LN is exactly one quarter of the total depth.
The paper's α selection for all experiments. Based on this sweep, the authors "choose $\alpha = 0.25$ for all model sizes, although we believe there is potential to further improve the performance of Mix-LN by searching for the optimal $\alpha$ for each individual model." This is a pragmatic choice — a single $\alpha$ across all model sizes from 71M to 1B simplifies the experimental design and already yields consistent gains. However, the scaling-up experiment (7B) reveals that this choice may not be optimal at larger scales.
Scaling behavior: α must shrink for larger models. In Section 4.2, the paper reports that "Mix-LN becomes more sensitive to $\alpha$ when scaling up to 7B training. In general, smaller values of $\alpha$ and a longer warm-up period are required to stabilize the training of Mix-LN at a larger scale." The LLaMA-7B experiment uses $\alpha = 0.0625$ (6.25%, or 2 Post-LN layers out of 32 total) to "ensure stable training." This suggests that the optimal $\alpha$ is not constant with model depth — deeper models may need proportionally fewer Post-LN layers to avoid the accumulating attenuation problem from the Post-LN segment. The authors explicitly flag this as an open problem:
"We conjecture that this instability is relevant to Post-LN, as using Post-LN alone also leads to the same issue. Addressing this challenge presents an interesting research direction."
Why Not Other Hybridization Strategies?
The paper's choice of a hard depth boundary between Post-LN and Pre-LN is deliberate, and the alternatives — some of which appear in the literature — are less suited to the problem.
Why not interleave Post-LN and Pre-LN? One might imagine alternating Post-LN and Pre-LN layers rather than grouping them by depth. This would be suboptimal because the gradient attenuation from Post-LN is cumulative — inserting even a single Post-LN layer early in the network subjects all earlier layers to its Jacobian attenuation on the backward pass. Interleaving would spread this attenuation throughout the depth rather than containing it to a bounded early segment. Worse, a Post-LN layer late in the network would provide large gradients to its immediate predecessors but wouldn't help earlier layers, since the gradient from that late Post-LN would still pass through many intervening Pre-LN identity paths — which provide gradient passage but don't amplify it. The grouped design ensures that Post-LN's attenuating effect is limited to the first $\alpha L$ layers, where the accumulated multiplication $(1/\sigma)^{\alpha L}$ is still manageable.
Why not use a weighted combination within each layer? One could define a "mixed" normalization per layer: $\beta \cdot \text{Post-LN}(x) + (1-\beta) \cdot \text{Pre-LN}(x)$ for some mixing weight $\beta$. This would create a continuous interpolation between the two strategies but would also inherit both failure modes simultaneously in every layer — gradient attenuation from the Post-LN component and diminishing returns from the Pre-LN component would both operate at every depth, defeating the purpose of depth-dependent specialization. The hard boundary ensures that Post-LN's benefits (large deep-layer gradients) are realized in the deep layers (by improving gradient flow to them from the loss) while its costs (gradient vanishing) are confined to a short early segment where the accumulation is tolerable.
Why not use modified Post-LN variants (DeepNorm, Admin) for early layers? Methods like DeepNorm and Admin were designed to stabilize Post-LN at all depths by modifying initialization or scaling. They could, in principle, be used instead of vanilla Post-LN for the early layers in Mix-LN. However, the paper's empirical results (Table 1) show that DeepNorm alone still diverges at 1B scale (perplexity 1410.94 vs. 18.65 for Pre-LN), indicating it does not fully solve the instability. Vanilla Post-LN with the depth-limiting strategy of Mix-LN achieves stable training and better performance than DeepNorm, suggesting that depth-restriction is a more effective stabilization mechanism than parameter-space interventions. The paper does not experiment with DeepNorm-in-Mix-LN combinations, which remains an unexplored possibility.
Why Post-LN on early layers rather than middle layers? A plausible alternative would be to place Post-LN in the middle layers (say, layers 3–6 in a 12-layer model) and Pre-LN on both the earliest and latest layers. The paper does not explore this geometry, but the gradient analysis suggests why early-placement is correct. The gradient from the loss must travel through all later layers to reach the Post-LN segment. If Post-LN is in the middle, the gradient to the earliest (Pre-LN) layers must pass through both the later Pre-LN layers (identity path — no problem) and the middle Post-LN layers (attenuating path — problem). The early layers would thus still suffer from LN Jacobian multiplication from the middle Post-LN layers. Placing Post-LN at the very beginning means the only layers whose gradients are attenuated are the earliest ones, which are exactly the layers that Post-LN was designed to handle — and there are few enough of them to survive the attenuation.
Compatibility with RMSNorm and Modern LLM Architectures
The paper uses RMSNorm (Root Mean Square Layer Normalization; Zhang & Sennrich, 2019) rather than the original LayerNorm (Ba, 2016) in all experiments. RMSNorm simplifies standard LayerNorm by removing the mean-centering step and only scaling by the root-mean-square of the activations:
where $\gamma$ is a learnable scale parameter. This is the normalization used in LLaMA, Mistral, and most modern LLMs. The Jacobian analysis in Section 2.1 applies to RMSNorm as well — the critical property is the $1/\sigma$ scaling factor, which depends only on the input standard deviation, not on mean-centering.
Why RMSNorm matters for Mix-LN's practicality: the fact that Mix-LN works with RMSNorm (rather than requiring a specific LN variant) means it is a drop-in replacement for the normalization in existing architectures. The paper's experiments use "the LLaMA-based architecture with various sizes from 71M to 1B parameters, incorporating RMSNorm (Shazeer, 2020) and SwiGLU activations (Zhang & Sennrich, 2019)" — exactly the standard modern recipe. No architectural changes beyond the normalization placement are required.
Interaction with SwiGLU and other modern components. The paper uses SwiGLU activations in the FFN layers, following standard practice. There is no indication that Mix-LN interacts differently with SwiGLU than Pre-LN or Post-LN would — the normalization placement affects the gradient flow through the entire sub-layer (attention or FFN), regardless of the internal activation functions. The mechanism is purely through the LN Jacobian, which multiplies the gradient of the combined $(x + F(x))$ or $x + F(\text{LN}(x))$.
Training Configuration for Mix-LN Experiments
The paper trains all models with consistent hyperparameters, varying only the normalization strategy to isolate its effect. The key configuration details (from Appendix A, Table 9) are:
Model architectures across scales:
- 71M parameters: hidden size 512, intermediate size 1368, 8 attention heads, 12 layers.
- 130M parameters: hidden size 768, intermediate size 2048, 12 heads, 12 layers.
- 250M parameters: hidden size 1024, intermediate size 2560, 16 heads, 24 layers.
- 1B parameters: hidden size 2048, intermediate size 5461, 32 heads, 24 layers.
- 7B parameters: hidden size 4096, intermediate size 11008, 32 heads, 32 layers.
Training hyperparameters:
- Optimizer: Adam (Kingma, 2014) with cosine annealing learning rate schedule, decaying to 10% of the initial learning rate.
- Learning rate:
$1 \times 10^{-3}$for models with 250M parameters and below;$5 \times 10^{-4}$for the 1B model;$5 \times 10^{-4}$for the 7B model. - Batch size: 512 sequences, with max sequence length 256, yielding 131K tokens per batch.
- Warmup: learning rate warmup applied to the first 10% of training steps.
- Training data: C4 dataset (Raffel et al., 2020), with token amounts varying by model size: 1.1B tokens (71M, 10K steps), 2.2B tokens (130M, 20K steps), 3.9B tokens (250M, 40K steps), 5.0B tokens (1B, 100K steps), 1.7B tokens (7B, 13K steps — truncated due to compute constraints).
- Mix-LN α: 0.25 for 71M through 1B experiments; 0.0625 for 7B experiment.
Why these hyperparameters matter for the Mix-LN comparison: all models of the same size are trained with identical configurations except for the normalization type. This means the perplexity differences in Table 1 (e.g., 21.39 for Mix-LN vs. 21.92 for Pre-LN at 250M) are purely attributable to the normalization strategy, not to differences in optimization, data, or architecture. The consistent gains across scale suggest that the benefit of balanced gradients is not an artifact of a particular learning rate or batch size — it is a robust property of the normalization placement.
Gradient Norm Measurement Protocol
The gradient norm data presented in Figures 3e and 6a are measured at initialization (before any training steps), providing a clean view of how the normalization strategy affects the gradient flow geometry independent of learning dynamics. The paper states:
"Figure 3-(c) shows the gradient norm of each layer for Post-LN and Pre-LN at the beginning of the training."
This is an important methodological detail. At initialization, the weights are random, and the gradient norms reflect purely the architectural properties — how the normalization placement and residual connections route the loss gradient backward through the untrained network. This is precisely the regime where the LN Jacobian approximation (Equation 7) is cleanest to interpret, since the $\sigma_{x'}$ values are determined by the initialization distribution rather than by the complex statistics of partially-trained representations.
What the gradient norm plots show at initialization (Figure 3e):
- Post-LN: gradient norm is approximately 0.05 in layer 0, grows slowly through layers 1–3, accelerates through layers 4–8, and reaches approximately 0.6 in layer 11 — a 12× increase across depth. The early layers are severely gradient-starved.
- Pre-LN: gradient norm is approximately 0.4 in layer 0, stays roughly flat through layer 3, then gradually declines to approximately 0.1 by layer 11 — a 4× decrease across depth. The deep layers receive relatively weak gradients.
- Mix-LN (Figure 6a, for LLaMA-250M with 24 layers): gradient norms are more uniform — neither the extreme attenuation of Post-LN early layers nor the monotonic decay of Pre-LN. The profile shows higher gradients in the early Post-LN layers than pure Post-LN would produce (since fewer layers of attenuation), and higher gradients in the later Pre-LN layers than pure Pre-LN would produce (since fewer layers of accumulation). The paper describes this as "a more uniform distribution of gradient norms across layers."
Angular Distance Computation Protocol
The angular distance metric (Equation 8) is used extensively in Sections 2 and 5 to quantify how similar a layer's output is to its successors — a proxy for layer redundancy. While not part of the Mix-LN method, understanding how this metric is computed is essential to interpreting the paper's evidence for why Mix-LN works.
Setup: the metric measures the angular distance between the input to layer $\ell$ (denoted $x^{\ell}_T$) and the input to layer $\ell + n$ (denoted $x^{\ell + n}_T$) for a single token $T$. The "input to a layer" means the representation that enters that layer's first sub-layer — in a Pre-LN model, this is the residual stream after the previous layer's addition; in a Post-LN model, this is the normalized output of the previous layer.
The computation for a single token:
where $\cdot$ is the dot product, $||\cdot||$ is the L2 norm, and the factor $1/\pi$ scales the result to $[0, 1]$ (since arccos outputs values in $[0, \pi]$).
What this computes: the cosine similarity between the two representation vectors, converted to an angle via arccos, then normalized to $[0, 1]$. When the two vectors point in exactly the same direction, the cosine is 1, arccos is 0, and $d = 0$ — the representations are identical up to scale. When they are orthogonal, the cosine is 0, arccos is $\pi/2$, and $d = 0.5$. When they point in opposite directions, the cosine is -1, arccos is $\pi$, and $d = 1$.
Aggregation across tokens: to eliminate the effect of individual token variability, the angular distance is averaged over 256K tokens from the C4 dataset. This produces a stable estimate of how much the representation typically changes between two layers.
What the metric reveals about layer effectiveness. A small angular distance between layer $\ell$ and layer $\ell + n$ means that $n$ successive Transformer blocks have produced outputs that point in nearly the same direction in representation space. This is strong evidence that those blocks are learning near-identity transformations — their weights are not meaningfully transforming the representation. Such blocks can be removed without substantially affecting the model's output, which is exactly what the performance drop metric (Equation 9) tests: removing a layer with small angular distance to its neighbors should cause minimal accuracy loss.
The angular distance patterns for different LN types (Figure 5). For LLaMA-130M with 12 layers:
- Pre-LN: angular distance decreases with depth. The bottom-left of the heatmap (early layers) is purple (large distance, ~0.35), while the top-right (deep layers) is yellow (small distance, ~0.10–0.15). Neighboring deep layers have very similar representations — they are redundant.
- Post-LN: angular distance increases with depth. The bottom-left (early layers) is yellow (~0.10), while the top-right (deep layers) is purple (~0.40). Early neighboring layers are highly similar.
- Mix-LN: angular distance is more uniformly distributed, with most values in the 0.18–0.35 range. The Post-LN early layers have moderate similarity (not as extreme as pure Post-LN), and the Pre-LN later layers are more distinctive than in pure Pre-LN. The paper notes: "Mix-LN consistently exhibits larger distances (darker color) across layers compared to Pre-LN, except for the final two layers."
This pattern provides the representational evidence that Mix-LN succeeds in its goal: every layer is learning a meaningfully different transformation from its neighbors, maximizing the effective capacity of the network's parameters.
Performance Drop Computation Protocol
The performance drop metric (Equation 9) provides the functional counterpart to angular distance — it directly measures how much a specific layer matters for the model's output by removing it and observing the accuracy change.
where $P_{\text{original}}$ is the model's accuracy without any layer removal, $P(\ell)_{\text{pruned}}$ is the accuracy after removing exactly layer $\ell$ (by bypassing it — connecting the output of layer $\ell-1$ directly to the input of layer $\ell+1$), and $\Delta P(\ell)$ is the accuracy change (negative means performance drops).
What this computes: the marginal contribution of layer $\ell$ to the model's task performance. A $\Delta P(\ell)$ close to 0 means the layer is dispensable — removing it doesn't affect the model's answers. A large negative $\Delta P(\ell)$ means the layer is critical — the model heavily depends on the transformation that layer applies.
Task selection by model type. Different models use different evaluation tasks appropriate to their capabilities:
- LLaMA2-7B (open-weight): MMLU (Hendrycks et al., 2020), a broad knowledge benchmark.
- BERT-large (open-weight): SQuAD v1.1 (Rajpurkar, 2016), a reading comprehension task.
- In-house LLaMA-130M (small-scale): ARC-e (Clark et al., 2018), an easy science question set, evaluated after supervised fine-tuning (since the base model is too weak for zero-shot evaluation).
The performance drop patterns (Figures 3c, 3d, 6b). For LLaMA-130M:
- Pre-LN: removing early layers (0, 1) causes noticeable drops (roughly -3 to -5 points on ARC-e). Removing layers 2–11 causes minimal drops (most within -2 to +2 points), with some layers even showing slight improvements when removed (positive
$\Delta P$). This confirms that Pre-LN concentrates useful computation in the first 1–2 layers. - Post-LN: removing early layers (0–7) causes almost no performance change (values near 0), while removing deep layers (9–11) causes catastrophic drops (up to -15 points for layer 11). Useful computation is concentrated in the deepest layers.
- Mix-LN (Figure 6b): the performance drop profile is more balanced across layers. Early layers show moderate drops, middle layers show consistent contributions, and deep layers show larger drops than in pure Pre-LN — indicating that Mix-LN has successfully redistributed useful computation across the depth, with "deeper layers in Mix-LN models show a larger
$\Delta P$compared to Pre-LN."
The Complete Training and Evaluation Pipeline
Putting all components together, the Mix-LN methodology proceeds as follows:
1. Architecture definition. The model architecture (LLaMA-based with RMSNorm and SwiGLU) is specified with $L$ layers. The first $\lfloor \alpha L \rfloor$ layers are configured with Post-LN placement (LN after each residual addition), and the remaining layers with Pre-LN placement (LN before each sub-layer, residual untouched). The hyperparameter $\alpha$ is set to 0.25 for models up to 1B and to 0.0625 for 7B.
2. Pretraining. The model is trained on the C4 dataset using the Adam optimizer with cosine annealing, a learning rate of $10^{-3}$ (or $5 \times 10^{-4}$ for 1B+), batch size 512, sequence length 256, and 10% warmup steps. The training objective is standard next-token prediction with cross-entropy loss. All models of the same size use identical configurations except normalization type.
3. Perplexity evaluation. After pretraining, perplexity is measured on a held-out validation set from C4. This is the primary metric for comparing normalization strategies (Table 1).
4. Supervised fine-tuning (optional downstream evaluation). To assess how Mix-LN's improved layer quality transfers to downstream tasks, pretrained models are fine-tuned on Commonsense170K (Hu et al., 2023) and evaluated on eight tasks: MMLU, BoolQ, ARC-e, PIQA, HellaSwag, OpenBookQA, and WinoGrande (Table 4).
5. RLHF (optional alignment evaluation). For the 1B models, the InstructGPT RLHF workflow (Ouyang et al., 2022) is applied using the UltraFeedback dataset to optimize a preference reward (Table 5).
6. Analytical measurements. To verify the mechanism, the paper measures (a) gradient norms at initialization for each layer (Figures 3e, 6a), (b) angular distance between layer representations using 256K C4 tokens (Figures 2, 3, 5), and (c) performance drop from single-layer removal (Figures 2, 3, 6b). These measurements confirm that Mix-LN achieves its stated goal of more balanced gradients and more uniformly useful layers.
Design Choices Summary
| Choice | What Was Done | Why |
|---|---|---|
| Boundary position | Hard cutoff at $\lfloor \alpha L \rfloor$, not soft interpolation | Preserves clean gradient properties of each LN type per segment; avoids having both failure modes in every layer |
| Which layers get Post-LN | Earliest layers | Limits accumulated LN Jacobian attenuation to a short segment; Post-LN early = gradient boost to deep layers (which the loss gradient passes through first) without starving early layers |
| Which layers get Pre-LN | Remaining deeper layers | Prevents unbounded residual accumulation in late layers (since early Post-LN layers constrain the residual scale); benefits from healthy backward gradient flow from the identity residual path |
| α = 0.25 for ≤1B models | Determined by sweep on LLaMA-1B (Table 7) | Empirical optimum; represents the sweet spot where Post-LN benefits are realized without dangerous attenuation accumulation |
| α = 0.0625 for 7B | Reduced for stability at scale | Post-LN instability grows with total depth; fewer Post-LN layers needed to avoid training collapse |
| RMSNorm (not LayerNorm) | Used throughout | Compatibility with modern LLaMA architecture; no mean-centering but same $1/\sigma$ scaling property |
| Adam with cosine annealing | Standard optimizer | Isolates normalization effect; no special optimizer needed for Mix-LN |
| Gradient norms measured at initialization | Before any training | Clean architectural signal; not confounded by training dynamics or learned statistics |
| Angular distance averaged over 256K tokens | Statistical stability | Eliminates token-level noise; provides reliable per-layer-pair similarity estimate |
4. Key Insights and Innovations
Innovation 1: Reframing Deep-Layer Inefficacy from an Architectural Fact to a Training Pathology
The paper's most fundamental contribution is not Mix-LN itself but the diagnostic reframing that makes Mix-LN necessary. Prior to this work, the observation that deep layers in LLMs contribute minimally to output was treated as an inherent property to be exploited — the pruning literature (Yin et al., 2023; Gromov et al., 2024; Men et al., 2024) had documented that deep layers could be removed with negligible accuracy loss, and the natural response was to develop better compression algorithms (Siddiqui et al., 2024; Zhong et al., 2024; Sreenivas et al., 2024). The dominant implicit assumption was that this was normal — that the deepest layers of a Transformer are simply less useful, perhaps because lower-level features are more transferable or because the residual stream naturally saturates.
The paper directly challenges this assumption by asking a question the pruning literature never asked: is this property caused by a specific, addressable training deficiency rather than an architectural inevitability? This is a conceptual shift from "how do we compress away useless layers?" to "how do we train so that no layers are useless?" The distinction matters because it redirects the research agenda from mitigation (pruning, distillation) to root-cause analysis (gradient flow, normalization).
The evidence that supports this reframing is the controlled comparison in Figure 3, where two identically-architected LLaMA-130M models — differing only in normalization type — show opposite patterns of layer effectiveness. In the Pre-LN model, deep layers are most redundant (angular distance smallest, performance drop near zero). In the Post-LN model, early layers are most redundant. This inversion cannot be explained by any fixed architectural property (depth itself, residual connections, attention patterns) — it is exclusively a function of the normalization strategy. The paper effectively proves that deep-layer inefficacy is contingent, not necessary.
This reframing has downstream implications beyond this paper. It suggests that other "accepted" properties of deep Transformers — the tendency for representations to converge across layers (Yang et al., 2023), the observation that fine-tuning mostly updates early layers (Li et al., 2024), the success of layer-dropping regularization — may similarly be symptoms of suboptimal gradient flow rather than fundamental inductive biases. The paper opens an entire research direction: identifying which properties of trained LLMs are genuine architectural regularities and which are artifacts of specific training choices that we've implicitly accepted as defaults.
Innovation 2: The Depth-Aware Gradient Flow Diagnosis Through Layer Normalization Analysis
While the mathematical analysis of Pre-LN and Post-LN gradients originates with Xiong et al. (2020), the paper makes a distinct diagnostic contribution by connecting that analysis directly to the observable phenomenon of layer similarity and redundancy — and by using that connection to predict and verify depth-dependent failure modes.
The prior theoretical understanding (Xiong et al., 2020; Takase et al., 2023) established that Post-LN causes gradient vanishing in early layers (the $\prod 1/\sigma$ attenuation chain) while Pre-LN avoids this through the decoupled residual path. But this analysis was used primarily to explain training stability — why Post-LN diverges without careful warmup, why Pre-LN enables deeper architectures. The paper's novel move is to recognize that these gradient properties don't just affect whether training converges; they determine which layers actually learn useful transformations in the converged model.
The innovation is the prediction of inverse redundancy patterns. From the gradient analysis alone, one can predict:
- Pre-LN: deep layers receive small gradients → learn near-identity mappings → high angular similarity to neighbors → low performance impact when removed.
- Post-LN: early layers receive small gradients → learn near-identity mappings → high angular similarity to neighbors → low performance impact when removed.
This prediction is non-obvious — it's not simply "Post-LN is unstable so it's worse." It specifies exactly which layers will be affected and in what direction. The paper then validates this prediction with three independent metrics (gradient norms at initialization in Figure 3e, angular distance in Figures 3a-b and 5, and performance drop in Figures 3c-d and 6b), all of which confirm the predicted patterns.
What makes this a genuine insight rather than an incremental application of Xiong et al.'s math is the explanatory power it provides for otherwise puzzling results in the literature. Gromov et al. (2024) documented that LLaMA-2's deep layers are redundant without explaining why. Lad et al. (2024) found that deep layers in Pythia and GPT-2 are robust to deletion and swapping — again without a mechanistic account. This paper provides a unified explanation: all of these models use Pre-LN, and Pre-LN's gradient decay property predicts that deep layers will be undertrained and therefore dispensable. The diagnosis is falsifiable (a Post-LN model of the same architecture should show the opposite pattern, which the paper demonstrates it does) and actionable (fix the gradient flow, fix the layer quality).
Innovation 3: Depth-Conditioned Hybridization as a General Design Principle
The paper's third contribution is the design principle that normalization strategy should vary with depth, exploiting the fact that Pre-LN and Post-LN have complementary depth-dependent behaviors. This principle is more general than the specific Mix-LN instantiation and represents a departure from how normalization has been treated in the Transformer literature.
Prior work on Transformer normalization can be divided into two categories: (a) analyses of which single strategy works best at all depths (Pre-LN vs. Post-LN, the Xiong et al. debate), and (b) modifications that apply the same strategy to every layer but with tweaked parameters (Admin's extra parameters, DeepNorm's residual scaling, Sandwich-LN's dual normalization). In both cases, the underlying assumption is that the normalization strategy should be depth-uniform — every layer gets the same treatment.
The paper demonstrates that this assumption is suboptimal because the same normalization strategy has different effects depending on layer position. Post-LN at layer 2 (two LN Jacobians between it and the loss, minimal accumulated attenuation) is a very different beast from Post-LN at layer 20 (twenty LN Jacobians, catastrophic attenuation). Pre-LN at layer 20 (twenty layers of residual accumulation pushing it toward identity) is different from Pre-LN at layer 2 (only two layers of accumulation, still plenty of room to contribute). The optimal strategy for a layer depends on its depth, not just on its architecture.
This insight has implications beyond the specific Post-LN/Pre-LN combination. It suggests that other training-time design choices currently applied uniformly — initialization schemes, learning rate multipliers, dropout rates — might similarly benefit from depth-aware scheduling. Prior work on layer-specific learning rates (e.g., LIMA-style fine-tuning with frozen early layers) implicitly acknowledges this, but the principle has not been articulated as a general design philosophy. The paper provides a concrete template: identify a pair of methods with complementary depth-dependent failure modes, apply each where it succeeds, and verify that the transition point can be chosen to avoid both failure regimes simultaneously.
The significance of this principle is supported by the $\alpha$ sweep results in Table 7. The U-shaped relationship between $\alpha$ and perplexity shows that both extremes (pure Pre-LN at $\alpha = 0$ and pure Post-LN at $\alpha = 1$) underperform the hybrid — but more importantly, that the optimal $\alpha$ is not at the boundary. If the benefit came merely from "adding some Post-LN flavor," any non-zero $\alpha$ would be better than $\alpha = 0$, and performance would monotonically increase until Post-LN's instability kicked in. Instead, there's a clear interior optimum ($\alpha = 0.25$), confirming that the interaction between the two normalization regimes — the Post-LN early layers constraining residual scale for the Pre-LN later layers, the Pre-LN later layers protecting the early Post-LN layers from excessive depth — is what matters, not either regime alone.
Innovation 4: Empirical Demonstration That Deep Layers Can Be Redeemed Without Architectural Modification
While Innovations 1–3 are conceptual and diagnostic, Innovation 4 is evidentiary: the paper provides the first controlled demonstration that deep-layer quality can be substantially improved through a normalization change alone, without increasing model size, training data, or architectural complexity. This is significant because it transforms the problem from "deep layers are inherently less useful" (the pruning literature's implicit view) to "current training methods fail to utilize deep layers, but we can fix that."
The evidence for this claim is multi-layered and unusually thorough:
Pretraining quality (Table 1). Mix-LN achieves consistent perplexity improvements over Pre-LN across model sizes (71M through 1B): 33.12 vs. 34.77 (71M), 26.07 vs. 26.78 (130M), 21.39 vs. 21.92 (250M), 18.18 vs. 18.65 (1B). These are not enormous gains — roughly 0.25–1.65 perplexity depending on scale — but they are achieved with zero additional parameters and zero additional computation. Any method that improves perplexity without increasing model cost is effectively increasing parameter efficiency.
Layer quality metrics (Figures 5 and 6b). The angular distance heatmaps show that Mix-LN's layers produce more consistently distinct representations than Pre-LN's layers, with larger distances (darker colors) across most depth positions. The performance drop curves show that Mix-LN's deep layers matter more than Pre-LN's deep layers — removing them causes larger accuracy degradation, meaning they're contributing meaningful transformations rather than near-identity mappings.
Downstream transfer (Tables 4 and 5). Perhaps the most compelling evidence is that Mix-LN's improvements amplify during fine-tuning. At LLaMA-250M, Mix-LN's pre-training advantage over Pre-LN is 0.53 perplexity (21.39 vs. 21.92). After supervised fine-tuning, the gap widens to 4.26% average accuracy across eight tasks (41.34% vs. 36.93%), with dramatic improvements on specific tasks like BoolQ (56.12% vs. 38.35%). The RLHF results (Table 5) show a similar pattern: Mix-LN achieves a final reward of 1.32 vs. Pre-LN's 0.75. This amplification effect is consistent with the paper's core claim: deep layers trained under Mix-LN have developed richer, more transferable features that become valuable when the model is adapted to new tasks. Under Pre-LN, these layers are too similar to provide additional representational capacity — fine-tuning can't extract value from a layer that never learned to represent anything distinct.
Why the amplification effect matters conceptually. It addresses a potential counterargument: if Mix-LN only improved pre-training perplexity by a small amount, perhaps the deep layers are still not that important, and the pruning literature's conclusion (that they can be safely removed) remains practically valid. The fine-tuning results refute this: the deep layers trained under Mix-LN contribute substantially to downstream performance in ways that are not fully captured by pre-training perplexity alone. This suggests that pre-training perplexity — the standard metric for comparing language models — may be too coarse to detect the benefits of better-trained deep layers, and that downstream task evaluation reveals a larger quality gap than perplexity implies.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. All pretraining experiments use the C4 dataset (Raffel et al., 2020), a large-scale cleaned version of Common Crawl web text. The paper uses token amounts that scale with model size: 1.1B tokens for the 71M model (10K training steps), 2.2B tokens for 130M (20K steps), 3.9B tokens for 250M (40K steps), 5.0B tokens for 1B (100K steps), and 1.7B tokens for the 7B model (13K steps — truncated due to compute constraints). Supervised fine-tuning uses Commonsense170K (Hu et al., 2023), a compilation of 170K training examples spanning multiple reasoning and knowledge tasks. Evaluation during fine-tuning spans eight downstream benchmarks: MMLU (Hendrycks et al., 2020), BoolQ, ARC-e (Clark et al., 2018), PIQA, HellaSwag, OpenBookQA, and WinoGrande. RLHF experiments use the UltraFeedback dataset following the InstructGPT workflow (Ouyang et al., 2022). ViT experiments use ImageNet-1K with standard training (120 epochs, ConvNeXt configurations from Liu et al., 2022c).
-
Base model(s). All main experiments use the LLaMA architecture (Touvron et al., 2023) with RMSNorm (Zhang & Sennrich, 2019) and SwiGLU activations (Shazeer, 2020), spanning five scales: 71M, 130M, 250M, 1B, and 7B parameters. Architecture details are provided in Appendix A, Table 9 — for example, the 250M model uses 24 layers with hidden size 1024 and 16 attention heads, while the 1B model uses 24 layers with hidden size 2048 and 32 heads. The authors choose this family because LLaMA represents the dominant modern architecture used in production LLMs, making findings directly transferable. For open-weight comparisons, the paper evaluates LLaMA2-7B (Touvron et al., 2023, Pre-LN) and BERT-large (Devlin, 2018, Post-LN) — though these differ in training data, activation functions, and context length. For vision experiments, the ViT architecture (Dosovitskiy et al., 2021) is used at Tiny and Small scales.
-
Metrics. Perplexity on held-out C4 validation data is the primary pretraining metric — lower is better. For supervised fine-tuning, accuracy (%) is reported on each of the eight downstream tasks, with an unweighted average across tasks. For RLHF, the paper reports final reward (higher is better), following the InstructGPT reward modeling framework. For analytical measurements, three auxiliary metrics are used: gradient norm (L2 norm of the gradient with respect to each layer's parameters, measured at initialization), angular distance
$d(x^{\ell}, x^{\ell+n})$(Equation 8, ranging 0–1, with 0 indicating identical representation direction and 1 indicating opposite directions, averaged over 256K C4 tokens), and performance drop$\Delta P(\ell)$(Equation 9, the change in task accuracy after removing layer$\ell$, where more negative values indicate greater layer importance). -
Baselines. The paper compares against four normalization strategies: (1) Pre-LN (Dai, 2019; Baevski & Auli, 2018) — the dominant approach in GPT, LLaMA, and Mistral, applying LN before each sub-layer with a pure additive residual; (2) Post-LN (Ba, 2016; Vaswani, 2017) — the original Transformer normalization, applying LN after the residual addition; (3) DeepNorm (Wang et al., 2024) — a Post-LN variant that upscales the residual connection before LN to stabilize training at extreme depths; (4) Admin (Liu et al., 2020) — adaptive model initialization that controls residual dependencies in Post-LN through additional parameters. In Table 8, two additional baselines are evaluated: Sandwich-LN (Ding et al., 2021), which normalizes both the input and output of each sub-layer, and Group-LN (Wu & He, 2018; Ma et al., 2024), which applies group-based rather than layer-wise normalization. For the ViT experiments, the baseline is standard Pre-LN.
-
Generation budget / compute accounting. Training compute is measured in training tokens processed, equalized across normalization variants within each model size. For example, all 250M models are trained on exactly 3.9B tokens (40K steps at batch size 512 and sequence length 256 = 131K tokens/batch). The paper does not track wall-clock time or FLOPs explicitly, since Mix-LN introduces no additional operations per forward pass beyond what Pre-LN already computes — it only changes the placement of the normalization operations, not their number or computational cost. This means all normalization variants at a given model size have identical theoretical FLOP counts per training step.
-
Cross-validation / statistical protocol. The paper does not employ k-fold cross-validation — models are trained once from scratch at each scale due to computational cost. Statistical significance is not reported (no confidence intervals or standard errors appear in any table or figure). The
$\alpha$hyperparameter sweep (Table 7) evaluates each$\alpha$value with a single training run on LLaMA-1B, and the chosen value ($\alpha = 0.25$) is held fixed across all other model sizes without per-size tuning. The angular distance metric is averaged over 256K tokens from C4 to reduce sampling noise, providing a statistically stable estimate at the token level, but there is no averaging across multiple training seeds. The open-weight model evaluations (Figure 2) represent single measurements on publicly released checkpoints.
Main Quantitative Results
This section is organized into five axes of investigation: pretraining perplexity across normalization strategies (Table 1), scaling behavior to 7B (Figure 4), supervised fine-tuning performance (Table 4), RLHF reward optimization (Table 5), and vision domain generalization (Table 6).
Pretraining Perplexity: Mix-LN vs. Pre-LN, Post-LN, and DeepNorm Across Scales
The central quantitative claim of the paper is that Mix-LN achieves lower perplexity than Pre-LN at every model scale tested, while Post-LN and DeepNorm diverge at larger scales. Table 1 presents these results:
| Model | Training Tokens | Post-LN | DeepNorm | Pre-LN | Mix-LN |
|---|---|---|---|---|---|
| LLaMA-71M | 1.1B | 35.18 | 34.87 | 34.77 | 33.12 |
| LLaMA-130M | 2.2B | 26.95 | 27.17 | 26.78 | 26.07 |
| LLaMA-250M | 3.9B | 1409.09 | 22.77 | 21.92 | 21.39 |
| LLaMA-1B | 5.0B | 1411.54 | 1410.94 | 18.65 | 18.18 |
Several patterns emerge from this table:
Post-LN collapses at scale. At 71M and 130M parameters, Post-LN achieves perplexities that are competitive with Pre-LN (35.18 vs. 34.77 at 71M; 26.95 vs. 26.78 at 130M). However, at 250M and above, Post-LN diverges catastrophically — perplexity jumps to 1409.09 at 250M and 1411.54 at 1B. These values are essentially random-guess level (a uniform distribution over the C4 vocabulary would produce perplexity on this order), indicating that training has completely failed. This confirms the theoretical prediction from Section 2.1: as depth increases (the 250M model has 24 layers), the accumulated LN Jacobian attenuation $\prod_{l=1}^{24} 1/\sigma_l$ starves early layers of gradient signal, preventing any meaningful learning.
DeepNorm postpones but does not fix the collapse. DeepNorm, which was designed specifically to stabilize Post-LN at depth (Wang et al., 2024), performs well at smaller scales: 34.87 at 71M (competitive with Pre-LN's 34.77) and 22.77 at 250M (slightly worse than Pre-LN's 21.92 but far from collapsed). However, at 1B parameters, DeepNorm also diverges — 1410.94, indistinguishable from Post-LN's 1411.54. The paper notes: "This observation confirms severe training instability of Post-LN, where gradients in early layers vanish, preventing proper model convergence." DeepNorm's residual scaling and initialization modifications can handle moderate depth (24 layers at 250M works), but at larger hidden dimensions and higher total compute, the underlying gradient vanishing problem re-emerges.
Mix-LN gains are consistent and scale with model size. The absolute perplexity reduction from Mix-LN over Pre-LN is: 1.65 at 71M, 0.71 at 130M, 0.53 at 250M, and 0.47 at 1B. While the raw gap shrinks with model size (as perplexity itself decreases — it's harder to improve a better model), the relative improvement remains meaningful. At 71M, Mix-LN reduces perplexity by 4.7% relative to Pre-LN; at 1B, by 2.5%. The paper attributes these gains to "combining the benefits of Pre-LN and Post-LN" — the Post-LN early layers improve gradient flow to middle and deep layers (which are starved under pure Pre-LN), while the Pre-LN later layers prevent the catastrophic vanishing that affects pure Post-LN.
A noteworthy observation from Table 1: DeepNorm actually underperforms Pre-LN at 130M (27.17 vs. 26.78) and underperforms Mix-LN substantially at all scales where it trains successfully. This suggests that DeepNorm's stabilization mechanism (upscaling residuals before LN) trades off some model quality for stability — it prevents collapse but doesn't optimize gradient flow as effectively as the Mix-LN hybrid approach.
Scaling Behavior: Mix-LN at 7B Parameters
The 7B experiment (Figure 4) is truncated — only 13,000 training steps (~1.7B tokens) due to computational constraints — but provides critical evidence about Mix-LN's scalability. The training curve shows evaluation perplexity as a function of updated steps for Mix-LN and Pre-LN, both using the LLaMA-7B architecture (32 layers, hidden size 4096, 32 attention heads).
What Figure 4 shows. The y-axis plots evaluation perplexity (ranging from roughly 27 to 38), and the x-axis plots training steps (from roughly 7,000 to 13,000). Both curves decrease monotonically. Mix-LN (blue curve) is consistently below Pre-LN (orange curve) throughout the plotted range — at step 7,000, Mix-LN is at approximately 28.5 vs. Pre-LN at approximately 29.0; by step 13,000, Mix-LN reaches approximately 27.8 vs. Pre-LN at approximately 28.2. The gap appears to be widening slightly — roughly 0.5 perplexity at 7,000 steps and roughly 0.4–0.5 at 13,000 steps.
Why this matters despite the truncated training. Training a 7B model to convergence is computationally prohibitive for academic groups. The fact that Mix-LN maintains a consistent advantage over Pre-LN throughout 13,000 steps — without any sign of the advantage diminishing — provides suggestive evidence that the benefits scale to models an order of magnitude larger than the 1B experiments. If the gap were an artifact of small-scale training dynamics (e.g., Mix-LN converges faster but Pre-LN catches up with more steps), we would expect the curves to converge. The persistent separation is consistent with a genuine quality improvement rather than a mere speedup.
The α sensitivity caveat. The 7B experiment uses $\alpha = 0.0625$ (6.25%, corresponding to 2 Post-LN layers out of 32 total), rather than the $\alpha = 0.25$ used for smaller models. The paper states that "Mix-LN becomes more sensitive to α when scaling up to 7B training. In general, smaller values of α and a longer warm-up period are required to stabilize the training of Mix-LN at a larger scale." This is an important qualification: the optimal $\alpha$ is not constant with model depth. At 32 layers and $\alpha = 0.25$, there would be 8 consecutive Post-LN layers — apparently enough accumulated LN Jacobian attenuation to cause training difficulties at this scale. Reducing to 2 Post-LN layers limits the attenuation chain to $(1/\sigma)^2$, which is manageable even at large hidden dimensions. The paper treats this as a limitation (Section 6) rather than a fundamental problem, but it means the $\alpha = 0.25$ rule determined from the 1B sweep (Table 7) does not generalize without modification.
Supervised Fine-Tuning Performance
Table 4 reports accuracy on eight downstream tasks after fine-tuning LLaMA-250M and LLaMA-1B models on Commonsense170K. The headline result: Mix-LN models substantially outperform Pre-LN models on downstream tasks, with gains that are larger than the pre-training perplexity differences would suggest.
LLaMA-250M results. Mix-LN achieves an average accuracy of 41.34% across the eight tasks, compared to 36.93% for Pre-LN — a gain of +4.41 percentage points absolute, or +11.9% relative. The gains are not uniform across tasks:
- BoolQ shows the most dramatic improvement: 56.12% for Mix-LN vs. 38.35% for Pre-LN — a +17.77 point gain. BoolQ is a binary yes/no question answering task requiring factual reasoning; the large gain suggests that Mix-LN's deeper layers are learning representations that are particularly useful for this type of inference.
- PIQA improves by +2.79 points (66.34% vs. 63.55%) — physical commonsense reasoning.
- ARC-e improves by +1.53 points (41.68% vs. 40.15%) — easy science questions.
- HellaSwag improves by +3.82 points (30.16% vs. 26.34%) — commonsense NLI.
- MMLU improves by +1.60 points (26.53% vs. 24.93%) — broad knowledge.
- OpenBookQA improves by +1.80 points (18.00% vs. 16.20%) — science understanding.
- WinoGrande improves by +1.55 points (50.56% vs. 49.01%) — pronoun resolution.
Post-LN and DeepNorm baselines are substantially worse: Post-LN achieves only 32.54% average (vs. Mix-LN's 41.34%), and DeepNorm achieves 35.63%. The poor Post-LN fine-tuning performance is expected given its near-random pretraining perplexity (1409.09, Table 1) — the model never learned useful representations in the first place.
LLaMA-1B results. Mix-LN achieves 44.66% average accuracy vs. 43.01% for Pre-LN — a gain of +1.65 points absolute, or +3.8% relative. The gains are smaller than at 250M, consistent with the smaller pre-training perplexity gap (0.47 vs. 0.53). However, the pattern of task-specific improvements differs:
- WinoGrande shows the largest gain: 55.93% vs. 50.51% (+5.42 points). This is notable because at 250M, WinoGrande showed only a +1.55 point improvement — the benefit shifts across tasks as model scale changes.
- ARC-e improves by +2.41 points (48.11% vs. 45.70%).
- MMLU improves by +1.45 points (27.99% vs. 26.54%).
- BoolQ shows essentially no gain: 61.93% vs. 62.20% (-0.27 points). This is the only task where Mix-LN underperforms Pre-LN at 1B.
- HellaSwag, PIQA, and OpenBookQA show modest gains of 0.39, 0.71, and 1.40 points respectively.
Post-LN and DeepNorm at 1B both collapse to near-random accuracy (31.96% and 32.67% respectively), consistent with their pretraining divergence.
The amplification effect. A key finding in Table 4 is that the downstream gains are proportionally larger than the pre-training perplexity improvements. At 250M, Mix-LN improves pre-training perplexity by 0.53 (21.39 vs. 21.92, a 2.4% relative improvement) but improves fine-tuning accuracy by 11.9% relative. The paper interprets this as evidence that "the superior middle and deeper layers produced by Mix-LN are better equipped to learn during supervised fine-tuning" because "these layers capture more diverse and rich features compared to those trained with Pre-LN." In other words, the pre-training perplexity metric partially masks the true quality difference — it's a next-token prediction metric that can be achieved through many representational strategies, some of which leave deep layers interchangeable while others (Mix-LN) make them distinct and transferable.
RLHF Reward Optimization
Table 5 presents RLHF results for LLaMA-1B models. Mix-LN achieves a final reward of 1.32 compared to Pre-LN's 0.75 — a +0.57 advantage, or +76% relative improvement. The paper frames this result as demonstrating that "the benefits of Mix-LN can be seamlessly transferred to RLHF," but the magnitude of the gain warrants scrutiny.
What the reward metric captures. Following the InstructGPT workflow (Ouyang et al., 2022), a reward model is trained on human preference data from UltraFeedback, and the language model is optimized via PPO to maximize this learned reward. The final reward is the average score assigned by the reward model to the policy's outputs. A higher reward means the model's responses are judged by the (learned) reward model to be more aligned with human preferences.
Interpreting the magnitude. A +0.57 reward improvement is substantial. For context, typical RLHF training runs see reward improvements on the order of 1–3 points from the initial policy to the converged policy (Ouyang et al., 2022). A +0.57 gap between two architectures at the end of RLHF training suggests that Mix-LN's representations are significantly better suited to learning the preference optimization task.
Caveat: single training run. As with all experiments in this paper, the RLHF result is from a single training run per method. Without error bars or replication across seeds, it is impossible to assess whether the +0.57 gap is statistically reliable or within the range of training noise. RLHF training is known to be high-variance — reward model quality, PPO hyperparameters, and random seed can all substantially affect final reward. The paper does not report any details about the RLHF training configuration (number of PPO steps, KL penalty coefficient, reward model architecture), making it difficult to assess the robustness of this result.
Vision Domain Generalization
Table 6 evaluates Mix-LN on non-language models by replacing Pre-LN in ViT architectures with Mix-LN ($\alpha = 0.25$). The results on ImageNet-1K (top-1 accuracy, 120 epochs of training):
| Model | Pre-LN | Mix-LN |
|---|---|---|
| ViT-Tiny | 67.30% | 67.34% |
| ViT-Small | 75.99% | 76.40% |
The gains are modest — +0.04% for ViT-Tiny and +0.41% for ViT-Small — but the paper notes that "performance gains are more pronounced in larger models (ViT-Small) compared to smaller ones (ViT-Tiny)." This pattern is consistent with the language model results, where Mix-LN's advantage grows (in absolute terms) with model scale — the 71M LM gained 1.65 perplexity while the 1B gained 0.47 (but perplexity itself is lower, making 0.47 a meaningful fraction).
Significance of the vision result. The key takeaway is not the absolute accuracy improvement but the domain generality — Mix-LN's benefits are not specific to autoregressive language modeling, causal attention, or text data. The ViT architecture uses bidirectional self-attention, image patch inputs, and a classification head rather than next-token prediction. The fact that Mix-LN still provides improvements (albeit small) suggests that the underlying mechanism — balancing gradient norms across depth — is a general property of deep Transformer training, not a quirk of language model pretraining.
Why are the vision gains smaller? The paper does not explore this, but several factors may contribute. ViT-Tiny and ViT-Small are shallow by LLM standards (12 layers for Tiny, depth not specified for Small but typically 12–16 layers). With fewer layers, the gradient imbalance problem is less severe — Pre-LN's deep-layer decay doesn't have as many layers to accumulate, so the room for Mix-LN to improve is smaller. Additionally, ImageNet-1K is a relatively small dataset by modern standards (1.28M images), and supervised training from random initialization with strong data augmentation may saturate performance before gradient flow differences fully manifest. The paper's language experiments use much more data (billions of tokens), which may be necessary for the benefits of better deep-layer training to materialize.
Ablation Studies and Robustness Checks
Post-LN ratio $\alpha$ sweep (Table 7): The paper sweeps $\alpha \in \{0, 0.167, 0.25, 0.33, 0.417, 0.5, 1.0\}$ on LLaMA-1B, finding a U-shaped relationship between $\alpha$ and perplexity. Perplexity improves from 18.65 ($\alpha = 0$, pure Pre-LN) to 18.34 ($\alpha = 0.167$), reaches a minimum of 18.18 at $\alpha = 0.25$, then degrades to 18.41 ($\alpha = 0.33$), 18.55 ($\alpha = 0.417$), 18.86 ($\alpha = 0.5$), and collapses to 1434 ($\alpha = 1.0$, pure Post-LN). The U-shape confirms that the hybrid strategy genuinely outperforms either extreme — the benefit is not from merely "adding some Post-LN" but from the specific balance where Post-LN early layers provide gradient benefits to deeper layers without accumulating enough LN Jacobian attenuation to cause vanishing. The paper uses $\alpha = 0.25$ for all subsequent experiments up to 1B scale.
Scaling $\alpha$ at 7B: For LLaMA-7B, the paper reduces $\alpha$ to 0.0625 (2 out of 32 layers) for training stability. This is an implicit ablation revealing that the optimal $\alpha$ is not invariant with model depth — the accumulation factor $(1/\sigma)^{\alpha L}$ means that deeper models require proportionally fewer Post-LN layers to stay below the vanishing threshold. The paper does not sweep $\alpha$ at 7B (due to computational cost), leaving the optimal value at this scale unknown.
Angular distance across layers for Mix-LN vs. Pre-LN vs. Post-LN (Figure 5): For LLaMA-130M, the angular distance heatmaps show that Mix-LN produces consistently larger distances (darker purple colors, indicating more distinct representations) across most depth positions compared to Pre-LN. Specifically, Pre-LN's angular distance values cluster in the 0.10–0.25 range with deep layers showing very low distance (0.10–0.15), while Mix-LN's values span 0.18–0.35 with fewer extreme minima. Post-LN shows the opposite pattern to Pre-LN — very low angular distance (yellow) in the early layers (0.10–0.15 for the first 3–4 layers) and larger distance (purple) in deep layers (0.30–0.40). Mix-LN achieves the best of both: early layers are more distinct than in pure Post-LN, and deep layers are more distinct than in pure Pre-LN. The paper notes that Mix-LN's final two layers show similar distances to Pre-LN's final layers — even Mix-LN cannot completely prevent representational convergence in the very deepest Pre-LN layers, though the effect is substantially mitigated.
Gradient norm comparison at initialization (Figure 6a): For LLaMA-250M (24 layers), Mix-LN produces gradient norms that are more uniform across depth than either Pre-LN or Post-LN. The paper's plot shows that Pre-LN gradients decay monotonically with depth, Post-LN gradients increase monotonically with depth (with severe vanishing in the first few layers), and Mix-LN gradients show a flatter profile — higher in the middle layers than Pre-LN, lower in the deep layers than Post-LN, and without the extreme attenuation of Post-LN's early layers. The "uniform gradient norm" property is the paper's mechanistic explanation for why Mix-LN works: when all layers receive comparable gradient magnitudes, all layers learn meaningful transformations rather than some layers dominating while others become near-identity.
Performance drop from layer removal for Mix-LN vs. Pre-LN vs. Post-LN (Figure 6b): For LLaMA-130M (12 layers), removing individual layers from the Mix-LN model causes more uniform performance impact across depth compared to Pre-LN. In Pre-LN, layers 2–11 show performance drops near zero (some even positive — removing the layer improves accuracy slightly), while only layers 0–1 show meaningful negative drops. In Mix-LN, the performance drop curve is flatter: early layers still show drops, middle layers show moderate drops (roughly -3 to -5 ARC-e points), and deep layers show larger drops than in Pre-LN (roughly -5 to -8 for layers 8–10, compared to near-zero in Pre-LN). The paper interprets this as evidence that "Mix-LN achieves a more significant contribution from deeper layers" — layers that were essentially dead weight under Pre-LN become active contributors under Mix-LN.
Compatibility with Scaled Initialization and Scaled Embed (Table 10, Appendix B): The paper tests whether Mix-LN integrates with two stabilization techniques from prior work: Scaled Initialization (Nguyen & Salazar, 2019; Scao et al., 2022), which initializes weight matrices with smaller variance to control activation magnitudes, and Scaled Embed (Takase et al., 2023), which scales up embeddings to stabilize LN gradients. On LLaMA-130M:
| Normalization | Scaled Init | Scaled Embed | Perplexity |
|---|---|---|---|
| Pre-LN | ✗ | ✗ | 32.18 |
| Mix-LN | ✗ | ✗ | 29.95 |
| Pre-LN | ✓ | ✗ | 30.63 |
| Mix-LN | ✓ | ✗ | 29.77 |
| Pre-LN | ✓ | ✓ | 31.28 |
| Mix-LN | ✓ | ✓ | 31.19 |
Several findings: (1) Mix-LN consistently outperforms Pre-LN across all configurations, confirming robustness. (2) Scaled Initialization improves both Pre-LN (32.18 → 30.63) and Mix-LN (29.95 → 29.77), suggesting it provides complementary benefits. (3) Adding Scaled Embed on top of Scaled Initialization degrades both methods: Pre-LN goes from 30.63 to 31.28, Mix-LN from 29.77 to 31.19. The gap between Mix-LN and Pre-LN stays roughly constant (about 0.09–0.86 perplexity in Mix-LN's favor) regardless of these auxiliary techniques, suggesting that Mix-LN's benefits are additive with (and not superseded by) existing stabilization methods.
Comparison against other normalization methods (Table 8): On LLaMA-250M, the paper evaluates Mix-LN alongside Admin (Liu et al., 2020), Group-LN (Wu & He, 2018; Ma et al., 2024), and Sandwich-LN (Ding et al., 2021). Results: Pre-LN achieves 23.39 perplexity, Admin performs worse at 24.82, Group-LN marginally improves to 23.10, Sandwich-LN achieves 23.26, and Mix-LN achieves 22.33 — the lowest by a clear margin. This ablation establishes that Mix-LN is not simply benefiting from being "not Pre-LN" — other normalization variants that also modify the standard recipe do not achieve comparable gains. Admin's underperformance relative to Pre-LN is notable given that it was designed specifically to stabilize Post-LN training; its failure here (in a Pre-LN model at moderate scale) suggests its mechanisms (additional parameters to control residual dependencies) may be poorly suited to the LLaMA architecture.
Critical Assessment
Does the evidence support the claim that Pre-LN is the root cause of deep-layer inefficacy?
The controlled comparison (Figure 3) provides strong supporting evidence but has an important scope limitation. The paper demonstrates that two identically-architected LLaMA-130M models differing only in normalization type show inverse patterns of layer effectiveness — Pre-LN concentrates useful computation in early layers, Post-LN concentrates it in deep layers. This is compelling evidence that normalization choice causes the pattern, not merely correlates with it. The use of three independent metrics (gradient norms, angular distance, performance drop) converging on the same conclusion strengthens the case.
However, the demonstration is at a single small scale (130M parameters, 12 layers) and on a single architecture family (LLaMA with RMSNorm and SwiGLU). Modern frontier LLMs have 70B+ parameters and 80+ layers. The paper's open-weight analysis (LLaMA2-7B vs. BERT-large, Figure 2) attempts to bridge this gap, but these models differ in far more than normalization — architecture (decoder-only vs. encoder-only), training objective (causal LM vs. masked LM), training data, activation functions, and context length are all confounded. The BERT-large result showing early-layer redundancy under Post-LN is consistent with the hypothesis but does not isolate normalization as the cause the way the controlled 130M experiment does. A more rigorous test would involve training a modern decoder-only LLM at >1B scale with Post-LN (plus stabilization techniques like warmup and careful initialization to prevent divergence) and demonstrating the inverted redundancy pattern. This experiment is not attempted, likely because Post-LN diverges at scale (as shown in Table 1).
Additionally, the claim that Pre-LN is the root cause — implying it is the primary or sole factor — is stronger than what the evidence strictly supports. The paper demonstrates that normalization differences can cause the pattern, but does not rule out contributions from other factors that correlate with normalization choice in practice (initialization schemes, learning rate schedules, model depth, width). The gradient norm analysis (Section 2.1) provides a plausible mechanism, but the causal chain "Pre-LN → small deep-layer gradients → near-identity learning → redundancy" has only been demonstrated end-to-end at 130M.
Does the evidence support the claim that Mix-LN consistently outperforms Pre-LN?
The pre-training perplexity results (Table 1) consistently support this claim across the tested scales (71M–7B). At every model size where both methods are evaluated, Mix-LN achieves lower perplexity — with margins ranging from 0.47 (1B) to 1.65 (71M). The 7B training curve (Figure 4) shows a persistent gap through 13,000 steps, suggesting the benefit survives scaling. However, three limitations temper this evidence:
First, the 7B experiment is truncated. 13,000 steps on 1.7B tokens is far from convergence for a 7B model. It is possible (though seems unlikely given the stable gap) that Pre-LN would catch up with additional training. The Chinchilla scaling laws (Hoffmann et al., 2022) suggest that a 7B model should be trained on ~140B tokens for compute-optimality — the paper's 1.7B tokens is roughly 1% of that. Whether Mix-LN's advantage persists or diminishes at Chinchilla-optimal token counts is unknown.
Second, the $\alpha$ hyperparameter requires per-scale tuning for stability at 7B. The paper's $\alpha = 0.25$ rule—determined from a 1B sweep—does not transfer to 7B (which requires $\alpha = 0.0625$). This means deploying Mix-LN at a new scale requires either a costly $\alpha$ sweep or a yet-undiscovered scaling law for the optimal $\alpha$ as a function of depth and width. The paper acknowledges this as a limitation (Section 6) but does not characterize how $\alpha$ should be chosen for arbitrary model sizes.
Third, only one architecture family (LLaMA) is tested for language. While the ViT experiments (Table 6) demonstrate domain generality to vision, the sample size is small (two model sizes) and the gains are tiny at the smaller scale (+0.04% for ViT-Tiny). It is unclear whether Mix-LN would benefit encoder-decoder architectures (T5), mixture-of-experts models, or models with different attention mechanisms.
Does the evidence support the claim that Mix-LN improves supervised fine-tuning and RLHF?
The fine-tuning results (Table 4) show substantial gains, but the evidence is strongest at smaller scales. At 250M, Mix-LN's +4.41% average accuracy gain is large and consistent across seven of eight tasks. At 1B, the gain shrinks to +1.65% and is less uniform — BoolQ actually slightly favors Pre-LN. The paper attributes the shrinking gap to a smaller pretraining perplexity difference at 1B (0.47 vs. 0.53), which is a plausible explanation — if the pretrained representations are more similar between Mix-LN and Pre-LN at larger scales, the downstream transfer gap should narrow.
However, the paper only evaluates fine-tuning on Commonsense170K, a relatively small dataset of 170K examples. It is unclear whether the gains would persist with larger fine-tuning datasets (e.g., FLAN, 1M+ examples) or with parameter-efficient fine-tuning methods (LoRA, adapters) rather than full fine-tuning. If the benefit comes primarily from deep layers being more adaptable rather than from better pretrained features, then data-limited fine-tuning (where the model must rely on pretrained knowledge rather than learning new patterns) might show larger Mix-LN advantages than data-rich fine-tuning. The paper does not explore this dimension.
The RLHF result (Table 5) is the weakest evidence in the paper. A single training run per method, with no reported hyperparameters, no error bars, and a metric (learned reward) that is itself a function of a separately trained reward model with its own variance, provides minimal statistical confidence. The +0.57 reward gap is numerically large, but without knowing the standard deviation of the RLHF training process under different random seeds, it is impossible to assess significance. RLHF is notoriously unstable — PPO training with a learned reward model can show reward improvements that do not correlate with human preference ratings (reward hacking). The paper does not report whether the Mix-LN model's higher reward translates to actually preferred outputs (e.g., via human evaluation or an independent evaluation benchmark like AlpacaEval). This is a significant omission for a claim about alignment quality.
Does the evidence support the mechanism (balanced gradients → better layer quality)?
The gradient norm measurements (Figures 3e, 6a) are at initialization — before any training occurs. These plots demonstrate that Mix-LN's architecture produces more uniform gradient norms at the start of training, but they do not show that this uniformity persists throughout training or that it is the cause of improved layer quality in the converged model. The gradient dynamics during training are complex — gradient norms evolve as the model learns, and the distribution of LN input variances $\sigma_{x'}$ (which control the attenuation factor $1/\sigma$) changes as representations develop structure. It is possible that the initialization-time gradient balance is a good proxy for training-time behavior, but the paper provides no evidence for this (no gradient norm measurements at intermediate or final training steps).
The angular distance and performance drop measurements are on the final trained model and do provide evidence that Mix-LN achieves its stated goal of more uniformly useful layers. Figure 5 shows that angular distances are more balanced across depth under Mix-LN. Figure 6b shows that performance drops from layer removal are more uniform. These results are consistent with the claimed mechanism (balanced gradients → balanced learning → balanced importance), but they are correlational, not causal. The paper does not, for example, show that artificially equalizing gradient norms in a Pre-LN model (e.g., through layer-specific learning rate multipliers) produces the same effect as Mix-LN — which would be a stronger test of the gradient balance hypothesis.
Experiments that would have strengthened the paper
Training a Post-LN model at >1B scale with stabilization. The paper's central diagnostic claims rely on the Post-LN baseline to show the inverted redundancy pattern, but Post-LN diverges at scales where Pre-LN's deep-layer problem is most practically relevant (250M+). Using DeepNorm or Admin to stabilize Post-LN at 1B+ and demonstrating that deep layers become more useful under stabilized Post-LN would directly validate the gradient-flow explanation at a scale where it matters. The paper shows DeepNorm diverging at 1B (Table 1), but does not explore whether other stabilization techniques (longer warmup, smaller learning rates, gradient clipping) could enable a successful Post-LN run.
Measuring gradient norms throughout training, not just at initialization. The paper's mechanistic story is dynamic — LN input variances grow during training ($\sigma_{x'}$ "gradually grows larger than one"), and this accumulation causes the depth-dependent gradient effects. Tracking how gradient norms evolve over the course of training for Pre-LN, Post-LN, and Mix-LN would reveal whether the initialization-time patterns (Figure 6a) persist or whether training dynamics amplify/attenuate the differences between methods.
An $\alpha$ sweep at multiple scales. The paper determines $\alpha = 0.25$ on LLaMA-1B and uses it for all smaller models, but the 7B experiment reveals this value doesn't transfer. Sweeping $\alpha$ at 130M, 250M, and 7B (even coarsely) would reveal whether the optimal $\alpha$ follows a predictable trend with depth or width — information that would be practically necessary for anyone wanting to use Mix-LN at an arbitrary scale.
Multiple training seeds for key comparisons. All reported results are from single training runs. Given known variance in large-scale training (different random seeds can produce perplexity differences of 0.1–0.5 even with identical hyperparameters), the reported gaps between Mix-LN and Pre-LN at 1B (0.47 perplexity) could potentially be within the noise range. Multi-seed results would establish whether the improvements are statistically reliable.
Fine-tuning with more diverse datasets and methods. Evaluating Mix-LN's fine-tuning benefits only on Commonsense170K leaves open whether the advantage is specific to small-scale, knowledge-intensive fine-tuning. Testing on larger instruction-tuning datasets, on code generation tasks, or with parameter-efficient methods would clarify the scope of the downstream benefit.
Human evaluation or benchmark-based evaluation for RLHF. The paper reports only the learned reward score, which is susceptible to reward model over-optimization. Evaluating the RLHF-tuned models on a standard benchmark like AlpacaEval or MT-Bench, or conducting human preference judgments, would establish whether the higher reward translates to genuinely better outputs.
Where claims hold conditionally
-
Mix-LN outperforms Pre-LN on pretraining perplexity: holds across all tested scales (71M–7B), but the 7B result is from truncated training and uses a different
$\alpha$than smaller models. Whether the advantage persists at 7B with full Chinchilla-optimal token counts is unknown. -
Mix-LN improves fine-tuning performance: holds strongly at 250M (+4.41% average), more modestly at 1B (+1.65% average), and has only been tested on a single fine-tuning dataset (Commonsense170K) with full fine-tuning. Generalization to other fine-tuning regimes is untested.
-
Mix-LN improves RLHF: reported for a single 1B run with no statistical reliability assessment and no independent evaluation of output quality beyond the learned reward score.
-
Mix-LN improves vision models: demonstrated for ViT-Tiny and ViT-Small on ImageNet-1K, with gains that are small in absolute terms (+0.04% and +0.41%) and increasing with model size. Whether the benefit extends to larger ViTs or to more modern vision architectures (Swin, ConvNeXt) is unknown.
-
Mix-LN balances gradient norms: demonstrated at initialization for specific model sizes. Whether this balance persists throughout training and is the causal mechanism for improved layer quality is supported by correlation (angular distance, performance drop) but not by direct causal intervention.
6. Limitations and Trade-offs
The α Hyperparameter Does Not Transfer Across Model Scales
The assumption or constraint. Mix-LN introduces a single hyperparameter α that controls the fraction of early layers using Post-LN. The paper determines the optimal value via a sweep on LLaMA-1B (Table 7), finding α = 0.25 minimizes perplexity at that scale, and applies this same value to all smaller models (71M through 250M). However, the scaling-up experiment reveals that this value does not generalize. Section 4.2 states:
"Mix-LN becomes more sensitive to α when scaling up to 7B training. In general, smaller values of α and a longer warm-up period are required to stabilize the training of Mix-LN at a larger scale."
The LLaMA-7B experiment uses α = 0.0625 (6.25%, or 2 Post-LN layers out of 32) rather than α = 0.25, which would have placed 8 Post-LN layers. The paper provides no guidance for how to select α at arbitrary model depths or widths beyond this qualitative observation.
The consequence. A practitioner wanting to deploy Mix-LN at a new model scale faces a significant practical barrier: they must either conduct an expensive α sweep at the target scale (training multiple full-size models to convergence on billions of tokens) or guess a value and risk training instability. Worse, the "safe" value appears to shrink with depth — α = 0.25 works at 24 layers (1B) but is unstable at 32 layers (7B) — and the paper provides no scaling law relating optimal α to model depth, width, training tokens, or any other architecture parameter. For models at the 13B, 30B, or 70B scale that are of practical interest, the required α might be even smaller (perhaps 1–2%), and the cost of determining this via sweep is prohibitive. A guess that is too large risks training collapse; a guess that is too small leaves performance on the table. This uncertainty substantially increases the risk of adopting Mix-LN at production scale relative to the more predictable Pre-LN baseline.
What evidence exists in the paper. The α sweep is conducted only at 1B (Table 7), producing the U-shaped curve with optimum at 0.25. The 7B result (Figure 4, Section 4.2) uses a different α (0.0625) explicitly to avoid instability, but no sweep is performed at 7B to confirm this is optimal. The paper acknowledges the sensitivity in Section 4.2 and reiterates the limitation in Section 6:
"We conjecture that this instability is relevant to Post-LN, as using Post-LN alone also leads to the same issue. An interesting research direction is to explore how to mitigate the ineffectiveness of later layers without relying on Post-LN or to address the instability associated with Post-LN."
Mitigation status. Not addressed. The paper treats the α-scaling problem as an open question and a direction for future work. It does not propose a heuristic formula, a theoretical estimate based on depth and width, or an adaptive scheme that would determine α without full-scale training runs. The practitioner is left with the advice to use α = 0.25 for models up to ~24 layers and to reduce α for deeper models, with no quantitative guidance.
The Headline Gains Are Measured Without Accounting for Experimentation Overhead
The assumption or constraint. The paper's primary efficiency claim is that Mix-LN improves perplexity and downstream accuracy "without increasing model size" — it is a pure architectural change that adds no parameters or FLOPs to the forward or backward pass. This is true at deployment time: a Mix-LN model has identical inference cost to a Pre-LN model of the same architecture. However, the experimentation cost to adopt Mix-LN — the work required to find a working α, verify training stability, and tune complementary hyperparameters like warmup length — is substantial and is not amortized in any of the reported comparisons. The paper's own process involved: a 6-value α sweep at 1B (each training run consuming 5 billion tokens, Table 7); discovering instability at 7B with α = 0.25 and reducing to α = 0.0625 through undisclosed trial and error; and noting that "a longer warm-up period" is required at 7B (Section 4.2) without specifying how this was determined.
The consequence. The reported perplexity improvements — e.g., 18.18 vs. 18.65 at 1B, or the 7B curve in Figure 4 — represent the best configuration after hyperparameter optimization that would not be available to a team using Mix-LN for the first time on a new architecture. If a practitioner must spend 5–10 training runs to tune α and warmup duration, the total compute to reach a final Mix-LN model is (1 + k) × the compute of training a Pre-LN baseline, where k is the number of failed or suboptimal runs. The paper implicitly assumes k = 0 — that the practitioner already knows the correct α and supplementary hyperparameters for their specific model scale and training configuration. For a method whose primary advantage is improved parameter efficiency, this unaccounted experimentation cost partially undermines the efficiency argument.
What evidence exists in the paper. The paper is transparent about the hyperparameters used (Table 9 in Appendix A specifies α = 0.25 for models up to 1B and α = 0.0625 for 7B), but does not report how many failed or suboptimal runs occurred before arriving at these values. The total compute spent on α sweeps (Table 7), stability debugging at 7B, and compatibility testing (Table 10) is not quantified. The computational cost of producing Figures 2–6 (angular distance analyses, performance drop measurements, gradient norm measurements) — which are diagnostic, not deployment costs — is similarly unaccounted.
Mitigation status. Not addressed. The paper does not frame this as a limitation, does not estimate the total experimentation compute budget, and does not propose methods to reduce the hyperparameter search cost (e.g., via α prediction from architecture parameters, or via small-scale proxy experiments that transfer to larger scales). The transparency about hyperparameter values is helpful but does not retroactively account for the cost of finding them.
All Claims Are Based on Single Training Runs Without Statistical Replication
The assumption or constraint. Every perplexity value in Table 1, every accuracy value in Table 4, every reward result in Table 5, and every point on the 7B training curve (Figure 4) represents a single training run — one random seed, one initialization, one data order. The paper reports no error bars, confidence intervals, or standard deviations on any quantitative result. It does not conduct multi-seed experiments at any model scale, nor does it discuss the expected variance of the training process under different random initializations. This is methodologically standard in large-scale LLM research (due to computational cost), but it means the reader cannot distinguish between genuine method improvements and seed-dependent noise.
The consequence. Some of the reported gains are small enough that they could plausibly fall within the variance of single-run training. At LLaMA-250M, Mix-LN's pre-training perplexity advantage is 0.53 (21.39 vs. 21.92, Table 1) — approximately a 2.4% relative improvement. At LLaMA-1B, the advantage is 0.47 (18.18 vs. 18.65). While the consistency of the gain across model sizes (0.47–1.65 perplexity) argues against pure noise, the magnitude of the gain at any individual scale is not statistically established. A practitioner evaluating whether Mix-LN is worth adopting at their target scale needs to know whether the expected improvement reliably exceeds seed-to-seed variance — the paper provides no data to answer this question.
The RLHF result (Table 5) is particularly vulnerable to this critique. The reported reward gap is +0.57 (1.32 vs. 0.75) from a single RLHF training run per method. RLHF with PPO is known to be high-variance — the reward model is itself learned from finite preference data, and policy optimization exhibits substantial run-to-run variability. The paper reports no details about the RLHF configuration (PPO steps, KL penalty coefficient, reward model architecture or training), making it impossible to assess whether this result is replicable.
For the fine-tuning results (Table 4), the task-level gains are uneven — at 1B, Mix-LN outperforms Pre-LN by +5.42 points on WinoGrande but underperforms by -0.27 points on BoolQ. Without error bars, it is unclear which of these task-level differences are statistically meaningful versus within the noise of a single fine-tuning run.
What evidence exists in the paper. The consistency of Mix-LN's advantage across four model sizes (71M, 130M, 250M, 1B) in Table 1 provides informal evidence that the improvement is not a single-run fluke — the probability of Mix-LN randomly outperforming Pre-LN at all four scales under the null hypothesis of no true difference is 1/16 if runs were independent. However, the runs are not fully independent (they share architecture family, training data, optimizer, and hyperparameter selection methodology), and the gains fluctuate in magnitude (1.65, 0.71, 0.53, 0.47 perplexity), which could reflect true scale-dependent variation or seed noise.
Mitigation status. Not addressed. The paper does not acknowledge the single-run limitation, does not provide uncertainty quantification, and does not justify why the reported gains should be considered reliable despite the lack of replication. This is consistent with norms in the LLM pretraining literature (where training multiple seeds at scale is often computationally infeasible), but readers should calibrate their confidence in the precise numerical gains accordingly.
The Method Is Evaluated on a Single Architecture Family and a Single Pre-Training Corpus
The assumption or constraint. All language model experiments use the LLaMA architecture — decoder-only Transformers with RMSNorm, SwiGLU activations, and rotary position embeddings — trained on the C4 dataset. The paper states that this architecture "represents the dominant modern architecture used in production LLMs" (implicit in Section 4), which is true and justifies the choice, but it also means the findings are unvalidated for other widely-used architectures. Encoder-decoder models (T5, BART), encoder-only models (BERT, RoBERTa), mixture-of-experts architectures (Mixtral, DeepSeek-MoE), and models with different normalization variants (standard LayerNorm rather than RMSNorm, DeepNorm-stabilized architectures) are not tested. The ViT experiments (Table 6) partially address domain generality to vision but use the same basic Transformer block structure.
The pre-training corpus is similarly narrow: all language experiments use C4, a cleaned Common Crawl dataset. While C4 is a standard benchmark corpus, its properties (web text, English-dominant, specific deduplication and filtering choices) may interact with normalization strategy in ways that other corpora (code-heavy data like The Stack, multilingual data like mC4, curated high-quality data like FineWeb) might not.
The consequence. A practitioner using Mix-LN with a non-LLaMA architecture — for example, a T5-based model for translation, a BERT-based model for retrieval, or a Mixture-of-Experts model for efficient inference — cannot assume the benefits will transfer. The gradient flow analysis in Section 2.1 is architecture-agnostic (it depends only on the residual connection structure and the LN Jacobian approximation, which apply to any Transformer), so there is theoretical reason to expect benefits. However, the α sensitivity observed at 7B (Section 4.2) suggests that the optimal Post-LN fraction depends on depth and possibly width — encoder-decoder architectures with asymmetric encoder and decoder depths, or MoE architectures with different routing dynamics, might require substantially different α configurations that are unexplored.
Similarly, the C4 results may not predict performance on other data distributions. If C4's statistical properties (token frequency distribution, sequence length distribution, domain composition) affect the rate at which σ_{x'} grows during training (which controls the LN Jacobian attenuation magnitude), then the Mix-LN advantage on C4 might not hold on corpora with different properties. The paper provides no evidence one way or the other.
What evidence exists in the paper. The ViT experiments (Table 6) demonstrate that Mix-LN applies to at least one non-language domain and one non-autoregressive architecture, with small positive results (+0.04% for ViT-Tiny, +0.41% for ViT-Small). This is encouraging but limited: ViTs are still Transformer-based, the gains are small at this scale, and the training setup (ImageNet-1K, 120 epochs) is far from LLM-scale pretraining. The comparison against other normalization methods (Table 8) shows that Mix-LN outperforms Admin, Group-LN, and Sandwich-LN on LLaMA-250M, but all of these are tested only within the LLaMA/C4 setting.
Mitigation status. The paper acknowledges the scope limitation implicitly by qualifying claims to the LLaMA architecture and by testing generalization to vision as a separate experiment. Section 6 does not identify the narrow architecture/data scope as a limitation. Future work is not suggested on this dimension beyond the general need to replicate findings.
Training Instability from Post-LN Persists at Scale and Is Unresolved
The assumption or constraint. Mix-LN's design explicitly depends on applying Post-LN to a subset of early layers. The paper's own experiments (Table 1) confirm that pure Post-LN diverges at 250M+ parameters (perplexity 1409+), and DeepNorm — designed specifically to stabilize Post-LN — also diverges at 1B (1410.94). This means the very component that provides Mix-LN's gradient benefits (Post-LN on early layers) is known to cause training instability at scale. The paper addresses this by reducing α at 7B to 0.0625 (2 Post-LN layers) and using "a longer warm-up period," but this treatment is symptomatic rather than curative — the underlying instability is still present and constrains how aggressively Post-LN can be used.
The paper's own analysis (Section 2.1) predicts that as L and σ_{x'} increase, the accumulated LN Jacobian attenuation (1/σ)^{αL} will eventually destabilize training even with small α. For sufficiently deep or wide models, even a single Post-LN layer might introduce problematic gradient attenuation at the earliest layers (since the gradient from the loss must pass through all subsequent Pre-LN layers and the Post-LN layer to reach the input).
The consequence. There exists a scale — unknown, but likely above 7B given the paper's successful 7B run — at which Mix-LN becomes either untrainable (even with α approaching 0) or provides vanishing benefit (because α must be so small that the Post-LN segment is too short to meaningfully affect deep-layer gradients). This places a cap on Mix-LN's applicability that the paper does not characterize. The method may work well at the 1B–7B scale tested, but fail or provide no benefit at the 70B–405B scale of frontier production models. Since the paper's value proposition is precisely that Mix-LN improves the utilization of model parameters — which matters more at larger scales where inefficient parameter usage is most costly — this unresolved stability question directly affects the method's relevance to the settings where it would be most valuable.
Additionally, the stability problem interacts with other training design choices. The paper shows (Table 10, Appendix B) that Mix-LN is compatible with Scaled Initialization but that Scaled Embed degrades performance for both Pre-LN and Mix-LN. Practitioners using other stabilization techniques (e.g., gradient clipping, z-loss regularization, different optimizer configurations) would need to verify compatibility independently, and the paper provides no guidance for such verification.
What evidence exists in the paper. The evidence comes from three sources: (1) the Post-LN and DeepNorm divergence results at 250M and 1B (Table 1), establishing that Post-LN instability is real and not fully solved by existing stabilization methods; (2) the α sweep at 1B (Table 7), showing that performance degrades as α increases beyond 0.25 and collapses at α = 1.0; and (3) the 7B experience (Section 4.2), where α = 0.25 was unstable and had to be reduced to 0.0625. Together, these establish a clear pattern: Post-LN's instability grows with model scale, and Mix-LN's ability to harness Post-LN's benefits shrinks correspondingly.
Mitigation status. The paper explicitly identifies this limitation in Section 6:
"One limitation of Mix-LN is the introduction of Post-LN leads to instability issues when scaling up training, i.e., >7B. We conjecture that this instability is relevant to Post-LN, as using Post-LN alone also leads to the same issue. An interesting research direction is to explore how to mitigate the ineffectiveness of later layers without relying on Post-LN or to address the instability associated with Post-LN."
It proposes no partial solution — no modified Post-LN formulation, no adaptive α scheduling, no alternative mechanism for boosting deep-layer gradients that might replace Post-LN at extreme scales. The mitigation is entirely deferred to future work. This is an honest acknowledgement but also the paper's most fundamental limitation: the very mechanism that enables Mix-LN's benefits (Post-LN on early layers) is itself the source of a scaling bottleneck that will eventually cap those benefits.
Evaluation of RLHF and Downstream Gains Lacks External Validation
The assumption or constraint. The paper claims that Mix-LN improves RLHF performance (Table 5) and supervised fine-tuning (Table 4), but the RLHF evaluation relies entirely on the internal reward model score, and the fine-tuning evaluation is limited to Commonsense170K training with accuracy on eight multiple-choice benchmarks. Neither evaluation includes human preference judgments, generation quality assessments (e.g., AlpacaEval or MT-Bench for instruction-following), or evaluation on open-ended generation tasks where deep-layer representational quality might matter differently than in multiple-choice accuracy.
The RLHF result is particularly concerning because the paper reports only the final reward — a scalar produced by a separately trained reward model that is itself an imperfect proxy for human preferences. It is well-documented in the RLHF literature that language models can over-optimize the reward model, achieving high reward scores without corresponding improvements (or even with degradation) in actual output quality as judged by humans (Gao et al., 2023; Coste et al., 2024). The paper provides no guardrails against this interpretation: no KL divergence from the initial policy (to check whether the model is drifting into reward-hacking behavior), no human evaluation, and no independent benchmark evaluation of the RLHF-tuned model's outputs.
The consequence. A practitioner evaluating Mix-LN for an alignment-sensitive application cannot determine whether the +0.57 reward improvement translates to meaningfully better outputs. It is possible that Mix-LN's representations enable more effective reward optimization but that this optimization exploits idiosyncrasies in the reward model rather than genuinely improving response quality. If the reward improvement is partially or wholly attributable to reward over-optimization, then Mix-LN's advantage in RLHF may not hold under human evaluation or with a different reward model. Without external validation, the RLHF claim should be treated as suggestive rather than established.
For supervised fine-tuning, the eight multiple-choice benchmarks are standard but narrow — they measure factual and commonsense reasoning accuracy but not generation fluency, instruction-following ability, safety, or other dimensions critical to deployed LLMs. The paper's finding that Mix-LN improves BoolQ by +17.77 points at 250M but has essentially no effect on BoolQ at 1B (-0.27 points) illustrates the task-dependence and scale-dependence of the gains. A practitioner cannot predict from this data whether Mix-LN will improve performance on their specific downstream task.
What evidence exists in the paper. Table 5 reports only final reward values with no auxiliary metrics. Table 4 reports accuracy on eight benchmarks with consistent Mix-LN advantages at 250M and more mixed results at 1B. The paper does not report generation diversity, calibration, or any quality metric beyond accuracy. No example outputs are provided for qualitative comparison.
Mitigation status. Not addressed. The paper does not acknowledge the lack of external validation for the RLHF results, does not report standard RLHF diagnostic metrics (KL divergence, reward model accuracy), and does not suggest that future work should validate the RLHF gains with human evaluation or benchmark-based generation quality assessment. The fine-tuning evaluation is treated as sufficient, though its narrow scope (multiple-choice accuracy only, single dataset for SFT training) limits the conclusions that can be drawn about downstream transfer.
7. Implications and Future Directions
How This Work Changes the Landscape
This paper introduces a new design axis for Transformer training: normalization strategy need not be depth-uniform. Prior to this work, the literature treated layer normalization placement as a global architectural choice — a model uses Pre-LN or Post-LN or a modified variant (DeepNorm, Sandwich-LN, Admin) applied identically to every layer. The theoretical understanding from Xiong et al. (2020) explained why these choices have depth-dependent gradient effects, but no prior work exploited that knowledge to vary the normalization placement with depth. Mix-LN demonstrates that a simple depth-conditioned assignment — Post-LN in the first ~25% of layers, Pre-LN in the remainder — yields consistent improvements across model scales (71M to 7B), downstream tasks, and even modalities (language and vision). This is not merely a new normalization variant; it is evidence that the optimal inductive bias for early layers differs from the optimal inductive bias for later layers, and that architectures should reflect this asymmetry.
The conceptual shift is modest in magnitude but precise in implication: the field has implicitly assumed that layer uniformity is desirable (all layers share the same architecture, same normalization, same initialization), and Mix-LN demonstrates that breaking this uniformity along the depth axis — specifically, matching the normalization strategy to each depth regime's gradient flow properties — produces strictly better results. This parallels earlier realizations in the field: just as different attention heads learn different functions, and different layers in a CNN learn features at different levels of abstraction, different depth ranges in a Transformer face different gradient flow conditions and benefit from different normalization strategies. Mix-LN is the first method to operationalize this insight for layer normalization.
This work also resolves the tension between the pruning literature's findings and the goal of efficient pretraining. Gromov et al. (2024), Men et al. (2024), and others documented that deep layers in Pre-LN LLMs are dispensable — they can be pruned or removed with minimal accuracy loss. The pruning community interpreted this as an opportunity for compression. This paper demonstrates that the dispensability is an artifact of suboptimal training, not an architectural property. By fixing the gradient flow problem with Mix-LN, deep layers become more important (Figure 6b shows larger performance drops when deep Mix-LN layers are removed compared to Pre-LN), and the model achieves lower perplexity (Table 1) and better downstream performance (Tables 4, 5). The implication is that the research effort invested in pruning deep layers should be partially redirected toward training them better in the first place — compression is valuable, but preventing the redundancy from arising is even more valuable, since it increases total model capacity rather than merely recovering wasted parameters.
Practically, Mix-LN shifts the cost-benefit calculus for architecture modifications in LLM training. The method introduces zero additional parameters, zero additional FLOPs, and zero changes to training data or objectives — it is a pure rearrangement of existing operations. This means the barrier to adoption is low: any team training a Pre-LN Transformer can switch to Mix-LN by changing the normalization placement in a fraction of layers, with no downstream pipeline changes. The gains are modest in absolute terms (0.25–1.65 perplexity across scales, Table 1) but accumulate across the full training budget, and — critically — amplify during fine-tuning (the 4.41% downstream accuracy gain at 250M in Table 4 is proportionally larger than the 2.4% pre-training perplexity improvement). This amplification pattern suggests that better-trained deep layers provide representational capacity that downstream tasks can exploit even when pre-training metrics show only modest improvement, making Mix-LN's benefits partially hidden from standard pre-training evaluation.
The work also makes post-hoc sense of conflicting results in the layer effectiveness literature. Gromov et al. (2024) and Lad et al. (2024) found deep-layer redundancy in LLaMA and GPT models (all Pre-LN). Liu et al. (2020) and Takase et al. (2022) found early-layer redundancy in Post-LN models. These findings appeared contradictory — which layers are actually important? Mix-LN's diagnosis reconciles them: both patterns are caused by the same mechanism (LN Jacobian attenuation or residual accumulation affecting gradient flow), just applied at different depths depending on the normalization choice. There is no universal "important depth range" for Transformers; the importance pattern is a function of how the architecture routes gradients through depth. This reconciliation is valuable because it prevents the field from drawing incorrect general conclusions (e.g., "deep layers are always less useful") from observations that are specific to the dominant Pre-LN paradigm.
Finally, this paper opens the door to a broader research program of depth-heterogeneous Transformer design. If normalization placement should vary with depth, what else should? Learning rates? Initialization scales? Dropout rates? The number of attention heads? The paper's success with a single depth-dependent design choice suggests that the uniform-layer assumption — a near-universal default in Transformer architecture design — may be leaving performance on the table across multiple design dimensions. Mix-LN provides both a template (identify a design choice with depth-dependent effects, apply the variant suited to each depth regime) and a proof of concept (this template yields measurable gains in practice) for this broader program.
Follow-Up Research This Work Enables
Characterizing the scaling law for optimal α as a function of model depth and width. The paper's most pressing open question is how to choose α for arbitrary model scales without expensive per-scale sweeps. At 1B parameters (24 layers), α = 0.25 is optimal (Table 7). At 7B (32 layers), α = 0.25 is unstable and α = 0.0625 is used instead (Section 4.2). This suggests the optimal fraction of Post-LN layers decreases with total depth — consistent with the theory that Post-LN's accumulated attenuation (1/σ)^{αL} must stay below some threshold for training stability. A systematic study would train LLaMA models at 4–5 different depths (e.g., 12, 24, 36, 48, 64 layers) with multiple α values per depth, measuring both stability (does training converge?) and final perplexity. The goal would be to fit a functional form α_opt(L) = c / L or α_opt(L) = c / log(L) that predicts the largest stable α at a given depth. A successful scaling law would make Mix-LN immediately usable at production scales (70B, 405B) without prohibitive hyperparameter search costs. A negative result — e.g., α_opt depends on width, training tokens, or learning rate schedule in ways not captured by depth alone — would be equally informative, establishing the practical limits of Mix-LN's plug-and-play deployment.
Training a stabilized Post-LN model at 1B+ scale to directly test the gradient flow hypothesis. The paper's mechanistic argument — that Pre-LN causes deep-layer inefficacy through gradient decay — is supported by the controlled 130M comparison (Figure 3) and the gradient norm measurements at initialization (Figure 3e). But the paper cannot demonstrate the inverse pattern (Post-LN causing deep-layer efficacy and early-layer inefficacy) at a scale where it would be practically meaningful, because pure Post-LN diverges at 250M+ (Table 1). A strong confirmatory experiment would use stabilization techniques beyond those tested in the paper — perhaps combining DeepNorm with longer warmup, gradient clipping, or the Scaled Initialization that Table 10 shows is Mix-LN-compatible — to successfully train a Post-LN LLaMA-1B or larger model. If such a model converges, measuring angular distance (Figure 3a), performance drop (Figure 3c), and downstream fine-tuning performance would directly test whether Post-LN at scale produces inverse layer importance patterns to Pre-LN, as the theory predicts. If Post-LN cannot be stabilized at 1B+ despite aggressive tuning, this would suggest the gradient vanishing problem is fundamental and that Mix-LN's strategy of using limited Post-LN is the best available option — supporting the paper's implicit claim that pure Post-LN is not viable at scale.
Combining Mix-LN with depth-dependent learning rates. Mix-LN addresses the supply side of deep-layer training: it provides larger gradient norms to deeper layers (Figure 6a). But the paper's angular distance results (Figure 5) show that even under Mix-LN, the deepest layers still show somewhat higher similarity to their neighbors than middle layers — the gradient improvement helps but does not fully equalize layer contributions. A natural extension is to address the demand side: use layer-specific learning rate multipliers that are higher for deeper layers (compensating for any remaining gradient diminishment) and lower for early layers (where Mix-LN's Post-LN segment already provides strong gradients). This could be implemented as a simple schedule, e.g., η_ℓ = η_base × (1 + β × ℓ/L) for some β determined by a small grid search. The experiment would compare Mix-LN with uniform learning rates (the current paper) against Mix-LN with depth-dependent learning rates, measuring whether the combination eliminates the final angular distance convergence in the deepest layers (Figure 5) and whether downstream fine-tuning gains increase further. A positive result would establish that gradient supply (via Mix-LN) and learning rate demand are complementary levers for deep-layer quality.
Evaluating Mix-LN on encoder-decoder and mixture-of-experts architectures. The paper demonstrates Mix-LN on decoder-only LLaMA models and on encoder-only ViTs (Table 6), but the most widely used architecture in production translation, summarization, and multilingual models remains encoder-decoder (T5, BART). In an encoder-decoder model, the encoder and decoder have different roles (encoding input → representing meaning; decoding → generating output) and different depth considerations (encoder depth affects input representation quality; decoder depth affects generation quality). The optimal Post-LN ratio might differ between encoder and decoder — perhaps the encoder benefits from more Post-LN (since its outputs must compress information into a fixed-length representation) while the decoder needs Pre-LN stability for autoregressive generation. A controlled experiment would train T5-base or T5-large variants with Mix-LN applied independently to the encoder and decoder, sweeping α_enc and α_dec on a modest scale, and evaluating on both understanding (GLUE/SuperGLUE) and generation (XSum, WMT) tasks. For mixture-of-experts architectures (Mixtral, DeepSeek-MoE), the interaction between expert routing and the Post-LN residual constraint is unexplored — Post-LN early layers might affect how tokens are routed to experts, potentially changing the load-balancing dynamics. A small-scale MoE experiment (e.g., OLMoE-1B) with Mix-LN would reveal whether expert specialization and normalization strategy interact, which is practically important as MoE becomes the dominant paradigm for efficient scaling.
Stress-testing Mix-LN's RLHF claims with human evaluation and reward model diagnostics. The paper's RLHF result (Table 5) reports only a final reward score from a single PPO run, with no auxiliary metrics. This is insufficient to establish that Mix-LN improves alignment quality rather than enabling more effective reward over-optimization. A rigorous follow-up would replicate the 1B RLHF experiment with: (a) multiple random seeds (at least 3) for both Pre-LN and Mix-LN to establish the variance of the reward improvement; (b) KL divergence from the initial policy tracked throughout PPO training, to detect whether Mix-LN achieves higher reward by drifting further from the base model (which would suggest over-optimization); (c) evaluation on an independent benchmark like AlpacaEval 2.0 or MT-Bench, which uses a separate evaluator (GPT-4 as judge) that is not the same reward model used in training; and (d) a small-scale human preference study (even 100 judgments on a held-out set) comparing Mix-LN and Pre-LN outputs at matched reward model scores, to calibrate whether equal reward corresponds to equal human preference quality. If Mix-LN shows higher reward and higher independent benchmark scores and higher human preference at matched reward, the alignment benefit is genuine. If Mix-LN shows higher reward but equivalent or worse independent scores, the benefit is likely reward hacking — an important negative result that would substantially change how the community interprets Mix-LN's downstream value.
Extending the depth-heterogeneous design principle beyond normalization. Mix-LN demonstrates that varying a single design choice with depth yields measurable gains. This invites systematic exploration of other design choices that might benefit from depth-conditioning. Candidates include: (1) Initialization scale — perhaps deeper layers should use smaller initial weights (since Pre-LN later layers operate on larger residual magnitudes) while early Post-LN layers use standard initialization. (2) Attention head count — if deep layers under Pre-LN learn more specialized features, they might benefit from more heads to capture finer-grained patterns, while early layers (which learn broader features under Post-LN) might need fewer heads. (3) FFN expansion ratio — the optimal FFN hidden size might vary with depth, since different depth ranges have different gradient magnitudes and learn different types of transformations. Each of these could be tested by modifying the LLaMA architecture to vary the design choice with depth (e.g., heads per layer increasing from 8 in early layers to 16 in late layers) and comparing to a uniform baseline at matched total parameters. A systematic study sweeping multiple depth-heterogeneous axes simultaneously (possibly with a small neural architecture search over depth-dependent configurations) could reveal whether the gains from depth-conditioning in different design dimensions are additive, synergistic, or redundant — establishing the ceiling for depth-heterogeneous Transformer design as a research program.
Practical Applications and Downstream Use Cases
Cost-efficient pretraining of custom LLMs at the 1B–8B scale. The paper's results (Table 1) show that Mix-LN reduces perplexity by 0.47–1.65 across scales without any increase in parameters or training compute. For teams training custom LLMs in the 1B–8B range — a common scenario in industry for domain-specific models — this translates directly to better models for the same training budget. At LLaMA-1B, Mix-LN achieves 18.18 perplexity vs. Pre-LN's 18.65 on the same 5B tokens of training. To match Mix-LN's perplexity, a Pre-LN model would need to train longer (more tokens, more compute) or be scaled up (more parameters, more inference cost). The paper's fine-tuning results amplify this value: at 250M parameters, Mix-LN's +4.41% average downstream accuracy over Pre-LN (Table 4) means the model is strictly better at the tasks users care about, not just at next-token prediction. The zero-cost nature of the modification — change a few lines of model definition code, no new operations, no data pipeline changes — makes the adoption risk low for teams already training Transformer LMs. The main practical barrier is the α sensitivity at scale (Section 4.2): a team training at 8B parameters would need to budget for a small α sweep (perhaps 3–4 values around 0.05–0.10) to find a stable configuration, adding 3–4× the training cost of a single run. Whether this overhead is worth it depends on how many models the team plans to train — for a single production model, the sweep cost is amortized over the lifetime of the deployed system, likely making it worthwhile.
Improving data efficiency in fine-tuning and alignment workflows. The paper's supervised fine-tuning (Table 4) and RLHF (Table 5) results suggest that Mix-LN pre-trained models extract more value from limited downstream data. At 250M, Mix-LN's fine-tuning gain over Pre-LN (+4.41% average) is proportionally much larger than its pre-training gain (2.4% perplexity reduction), indicating that the representations learned in deeper layers under Mix-LN are particularly valuable for adapting to new tasks. For applications where fine-tuning data is scarce (domain-specific medical/legal QA, low-resource language tasks, specialized coding tasks), starting from a Mix-LN pre-trained checkpoint could yield better performance than starting from a Pre-LN checkpoint with the same pre-training budget — or equivalent performance with less fine-tuning data. The RLHF result (1.32 vs. 0.75 reward, Table 5), while needing further validation (see Limitation in prior sections), suggests that alignment procedures might also benefit from Mix-LN's deeper representations, potentially reducing the amount of expensive human preference data needed to reach a target alignment quality. This is particularly relevant for organizations building aligned models in specialized domains where preference data collection is costly.
A drop-in upgrade for existing LLaMA-based training pipelines. Because Mix-LN works with RMSNorm, SwiGLU, and the standard LLaMA architecture — the most common recipe in open-source LLM training — it requires no ecosystem changes. Training frameworks (Megatron-LM, DeepSpeed, HuggingFace Transformers), inference engines (vLLM, TensorRT-LLM), and fine-tuning libraries (LoRA, QLoRA) all operate on the same linear algebra operations regardless of normalization placement. The paper's ViT results (Table 6) further demonstrate that the method is not specific to autoregressive language modeling, suggesting it could be applied in multimodal training pipelines (e.g., LLaVA-style vision-language models) where both the language backbone and the vision encoder could use Mix-LN. The combination of zero inference overhead and broad architectural compatibility makes Mix-LN one of the lowest-friction improvements available to teams training Transformers at scale — on par with switching activation functions or tuning learning rate schedules, but with a mechanism (gradient rebalancing) that addresses a specific, well-characterized pathology rather than offering a generic hyperparameter tweak.
When to Prefer This Method
The paper does not position Mix-LN against a named, widely-used alternative normalization strategy in a clear "prefer A over B when X" framework — Pre-LN is the default, Post-LN is known to be unstable, and the other variants (DeepNorm, Admin, Sandwich-LN) are not positioned as competitive alternatives that a practitioner would seriously consider in production. Rather, Mix-LN is presented as a strict improvement over the universal default (Pre-LN) at the scales tested, with the caveat that α requires scale-dependent tuning. The decision rule implied by the paper is therefore simple:
-
Prefer Mix-LN over Pre-LN when training a Transformer model at any scale where a stable
αcan be found (demonstrated up to 7B withα = 0.0625), and when the cost of theαhyperparameter search is acceptable relative to the training budget. The expected benefit is 0.25–1.65 perplexity improvement during pretraining (Table 1) with proportionally larger gains during supervised fine-tuning (Table 4) and potential gains during RLHF (Table 5), at zero additional parameter or inference cost. -
Prefer standard Pre-LN when training at a scale or architecture where the
αsearch cost is prohibitive relative to the total training budget, when training stability is paramount and even the small Post-LN segment introduces unacceptable risk, or when the model depth exceeds the regime where Mix-LN has been validated (the paper tests up to 32 layers at 7B; behavior at 64–128 layers is unknown). -
Do not use pure Post-LN or DeepNorm at scales above ~250M parameters based on this paper's evidence (Table 1 shows both diverging at 1B), unless combined with the depth-limiting strategy of Mix-LN.
The paper does not provide enough information to compare Mix-LN against other normalization variants (Group-LN, Sandwich-LN) at production scale, since those comparisons are limited to a single 250M experiment (Table 8). A practitioner choosing between Mix-LN and a method from that literature would need additional evidence at their target scale.