---
title: LayerNorm After First Feed-Forward Linear Layer
url: https://www.emergentmind.com/topics/layernorm-after-first-feed-forward-linear-layer
type: topic
---

# LayerNorm After First Feed-Forward Linear Layer

LayerNorm After First Feed-Forward Linear Layer refers to the architectural practice of inserting a Layer Normalization (LayerNorm) operation immediately after the initial linear transformation in a feed-forward block (typically $W_1x + b_1$), often before or after the nonlinear activation function. This design choice has concrete mathematical, optimization, geometric, and initialization implications in deep neural networks, especially in transformer and MLP-based architectures.

## 1. Mathematical Formulation and Placement

The canonical setting for "LayerNorm after first feed-forward linear layer" appears in both standard MLPs and Transformer FFN blocks. Consider an input $x \in \mathbb{R}^d$, and a two-layer FFN:

\[
\begin{align*}
a &= W_1x + b_1 \\
h &= \phi(a) \\
y &= W_2h + b_2
\end{align*}
\]

Layer Normalization can be interposed:
- **After the first linear layer**: $h = \phi(\text{LayerNorm}(W_1x + b_1))$
- **After activation**: $h = \text{LayerNorm}(\phi(W_1x + b_1))$

In "NormFormer: Improved Transformer Pretraining with Extra Normalization," the extra LayerNorm is applied _after the activation_ within the FFN:
\[
x_\ell \rightarrow \text{LN}(x_\ell) \rightarrow W_1 + b_1 \rightarrow \text{GELU}(~) \rightarrow \mathbf{LN}(~) \rightarrow W_2 + b_2 \rightarrow x_{\ell+1} = x_\ell + \cdots
\]
with
\[
\text{LayerNorm}(h) = \gamma \circ \frac{h - \mathbb{E}[h]}{\sqrt{\mathrm{Var}[h] + \epsilon}} + \beta
\]
where learnable scale $\gamma$ is initialized to $1$, shift $\beta$ to $0$, and $\epsilon$ is a small constant (e.g., $10^{-5}$) [2110.09456].

## 2. Impact on Optimization and Gradient Flow

In Pre-LayerNorm Transformers, gradients to early layers are much larger than to deeper layers, causing a gradient magnitude mismatch and inefficient utilization of parameters. The insertion of an extra LayerNorm after the first feed-forward (linear+GELU) addresses this by:
- Downscaling gradients in early FFN layers and upscaling them in later layers.
- Aligning the magnitude of gradient updates across all layers, which improves both convergence speed and final validation perplexity.

Empirical evidence from [2110.09456] demonstrates:
- The gradient norm band is significantly tightened across layers.
- Removing this single FFN LayerNorm (with other modifications held) raises validation perplexity (e.g., $15.88 \to 16.14$ for a 125M param language model), confirming its unique utility.

## 3. Effects on Initialization and Predictive Bias

Normalization placement fundamentally alters the distribution of network outputs at initialization. As rigorously studied in [2505.11312]:

- **Norm Before Activation** ("Linear → LayerNorm → ReLU") produces _prejudiced initializations_: initial predictions are biased toward a subset of classes, with per-class output distributions (\(G_0\)) spread toward the extremes.
- **Norm After Activation** ("Linear → ReLU → LayerNorm") yields _neutral initializations_: all output logits are symmetrically distributed, with $G_0$ tightly centered at $1/C$ (where $C$ is the number of classes), regardless of depth.

This is mathematically guaranteed by LayerNorm's per-example feature centering, eliminating the mean fluctuation term from the output logit distribution at all depths and collapsing inter-initialization variance ($\Gamma = 0$).

Empirical validation [2505.11312]:
- In 20-layer MLPs, "LN after ReLU" produces uniform class distributions at initialization ($G_0 \approx 0.5$ for $C=2$), while "LN before ReLU" amplifies bias with depth.
- On CIFAR-10, the neutralization effect persists in real structured datasets.

## 4. Geometric and Bayesian Complexity Costs

LayerNorm imposes strict geometric constraints on the activations it normalizes. By mean-centering, it projects representations to a codimension-1 hyperplane ($\mathbf{1}^\top x = 0$). This constrains the downstream weight matrix $W \in \mathbb{R}^{m \times d}$ and reduces the Local Learning Coefficient (LLC) [2603.27432], or effective Bayesian parameter count, by exactly $m/2$:
\[
\Delta\lambda = \frac{m}{2}
\]
where $m$ is the output dimension of the subsequent weight matrix. No such drop occurs for RMSNorm.

Controlled experiments confirm that this LLC reduction is:
- Exact and binary for affine (flat) hyperplanes created by LayerNorm.
- Absent for full-dimensional spheres (RMSNorm) or for non-affine, curved surfaces.
- Robust across varying $m$, $d$ and in presence of affine bias with simplex data.

This structural reduction in model complexity occurs before any training and reflects true reductions in the “observable” directions in parameter space.

## 5. Comparison to Standard and Alternative Normalization Placements

Baseline transformer architectures ("Pre-LN" and "Post-LN") insert only one LayerNorm per FFN: either immediately before the sublayer or after the residual addition [2204.04477]. No LayerNorm is present between the two FFN linear layers ($W_1$ and $W_2$) by default.

- **FoundationLayerNorm** [2204.04477]: Scales the residual input by $\alpha<1$ before post-FFN LayerNorm, enabling deep (1,000 layer) training, but does not place an extra LayerNorm after the first FFN linear. No empirical evidence is shown for gradient norm balancing at the level achieved by NormFormer's FFN LN insertion.
- **RMSNorm**: Projects onto a sphere, preserving the LLC ($\Delta\lambda = 0$), and does not enforce mean centering.

This highlights the uniqueness of the LayerNorm-after-FFN design used in NormFormer [2110.09456] and the "LN after ReLU" variant [2505.11312] for balancing gradient flow, promoting neutrality, and structurally reducing complexity.

## 6. Implementation, Initialization, and Practical Training Effects

Implementation of LayerNorm after the first FFN linear uses framework-native fused LayerNorm modules with standard parameters ($\epsilon$ for numerical stability, per-feature $\gamma$, $\beta$). No special initialization or training tricks are needed [2110.09456].

This extra normalization has measurable practical benefits:
- Low computational overhead (∼$0.4\%$ in parameters).
- Faster pretraining convergence: NormFormer reaches the baseline perplexity in 60% of the GPU-time, or achieves 0.27 lower perplexity in fixed compute [2110.09456].
- Consistent gains in downstream transfer: GLUE average score increases by $\sim1.9$ points for a 125M param model.
- Streamlined learning dynamics: bias-neutral initializations avoid two-phase dynamics and accelerate transition to informative gradients [2505.11312].

## 7. Design Principles and Architectural Guidelines

Key empirically and theoretically supported guidelines [2505.11312, 2110.09456]:
- Use LayerNorm immediately after nonlinearity (Linear → ReLU/GELU → LN) to guarantee neutral initial state and balanced convergence.
- In transformer FFN blocks, extra LN after the first linear+GELU (NormFormer's scheme) stabilizes gradients at all depths and improves both speed and accuracy.
- RMSNorm may be preferred if preservation of full LLC is desired, but does not provide gradient norm balancing at the FFN sublayer.
- The geometric cost (LLC penalty) of LayerNorm is structural and should be considered if Bayesian model complexity is critical to application.
- In highly correlated datasets (e.g., MNIST), residual bias after LN is possible; introducing small random input shifts can mitigate this effect.

---

**Summary Table: LayerNorm Placement Effects**

| Placement Variant                            | Initialization Bias | Gradient Norm Balancing | LLC Drop      |
|----------------------------------------------|--------------------|------------------------|---------------|
| Linear → LayerNorm → ReLU                    | Prejudiced         | No                     | $m/2$         |
| Linear → ReLU → LayerNorm                    | Neutral            | Yes                    | $m/2$         |
| Linear → RMSNorm (any point)                 | Neutral            | No                     | $0$           |
| Linear → ReLU (no normalization)             | Prejudiced         | No                     | $0$           |
| NormFormer: Pre-LN + LN after FFN Linear+Act | Neutral            | Yes                    | $m/2$         |

The insertion of LayerNorm after the first feed-forward linear (and activation) is a principled and empirically validated strategy for promoting gradient stability, initialization neutrality, and predictable geometric constraints in deep neural network design [2110.09456, 2505.11312, 2603.27432, 1607.06450].

Source: https://www.emergentmind.com/topics/layernorm-after-first-feed-forward-linear-layer