LeetLLM
My PlanLearnGlossaryTracksPracticeBlog
LeetLLM

Your go-to resource for mastering AI & LLM systems.

Product

  • Learn
  • Glossary
  • Tracks
  • Practice
  • Blog
  • RSS

Legal

  • Terms of Service
  • Privacy Policy

© 2026 LeetLLM. All rights reserved.

All Topics
Your Progress
0%

0 of 196 articles completed

🛠️Computing Foundations0/9
Git, Shell, Linux for AIDocker for Reproducible AIPython for AI EngineeringNumPy and Tensor ShapesCUDA for ML TrainingMPS & Metal for ML on MacData Structures for AISQL and Data ModelingAlgorithms for ML Engineers
📊Math & Statistics0/9
Gradients and BackpropVectors, Matrices & TensorsLinear Algebra for MLAdam, Momentum, SchedulersProbability for Machine LearningBayesian Inference from EvidenceStatistics and UncertaintyDistributions and SamplingHypothesis Tests, Intervals, and pass@k
📚Preparation & Prerequisites0/13
Neural Networks from ScratchCNNs from ScratchTraining & BackpropagationSoftmax, Cross-Entropy & OptimizationRNNs, LSTMs, GRUs, and Sequence ModelingAutoencoders and VAEsThe Transformer Architecture End-to-EndLanguage Modeling & Next TokensFrom GPT to Modern LLMsPrompt Engineering FundamentalsCalling LLM APIs in ProductionFirst AI App End-to-EndThe LLM Lifecycle
🧮ML Algorithms & Evaluation0/15
Linear Regression from ScratchLogistic Regression and MetricsDecision Trees, Forests, and BoostingReinforcement Learning BasicsValidation and LeakageStatistical Learning and GeneralizationConformal Prediction and UncertaintyClustering and PCAProbabilistic Generative ModelsCore Retrieval AlgorithmsDecoding AlgorithmsExperiment Design and A/B TestingCausal Inference and CounterfactualsPyTorch Training LoopsDataset Pipelines and Data Quality
📦Production ML Systems0/6
Feature Engineering for Production MLBatch and Streaming Feature PipelinesGradient Boosted Trees in ProductionRanking and Recommendation SystemsForecasting and Anomaly DetectionMonitoring Predictive Models
🧪Core LLM Foundations0/8
The Bitter Lesson & ComputeBPE, WordPiece, and SentencePieceStatic to Contextual EmbeddingsPerplexity & Model EvaluationFile Ingestion for AIChunking StrategiesLLM Benchmarks & LimitationsInstruction Tuning & Chat Templates
🧰Applied LLM Engineering0/25
Dimensionality Reduction for EmbeddingsCoT, ToT & Self-Consistency PromptingFunction Calling & Tool UseMCP & Tool Protocol StandardsContext EngineeringPrompt Injection DefenseResponsible AI GovernancePrivacy-Preserving Machine LearningData Labeling and Human FeedbackEvaluating AI AgentsProduction RAG PipelinesHybrid Search: Dense + SparseReranking and Cross-Encoders for RAGRAG Evaluation for Reliable AnswersLLM-as-a-Judge EvaluationBias & Fairness in LLMsHallucination Detection & MitigationLLM Observability & MonitoringExperiment Tracking with MLflow and W&BPrompt Optimization with DSPyModel Versioning & DeploymentSemantic Caching & Cost OptimizationLLM Cost Engineering & Token EconomicsModel Gateways, Routing, and FallbacksDesign an Automated Support Agent
🎓Portfolio Capstones0/9
Capstone: Delivery ETA PredictionCapstone: Product RankingCapstone: Demand ForecastingCapstone: Image Damage ClassifierCapstone: Production ML PipelineCapstone: Document QACapstone: Eval DashboardCapstone: Fine-Tuned ClassifierCapstone: Reproducible ML Study
🧠Transformer Deep Dives0/8
Sentence Embeddings & Contrastive LossEmbedding Similarity & QuantizationScaled Dot-Product AttentionVision Transformers and Image EncodersPositional Encoding: RoPE & ALiBiLayer Normalization: Pre-LN vs Post-LNMechanistic InterpretabilityDecoding Strategies: Greedy to Nucleus
🧬Advanced Training & Adaptation0/17
Scaling Laws & Compute-Optimal TrainingPre-training Data at ScaleBuild GPT from Scratch LabJAX for PyTorch ResearchersContinued Pretraining for Domain ShiftSynthetic Data PipelinesSupervised Fine-Tuning PipelineMixed Precision TrainingDistributed Training: FSDP & ZeROLoRA & Parameter-Efficient TuningTraining Run OperationsReward Modeling from Preference DataRLHF & DPO AlignmentConstitutional AI & Red TeamingRLVR & Verifiable RewardsKnowledge Distillation for LLMsModel Merging and Weight Interpolation
🤖Advanced Agents & Retrieval0/16
Vector DB Internals: HNSW & IVFAdvanced RAG: HyDE & Self-RAGGraphRAG & Knowledge GraphsRAG Security & Access ControlStructured Output GenerationReAct & Plan-and-ExecuteGuardrails & Safety FiltersCode Generation & SandboxingComputer-Use / GUI / Browser AgentsHuman-in-the-Loop Agent ArchitectureAI Coding Workflow with AgentsAgent Memory & PersistenceAgent Failure & RecoveryRecursive Language Models (RLM)Multi-Agent OrchestrationCapstone: Production Agent
⚡Inference & Production Scale0/31
Inference: TTFT, TPS & KV CacheMulti-Query & Grouped-Query AttentionKV Cache & PagedAttentionPrefix Caching and Prompt CachingFlashAttention & Memory EfficiencyContinuous Batching & SchedulingScaling LLM InferenceModel Parallelism for LLM InferenceAccelerator Architecture Field GuideCorrect GPU Profiling and BenchmarkingGPU Kernel Performance Engineering LabGPU Parallel Primitives LabGPU Compilers and Kernel LanguagesGEMM From Naive CUDA to Tensor CoresLow-Precision GPU ProgrammingGPU Runtime Execution LabAI-Generated Kernel EngineeringModel Quantization: GPTQ, AWQ & GGUFLocal LLM DeploymentSLM Specialization & Edge DeploymentSpeculative DecodingLong Context Window ManagementMixture of Experts ArchitectureMamba & State Space ModelsReasoning & Test-Time ComputeAdvanced MLOps & DevOps for AIGPU Serving & AutoscalingGPU Failure OperationsDistributed Inference Data PlaneServing Benchmark EngineeringA/B Testing for LLMs
🏗️System Design Capstones0/9
Content Moderation SystemCode Completion SystemMulti-Tenant LLM PlatformLLM-Powered Search EngineVision-Language Models & CLIPMultimodal LLM ArchitectureDiffusion Models: Images & TextReal-Time Voice AI AgentReasoning Agent System Design
🎤AI Lab Interviewing0/4
AI Lab Coding Interview: Python SystemsAI Lab System Design InterviewAI Lab Behavioral InterviewAI Lab Technical Presentation
🔬Project Deep Dives0/17
Deep Dive - vLLMDeep Dive - SkyRLDeep Dive - FlashAttentionDeep Dive - FlashInferDeep Dive - DeepGEMMDeep Dive - NCCLDeep Dive - MegatronDeep Dive - DeepSpeedDeep Dive - RayDeep Dive - MLflowDeep Dive - PyTorchDeep Dive - TransformersDeep Dive - SGLangDeep Dive - slimeDeep Dive - DeepEPDeep Dive - TinkerDeep Dive - Light-PEFT
Back to Topics
LearnTransformer Deep DivesLayer Normalization: Pre-LN vs Post-LN
🧠HardTransformer Architecture

Layer Normalization: Pre-LN vs Post-LN

Understand LayerNorm mechanics, Pre-LN versus Post-LN placement, RMSNorm simplification, gradient stability, and hybrid normalization layouts for deep transformers.

34 min read
Learning path
Step 100 of 196 in the full curriculum
Positional Encoding: RoPE & ALiBiMechanistic Interpretability

Personalize this lesson

Adapt explanations and teaching visuals to your background and preferred voice.

A decoder keeps one hidden vector for each token as it moves through a stack. That shared per-token state is the residual stream. Suppose one state is [3.0, 1.0, -1.0, 5.0]: each transformer block adds an attention update and a feed-forward update to the same vector. After many additions, one coordinate can dominate later dot products or output scores, while the backward signal becomes uneven across layers.

LayerNorm gives each token a scale checkpoint. It measures that token's feature values, centers them, and rescales their spread; learned scale and shift can then restore useful asymmetry. The placement question is just as important as the arithmetic: does the residual shortcut bypass this checkpoint, or does it cross the checkpoint after every block?

The previous lesson's RoPE and ALiBi encode token position. They don't regulate residual magnitude. Keep the vector visible as we move from LayerNorm's arithmetic to Pre-LN and Post-LN blocks, then use small diagnostics to test what changes at depth.


One residual stream, many writes

Start with the block abstraction output = input + sublayer(input). The input takes a skip path, while the sublayer writes an attention or feed-forward update. LayerNorm can inspect the input before that write, or inspect the sum after it.

The distinction matters in both directions. During the forward pass, it controls which scale reaches the next block. During backpropagation, it controls whether the residual shortcut contributes a direct identity term or passes through a normalization Jacobian. Neither layout guarantees stable training under every learning rate or initialization.[1]Reference 1On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745

Residual updates aren't literal multipliers, but a toy multiplicative drift makes the risk easy to see: a scale factor of 1.1 repeated across fifty blocks gives 1.1^50, roughly 117. The real stack adds vectors, so use that number as a warning about repeated transformations, not as a prediction of a model's hidden-state norm.

LayerNorm's job is to re-center and rescale each token's hidden state. Before learned scale and shift are applied, the normalized vector sits near zero with a spread close to one. First, compute that operation for one vector; then the placement choices will have something concrete to act on.


LayerNorm by hand

Use the running vector as one token's hidden state:

layernorm-by-hand.py
1x = [3.0, 1.0, -1.0, 5.0]

First find its mean: (3 + 1 + (-1) + 5) / 4 = 2.0. Centering removes this shared offset.

Subtract that mean from each feature, giving [1.0, -1.0, -3.0, 3.0]. The distances between features stay the same; only their origin moves.

Now compute the population variance: (1^2 + (-1)^2 + (-3)^2 + 3^2) / 4 = (1 + 1 + 9 + 9) / 4 = 5.0. LayerNorm divides by the feature count here, rather than by d - 1 as an unbiased sample estimate.

The standard deviation is sqrt(5.0) ≈ 2.236, so dividing the centered vector gives [0.45, -0.45, -1.34, 1.34]. Its mean is zero and its spread is one, up to rounding.

Finally, learned scale and shift can change that standardized result. With γ = [1.0, 1.0, 1.0, 1.0] and β = [0.0, 0.0, 0.0, 0.0], the output stays the same. During training, the model learns γ and β that help it predict the next token.

The operation runs independently for each token. Other tokens in the batch don't enter this token's statistics, so batch-size-one autoregressive generation uses the same rule as larger-batch training.

If you feed 16 token hidden states through LayerNorm, which numbers determine the mean and variance for one token?

Answer

Only that token's own feature values determine its mean and variance. The other 15 tokens in the batch don't affect it, which is why LayerNorm behaves the same during batch-size-1 decoding and during larger-batch training.

Turn that hand calculation into a reusable function. It accepts one vector and optional gamma, beta, and eps values, then checks the rounded result and two edge cases. Predict why setting every scale entry to zero returns an all-zero vector before running it.

layernorm-by-hand-2.py
1from math import isclose, sqrt 2 3def layer_norm(values, gamma=None, beta=None, eps=1e-5): 4 if gamma is None: 5 gamma = [1.0] * len(values) 6 if beta is None: 7 beta = [0.0] * len(values) 8 mean = sum(values) / len(values) 9 variance = sum((value - mean) ** 2 for value in values) / len(values) 10 scale = sqrt(variance + eps) 11 return [ 12 gamma_i * ((value - mean) / scale) + beta_i 13 for value, gamma_i, beta_i in zip(values, gamma, beta) 14 ] 15 16normalized = layer_norm([3.0, 1.0, -1.0, 5.0]) 17rounded = [round(value, 2) for value in normalized] 18assert rounded == [0.45, -0.45, -1.34, 1.34] 19assert isclose(sum(normalized), 0.0, abs_tol=1e-12) 20assert layer_norm([3.0, 1.0], gamma=[0.0, 0.0]) == [0.0, 0.0]

The general formula

The function contains no hidden sequence operation. For a vector x ∈ ℝ^d (one token's hidden state), LayerNorm computes:[2]Reference 2Layer Normalization.https://arxiv.org/abs/1607.06450

LayerNorm(x)=γ⊙x−μσ2+ϵ+β\text{LayerNorm}(x) = \gamma \odot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \betaLayerNorm(x)=γ⊙σ2+ϵ​x−μ​+β

where:

  • μ = (1/d) Σ x_i: mean over the hidden dimension
  • σ² = (1/d) Σ (x_i - μ)²: variance over the hidden dimension
  • γ, β ∈ ℝ^d: learnable scale and shift parameters
  • ε: tiny positive constant to prevent division by zero (10−510^{-5}10−5 in this example; defaults vary by implementation). In mixed precision (FP16 / BF16), keep ε large enough that the variance floor doesn't underflow to zero. Common production floors sit around 1e-5 to 1e-6, not aggressive 1e-12 values that can look fine in FP32 and blow up in half precision.

Subtracting the mean centers the values around zero. Dividing by the standard deviation squeezes them to a similar scale. The learned γ (scale) and β (shift) let the model decide whether the normalized values should be stretched, compressed, or offset for the next layer.

Key distinction from BatchNorm

LayerNorm normalizes across the feature dimension for each token independently. Batch Normalization (BatchNorm) instead computes each feature coordinate's statistics across examples in a mini-batch (and, for some data layouts, additional positions). This makes LayerNorm independent of batch size, which is useful for variable-length sequences and autoregressive generation.

PropertyBatchNormLayerNorm
Normalizes acrossMini-batch examples per feature coordinateFeature values within one token
Depends on batch sizeYesNo
Running statisticsYes (train/eval mismatch)No
Works for variable-length sequencesAwkwardNatural
Works for autoregressive decodingAwkwardNatural
Used in modern LLMsRareRoutine

Why is BatchNorm awkward for autoregressive decoding while LayerNorm is natural?

Answer

BatchNorm depends on batch-level statistics or stored running statistics, so behavior can change between training and inference. Autoregressive decoding often uses batch size 1 and variable sequence lengths. LayerNorm normalizes each token independently across its feature dimension, so population statistics are unnecessary.

Hold one token fixed and change only its neighbors. Predict which first output should stay unchanged, then compare it with a batch-normalized feature coordinate.

layernorm-is-per-token.py
1from math import sqrt 2 3def layer_norm(values, eps=1e-5): 4 mean = sum(values) / len(values) 5 variance = sum((value - mean) ** 2 for value in values) / len(values) 6 return [(value - mean) / sqrt(variance + eps) for value in values] 7 8def batch_norm_one_feature(values, eps=1e-5): 9 mean = sum(values) / len(values) 10 variance = sum((value - mean) ** 2 for value in values) / len(values) 11 return [(value - mean) / sqrt(variance + eps) for value in values] 12 13token = [1.0, 3.0, 5.0, 7.0] 14ln_before = [layer_norm(row) for row in [token, [2.0, 4.0, 6.0, 8.0]]][0] 15ln_after = [layer_norm(row) for row in [token, [100.0, 100.0, 100.0, 100.0]]][0] 16batch_a = batch_norm_one_feature([token[0], 2.0, 3.0]) 17batch_b = batch_norm_one_feature([token[0], 2.0, 100.0]) 18 19print("LayerNorm token:", [round(value, 3) for value in ln_before]) 20print("Batch feature with peers=2,3:", round(batch_a[0], 3)) 21print("Batch feature with peers=2,100:", round(batch_b[0], 3)) 22assert ln_before == ln_after 23assert round(batch_a[0], 3) != round(batch_b[0], 3)
Output
1LayerNorm token: [-1.342, -0.447, 0.447, 1.342] 2Batch feature with peers=2,3: -1.225 3Batch feature with peers=2,100: -0.718

The LayerNorm row is identical before and after the neighboring token changes. The batch-normalized feature moves from -1.225 to -0.718 because its peers changed. That's the dependency the table summarized.

PyTorch implementation

The hand function operated on one list. Now the PyTorch version accepts a tensor of shape (batch, seq_len, d_model), reduces only its final feature dimension, and broadcasts learned scale and shift over batch and sequence positions. The unbiased=False variance matches the formula above.

Before running it, predict whether the one-token tensor will match [0.45, -0.45, -1.34, 1.34] after rounding.

pytorch-implementation.py
1import torch 2import torch.nn as nn 3 4class LayerNorm(nn.Module): 5 """Layer Normalization from scratch.""" 6 def __init__(self, d_model: int, eps: float = 1e-5): 7 super().__init__() 8 self.gamma = nn.Parameter(torch.ones(d_model)) 9 self.beta = nn.Parameter(torch.zeros(d_model)) 10 self.eps = eps 11 12 def forward(self, x: torch.Tensor) -> torch.Tensor: 13 # x shape: (batch, seq_len, d_model) 14 mean = x.mean(dim=-1, keepdim=True) 15 var = x.var(dim=-1, keepdim=True, unbiased=False) 16 x_norm = (x - mean) / torch.sqrt(var + self.eps) 17 return self.gamma * x_norm + self.beta 18 19# Quick check with the hand-worked example 20x = torch.tensor([[[3.0, 1.0, -1.0, 5.0]]]) # shape (1, 1, 4) 21ln = LayerNorm(d_model=4) 22# Fix gamma=1, beta=0 to match the hand calculation 23nn.init.constant_(ln.gamma, 1.0) 24nn.init.constant_(ln.beta, 0.0) 25out = ln(x) 26print(out.round(decimals=2).detach()) 27expected = torch.tensor([[[0.45, -0.45, -1.34, 1.34]]]) 28assert torch.allclose(out.round(decimals=2), expected)
Output
1tensor([[[ 0.4500, -0.4500, -1.3400, 1.3400]]])
Parallel-coordinate plot tracing four token features from input values through mean centering and standard-deviation scaling, with mean lines and one-standard-deviation bands at each stage.
Subtracting μ = 2 shifts every feature by the same amount, moving the mean to zero without changing σ ≈ 2.24. Dividing by σ contracts the spread to one while preserving feature order. With γ = 1 and β = 0, the normalized vector is also the output.

The colored paths preserve each feature's identity across both operations. Centering translates the whole vector; scaling changes its spread. Learned γ and β can then reshape that standardized vector for the next layer.


Where you put LayerNorm matters

The arithmetic is now concrete. The live question is where LayerNorm sits inside a transformer block: before the residual add, or after it. That choice decides which path the next block and the backward pass see.

One block, two placements

There are two placements to compare. Post-LN normalizes after the residual update has been added, so its residual path crosses a normalization operation in every block. Pre-LN normalizes before the sublayer update, so its residual path keeps an identity contribution to the backward Jacobian.

Post-LN: the original transformer

In the original 2017 Transformer and BERT-style encoders, normalization happens after the residual addition:[3]Reference 3Attention Is All You Need.https://arxiv.org/abs/1706.03762[4]Reference 4BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.https://arxiv.org/abs/1810.04805

xl+1=LayerNorm(xl+Sublayer(xl))x_{l+1} = \text{LayerNorm}(x_l + \text{Sublayer}(x_l))xl+1​=LayerNorm(xl​+Sublayer(xl​))

In plain terms: compute the sublayer output (attention or FFN), add it to the original input along the skip connection, then normalize the sum. Normalization now sits directly on the main highway, so gradients through the residual path pass through a LayerNorm Jacobian in each Post-LN block.

Pre-LN: the GPT-2 layout

GPT-2 uses the other arrangement: normalization happens before the sublayer.[5]Reference 5Language Models are Unsupervised Multitask Learners.https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf

xl+1=xl+Sublayer(LayerNorm(xl))x_{l+1} = x_l + \text{Sublayer}(\text{LayerNorm}(x_l))xl+1​=xl​+Sublayer(LayerNorm(xl​))

Normalize the input first, feed it into the sublayer, then add the result to a direct copy of the input. The residual path from x_l to x_{l+1} keeps an identity contribution to the backward pass.

At the full-model level, GPT-2 and the Pre-LN architecture analyzed by Xiong et al. apply a final normalization before prediction.[5]Reference 5Language Models are Unsupervised Multitask Learners.https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf[1]Reference 1On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745 The final operation sets the scale of the accumulated residual stream before the output projection. Dropping it produces a different architecture, so its output scale needs a fresh check. Post-LN doesn't need that extra layer: every block already ends with LayerNorm, so its last hidden state is normalized.

The two full-model paths are easier to compare when the final normalization is visible. The diagram keeps the block formulas from the equations above and shows where the model reaches its logits.

Diagram showing Pre-LN model, Post-LN model, embed, and each block: x + F(LN(x)).
Pre-LN model, Post-LN model, embed, and each block: x + F(LN(x)).

The block graphs below show why that skip path changes the backward Jacobian.

Side-by-side Pre-LN and Post-LN block graphs with their exact backward Jacobians, marking the additive identity route in Pre-LN and the repeated LayerNorm Jacobian on the Post-LN residual path.
Pre-LN gives each block Jacobian an additive identity term, $I + J_FJ_{LN}$. Post-LN left-multiplies the residual contribution by $J_{LN}$, so the depth-wide route repeatedly crosses normalization Jacobians.

The left graph is Pre-LN: LayerNorm cleans the sublayer input, and the raw input still jumps to the add. The right graph is Post-LN: the sublayer and the skip meet first, then LayerNorm cleans the sum.

In one equation, what is the architectural difference between Post-LN and Pre-LN?

Answer

Post-LN is xl+1=LayerNorm(xl+F(xl))x_{l+1} = LayerNorm(x_l + F(x_l))xl+1​=LayerNorm(xl​+F(xl​)). Pre-LN is xl+1=xl+F(LayerNorm(xl))x_{l+1} = x_l + F(LayerNorm(x_l))xl+1​=xl​+F(LayerNorm(xl​)). The important difference is whether normalization sits on the residual highway after the add, or only on the sublayer input before the add.

Hold a sublayer update fixed to isolate placement. Predict which stream has mean zero before looking at the output: Post-LN centers and rescales the updated stream immediately, while Pre-LN lets the updated stream continue along the residual path.

pre-post-forward-placement.py
1from math import sqrt 2 3def layer_norm(values, eps=1e-5): 4 mean = sum(values) / len(values) 5 variance = sum((value - mean) ** 2 for value in values) / len(values) 6 return [(value - mean) / sqrt(variance + eps) for value in values] 7 8x = [1.0, 2.0, 4.0, 8.0] 9fixed_update = [0.5, -0.5, 1.0, -1.0] 10pre_ln_output = [value + update for value, update in zip(x, fixed_update)] 11post_ln_output = layer_norm(pre_ln_output) 12 13print("Pre-LN stream mean:", round(sum(pre_ln_output) / len(pre_ln_output), 3)) 14print("Post-LN stream mean:", round(sum(post_ln_output) / len(post_ln_output), 3)) 15print("Post-LN output:", [round(value, 3) for value in post_ln_output]) 16assert abs(sum(post_ln_output) / len(post_ln_output)) < 1e-12 17assert sum(pre_ln_output) / len(pre_ln_output) != 0.0
Output
1Pre-LN stream mean: 3.75 2Post-LN stream mean: 0.0 3Post-LN output: [-0.954, -0.954, 0.53, 1.378]

A two-layer walkthrough

The forward probe showed one update. Now repeat it twice. Assume each sublayer writes the same 0.5 to every feature, with γ = 1 and β = 0.

The simplification keeps trained attention and feed-forward details out of this pass; those updates are usually nonuniform. Before reading each stack, predict which one preserves the shared offset.

Post-LN stack

  • Start: x_0 = [1.0, 1.0, 1.0, 1.0]
  • After block 1: x_1 = LayerNorm(x_0 + 0.5) = LayerNorm([1.5, 1.5, 1.5, 1.5]) = [0.0, 0.0, 0.0, 0.0] (mean-centered, then scaled)
  • After block 2: x_2 = LayerNorm(x_1 + 0.5) = LayerNorm([0.5, 0.5, 0.5, 0.5]) = [0.0, 0.0, 0.0, 0.0]

The constant offset is reset at every layer. LayerNorm of a constant vector has variance zero, so with γ = 1 and β = 0 the output is the zero vector.

The residual stream therefore doesn't preserve that shared offset in this toy.

Pre-LN stack

  • Start: x_0 = [1.0, 1.0, 1.0, 1.0]
  • After block 1: x_1 = x_0 + Sublayer(LayerNorm(x_0)) = [1.0, 1.0, 1.0, 1.0] + 0.5 = [1.5, 1.5, 1.5, 1.5]
  • After block 2: x_2 = x_1 + Sublayer(LayerNorm(x_1)) = [1.5, 1.5, 1.5, 1.5] + 0.5 = [2.0, 2.0, 2.0, 2.0]

The running sum flows through the network because Pre-LN leaves the residual path outside block normalization. That's also why a Pre-LN stack still wants a final norm before the head: the stream was never reset at block boundaries.

The backward equation below shows why this placement includes an identity gradient route.

Why does the Pre-LN architecture analyzed by Xiong et al. include a final normalization before the output head?

Answer

Inside each Pre-LN block, the residual stream bypasses block-output normalization and accumulates updates. A final norm sets its scale before the vocabulary projection. Removing that norm changes the architecture and requires rechecking output-scale and training behavior.

Now let the sublayer depend on the normalized or raw input instead of returning a fixed offset. The next cell applies two steps of each layout and asserts the same contrast: Pre-LN's stream grows, while Post-LN's output remains centered. Predict which assertion would fail if the two placements were accidentally swapped.

pre-ln-stack.py
1from math import sqrt 2 3def layer_norm(values, eps=1e-5): 4 mean = sum(values) / len(values) 5 variance = sum((value - mean) ** 2 for value in values) / len(values) 6 scale = sqrt(variance + eps) 7 return [(value - mean) / scale for value in values] 8 9def sublayer(values): 10 # Tiny deterministic stand-in for attention or an FFN update. 11 return [0.2 * value + 0.1 for value in values] 12 13def pre_ln_step(values): 14 return [value + update for value, update in zip(values, sublayer(layer_norm(values)))] 15 16def post_ln_step(values): 17 raw = [value + update for value, update in zip(values, sublayer(values))] 18 return layer_norm(raw) 19 20start = [1.0, 2.0, 4.0, 8.0] 21pre_after_two = pre_ln_step(pre_ln_step(start)) 22post_after_two = post_ln_step(post_ln_step(start)) 23 24assert sum(abs(value) for value in pre_after_two) > sum(abs(value) for value in start) 25assert abs(sum(post_after_two) / len(post_after_two)) < 1e-12

Why Pre-LN can train more stably

The forward path explains what gets normalized. Training adds another question: how does a loss signal travel from xl+1x_{l+1}xl+1​ back to xlx_lxl​? In Post-LN it crosses the block's normalization Jacobian. In Pre-LN it also has a direct identity route. At depth, multiplying those block Jacobians can change which layers learn fastest.

Post-LN gradient problem

Xiong et al. show that Post-LN produces especially large expected gradients near the output layers at initialization under their analysis.[1]Reference 1On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745 For one block, ask where the residual contribution appears after differentiation. The backward Jacobian is:

∂xl+1∂xl=JLN,l⋅(I+JF,l)\frac{\partial x_{l+1}}{\partial x_l} = J_{\text{LN},l} \cdot (I + J_{F,l})∂xl​∂xl+1​​=JLN,l​⋅(I+JF,l​)

Here J_LN,l is the LayerNorm Jacobian and J_F,l is the sublayer Jacobian. Across many layers, the backward pass multiplies terms like this together. In practice that means:

  • Gradient scale becomes uneven across depth. Early layers may receive tiny updates while late layers receive huge ones.
  • Top layers can receive very large updates at initialization. A single step with a moderately high learning rate can push weights far from a good basin.
  • Learning-rate warmup addresses this risk in the evaluated Post-LN setups. Start the learning rate near zero and ramp it up so early updates are smaller.

Pre-LN gradient advantage

For Pre-LN, the same differentiation puts the identity term first:

∂xl+1∂xl=I+JF,l⋅JLN,l\frac{\partial x_{l+1}}{\partial x_l} = I + J_{F,l} \cdot J_{\text{LN},l}∂xl​∂xl+1​​=I+JF,l​⋅JLN,l​

Every block includes a direct residual contribution to the gradient, so the backward signal doesn't rely entirely on repeated normalization Jacobians. Xiong et al. report that, in their evaluated tasks:

  • Pre-LN trains without learning-rate warmup in experiments where Post-LN needed it for stable optimization.[1]Reference 1On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745
  • Its initialization-time gradients are better behaved in their analysis.

Use a scalar product to isolate that route. It isn't a training experiment: j_layer_norm and j_sublayer stand in for Jacobians, and the cell compares their repeated products. Predict which twelve-block product falls below one before running it.

jacobian-route-toy.py
1depth = 12 2j_layer_norm = 0.75 3j_sublayer = 0.10 4 5post_block = j_layer_norm * (1.0 + j_sublayer) 6pre_block = 1.0 + j_sublayer * j_layer_norm 7post_path = post_block ** depth 8pre_path = pre_block ** depth 9 10print("one block: post=", round(post_block, 3), "pre=", round(pre_block, 3)) 11print("twelve-block product: post=", round(post_path, 3), "pre=", round(pre_path, 3)) 12assert post_path < 1.0 13assert pre_path > 1.0
Output
1one block: post= 0.825 pre= 1.075 2twelve-block product: post= 0.099 pre= 2.382

The Post-LN factor is 0.825, so its repeated product shrinks to 0.099. Pre-LN's factor is 1.075, and its product grows to 2.382. Real Jacobians are matrices, not fixed scalars, but the toy makes the placement effect visible.


A reported Pre-LN risk: representation collapse

Pre-LN isn't free of trade-offs. Analyses including ResiDual report that its residual stream can dominate newer sublayer updates as depth increases. Under that behavior, late layers make smaller relative changes to the hidden representation.[6]Reference 6ResiDual: Transformer with Dual Residual Connections.https://arxiv.org/abs/2304.14802

The ResiDual authors call this reported behavior representation collapse.[6]Reference 6ResiDual: Transformer with Dual Residual Connections.https://arxiv.org/abs/2304.14802 Treat it as a design risk to measure, not a guarantee for every Pre-LN model or training run.

Keep the two trade-offs side by side:

PropertyPost-LNPre-LN
Warmup result in Xiong et al.Needed in evaluated stable runsRemoved in evaluated stable runs
Initialization gradients in Xiong et al.Large near output layersBetter behaved
Reported deep-layer riskGradient flow can become difficultResidual stream can dominate later updates
Representative architectureOriginal Transformer, BERTGPT-2

There's no universal winner. Xiong et al. demonstrate the Pre-LN optimization advantage in their evaluated tasks. DeepNorm later scales a Post-LN-derived design to a 1,000-layer machine-translation experiment by combining residual scaling with matched initialization.[1]Reference 1On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745[7]Reference 7DeepNet: Scaling Transformers to 1,000 Layers.https://arxiv.org/abs/2203.00555

What is the representation-collapse risk reported for Pre-LN?

Answer

The residual stream can dominate newer sublayer updates as depth grows. Late blocks still run, but their relative changes may become small compared with the accumulated stream. This is an observed and analyzed risk, rather than an inevitable outcome for every Pre-LN run.

Residual magnitude growth across depth

Use a fixed write to make the geometry visible: every block adds [1, 0, 0, 0]. Pre-LN lets the stream accumulate, so that write becomes a smaller fraction of ||x||. Post-LN re-normalizes after each add, so its stream magnitude stays fixed.

DeepNorm is still Post-LN-shaped after the add (LN(α x + G(x))), so its after-LN ||x|| matches Post-LN in this sketch. α changes the mix before LayerNorm by up-scaling the skip (α = (2N)^{1/4} = 2 for this 8-layer encoder-style sketch). The paper's other half, shrinking residual-branch initialization by β, needs weights and isn't represented in this table.

Predict how the Pre-LN write share and the three stream magnitudes will change with depth before running the cell.

residual-stream-magnitude-vs-depth.py
1from math import sqrt 2 3def l2(values: list[float]) -> float: 4 return sqrt(sum(v * v for v in values)) 5 6def layer_norm(values: list[float], eps: float = 1e-5) -> list[float]: 7 mean = sum(values) / len(values) 8 var = sum((v - mean) ** 2 for v in values) / len(values) 9 scale = sqrt(var + eps) 10 return [(v - mean) / scale for v in values] 11 12depth = 8 13update = [1.0, 0.0, 0.0, 0.0] 14alpha = (2 * depth) ** 0.25 # encoder-only residual scale from DeepNet 15x_pre = [1.0, 0.0, 0.0, 0.0] 16x_post = [1.0, 0.0, 0.0, 0.0] 17x_deep = [1.0, 0.0, 0.0, 0.0] 18 19print(f"alpha={alpha:.3f}") 20print("layer | Pre ||x|| | Pre write share | Post ||x|| | Deep ||αx|| | Deep ||x||") 21for layer in range(1, depth + 1): 22 pre_share = l2(update) / (l2(x_pre) + l2(update)) 23 x_pre = [a + b for a, b in zip(x_pre, update)] 24 x_post = layer_norm([a + b for a, b in zip(x_post, update)]) 25 scaled = [alpha * a for a in x_deep] 26 x_deep = layer_norm([a + b for a, b in zip(scaled, update)]) 27 print( 28 f"{layer:5d} | {l2(x_pre):8.3f} | {pre_share:15.3f} | " 29 f"{l2(x_post):10.3f} | {l2(scaled):9.3f} | {l2(x_deep):9.3f}" 30 ) 31 32assert l2(x_pre) > l2(x_post) 33assert abs(l2(x_post) - l2(x_deep)) < 1e-5 34assert alpha == 2.0
Output
1alpha=2.000 2layer | Pre ||x|| | Pre write share | Post ||x|| | Deep ||αx|| | Deep ||x|| 3 1 | 2.000 | 0.500 | 2.000 | 2.000 | 2.000 4 2 | 3.000 | 0.333 | 2.000 | 4.000 | 2.000 5 3 | 4.000 | 0.250 | 2.000 | 4.000 | 2.000 6 4 | 5.000 | 0.200 | 2.000 | 4.000 | 2.000 7 5 | 6.000 | 0.167 | 2.000 | 4.000 | 2.000 8 6 | 7.000 | 0.143 | 2.000 | 4.000 | 2.000 9 7 | 8.000 | 0.125 | 2.000 | 4.000 | 2.000 10 8 | 9.000 | 0.111 | 2.000 | 4.000 | 2.000

Pre-LN ||x|| grows from 2 to 9, while the fixed write's share falls from 0.5 to 0.11. Later blocks still run, but each write accounts for less of the accumulated stream. That's the geometry behind the reported representation-collapse risk.

Post-LN resets its magnitude every block, and DeepNorm matches it after LayerNorm (||x|| ≈ 2 here because a 4-d unit-variance vector has L2 norm d\sqrt{d}d​). α appears before LayerNorm: from layer 2 onward ||αx|| = 4 while the write remains 1. The paper then shrinks residual-branch initialization by β so the write doesn't explode at step one.

Real models use learned sublayers rather than a fixed update. Measure ||x_l|| and relative update size on your own stack before calling this behavior collapse.

DeepNorm: scaled residual + matched init

DeepNorm (DeepNet) keeps a Post-LN-shaped block but multiplies the residual branch by a depth-dependent constant α. It also down-scales residual-branch initialization so expected update magnitudes stay controlled at large depth.[7]Reference 7DeepNet: Scaling Transformers to 1,000 Layers.https://arxiv.org/abs/2203.00555 The paper's encoder update has the form:

xl+1=LN(α xl+Gl(xl))x_{l+1} = \mathrm{LN}\big(\alpha\, x_l + G_l(x_l)\big)xl+1​=LN(αxl​+Gl​(xl​))

Here GlG_lGl​ is the sublayer stack (attention and/or FFN writes). α grows slowly with depth so the residual highway still carries signal, while initialization of GlG_lGl​ is reduced. The paper pairs α with a gain β on residual-branch weights: FFN, value projection, and output projection use it, while query and key keep standard Xavier gain. Encoder-only uses α=(2N)1/4\alpha=(2N)^{1/4}α=(2N)1/4 and β=(8N)−1/4\beta=(8N)^{-1/4}β=(8N)−1/4 for NNN layers. Decoder-only uses the same pair with MMM in place of NNN; encoder-decoder stacks use different constants on each side. Choose the row that matches your architecture.

DeepNorm is therefore not "plain Post-LN with no residual scaling." The residual scale and matched init are the design, and the 1,000-layer machine-translation experiment is evidence under that recipe, not a free upgrade for an arbitrary Post-LN codebase.

What two mechanisms does DeepNorm add on top of a Post-LN-shaped block?

Answer

A depth-dependent residual scale α in the update xl+1=LN(αxl+Gl(xl))x_{l+1}=\mathrm{LN}(\alpha x_l + G_l(x_l))xl+1​=LN(αxl​+Gl​(xl​)), and matched down-scaling of residual-branch initialization so expected update size stays controlled as depth grows.


RMSNorm: scale-only normalization

Placement decides where the normalizer sits. Many decoder-only LLMs also change the statistic inside that box.

RMSNorm (Root Mean Square Layer Normalization) replaces LayerNorm's centered standard deviation with a root-mean-square scale. Zhang and Sennrich propose it as a cheaper normalization that keeps rescaling invariance while dropping mean-centering. Their evaluated models match LayerNorm quality, with the speedup depending on the stack.[8]Reference 8Root Mean Square Layer Normalization.https://arxiv.org/abs/1910.07467 The Qwen2.5 paper uses RMSNorm with pre-normalization. The Gemma 2 paper uses RMSNorm on both the input and the output of each sublayer.[9]Reference 9Qwen2.5 Technical Reporthttps://arxiv.org/abs/2412.15115[10]Reference 10Gemma 2: Improving Open Language Models at a Practical Sizehttps://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf Copy those layouts only after checking the checkpoint you actually serve.

RMSNorm(x)=γ⊙x1d∑i=1dxi2+ϵ\text{RMSNorm}(x) = \gamma \odot \frac{x}{\sqrt{\frac{1}{d}\sum_{i=1}^d x_i^2 + \epsilon}}RMSNorm(x)=γ⊙d1​∑i=1d​xi2​+ϵ​x​

Instead of centering (subtracting the mean) and then scaling (dividing by standard deviation), RMSNorm just divides by the root-mean-square. It doesn't subtract the mean, and many implementations also omit a learned bias term.

StepLayerNormRMSNorm
Compute meanYesNo
Subtract meanYesNo
Scale statisticCentered standard deviationRoot mean square
Divide by scale statisticYesYes
Learned scale γYesYes
Learned shift βYesOften omitted
Exact comparison of LayerNorm and RMSNorm on x equals 3, 1, negative 1, 5, showing LayerNorm output centered at zero, RMSNorm output with nonzero mean, and their different response to adding a constant offset.
With $\gamma = 1$, $\beta = 0$, and $\epsilon$ ignored, LayerNorm maps the token to mean zero and standard deviation one. RMSNorm maps it to RMS one but keeps mean $0.67$; adding a shared offset leaves LayerNorm unchanged but changes the RMSNorm output.

The original RMSNorm paper reports comparable quality to LayerNorm on its evaluated tasks and observed run-time reductions from 7% to 64% across different models and implementations.[8]Reference 8Root Mean Square Layer Normalization.https://arxiv.org/abs/1910.07467 Fused kernels, hardware, and the fraction of time spent in normalization determine the result in another stack, so those measurements are evidence from the paper rather than a promised speedup.

What does RMSNorm remove from LayerNorm, and what does it keep?

Answer

RMSNorm removes mean-centering and often removes the learned bias. It keeps per-token rescaling by a feature-wise magnitude estimate and keeps a learned scale vector.

The hand example already gives us a useful test. RMSNorm should keep the vector's RMS near one, but its mean needn't be zero. Predict both properties before running the list implementation.

rmsnorm-a-common-modern-shortcut.py
1from math import sqrt 2 3def rms_norm(values, weight=None, eps=1e-6): 4 if weight is None: 5 weight = [1.0] * len(values) 6 rms = sqrt(sum(value * value for value in values) / len(values) + eps) 7 return [value * scale / rms for value, scale in zip(values, weight)] 8 9values = [3.0, 1.0, -1.0, 5.0] 10normalized = rms_norm(values) 11 12assert round(sqrt(sum(value * value for value in normalized) / len(normalized)), 6) == 1.0 13assert round(sum(normalized) / len(normalized), 6) != 0.0 14assert rms_norm([3.0, 1.0], weight=[0.0, 0.0]) == [0.0, 0.0]

Now apply the same formula to a (batch, seq_len, d_model) tensor. The PyTorch module has one learned weight vector and no bias. Predict why its output mean is nonzero even though its feature RMS is one.

rmsnorm-a-common-modern-shortcut-2.py
1import torch 2import torch.nn as nn 3 4class RMSNorm(nn.Module): 5 """Root Mean Square Layer Normalization.""" 6 def __init__(self, d_model: int, eps: float = 1e-6): 7 super().__init__() 8 self.weight = nn.Parameter(torch.ones(d_model)) 9 self.eps = eps 10 11 def forward(self, x: torch.Tensor) -> torch.Tensor: 12 # x shape: (batch, seq_len, d_model) 13 rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) 14 return self.weight * (x / rms) 15 16# Quick sanity check 17x = torch.tensor([[[3.0, 1.0, -1.0, 5.0]]]) 18rms = RMSNorm(d_model=4) 19out = rms(x) 20print(out.round(decimals=4).detach()) 21feature_rms = torch.sqrt(out.pow(2).mean(dim=-1)) 22assert torch.allclose(feature_rms, torch.ones_like(feature_rms), atol=1e-5) 23assert not torch.allclose(out.mean(dim=-1), torch.zeros_like(out.mean(dim=-1)))
Output
1tensor([[[ 1.0000, 0.3333, -0.3333, 1.6667]]])

The output has RMS one and mean about 0.67, matching the list calculation. The removed centering step is doing visible work, not only saving an operation.


Reported normalization layouts

You can now separate two choices: where normalization sits around the residual add, and which statistic it computes. A third family targets attention logits rather than the residual stream. The table records reported examples; their experiments differ in model size, objective, and optimization setup.

RecipeReported exampleWhat changes
Post-LNBERTNormalize after the residual addition
Pre-LNGPT-2LayerNorm before attention and FFN, keeping an identity residual contribution
Pre-RMSNormQwen2.5Same Pre-LN residual layout, with RMSNorm in place of LayerNorm[9]Reference 9Qwen2.5 Technical Reporthttps://arxiv.org/abs/2412.15115
Pre + Post normGemma 2RMSNorm on both the input and the output of each sublayer[10]Reference 10Gemma 2: Improving Open Language Models at a Practical Sizehttps://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf
DeepNormDeepNet studyPost-LN-shaped update xl+1=LN(αxl+Gl(xl))x_{l+1}=\mathrm{LN}(\alpha x_l + G_l(x_l))xl+1​=LN(αxl​+Gl​(xl​)) with depth-dependent α and matched residual-branch init; evaluated up to 1,000 layers[7]Reference 7DeepNet: Scaling Transformers to 1,000 Layers.https://arxiv.org/abs/2203.00555
QK-Norm (Query-Key Normalization)Rybakov et al. studyNormalize query and key vectors before their dot product; evaluated alone and with softmax capping[11]Reference 11QK-Norm: Improving Transformer Training Stability.https://arxiv.org/abs/2410.16682

No row is a global winner. Peri-LN names an input-and-output normalization layout and reports experiments up to 3.2B parameters.[12]Reference 12Peri-LN: Revisiting Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2502.02732 Rybakov et al. evaluate attention-specific normalization and capping on an 830M-parameter model driven into instability with high learning rates.[11]Reference 11QK-Norm: Improving Transformer Training Stability.https://arxiv.org/abs/2410.16682

QK-Norm targets attention logits, not the residual stream. Gemma 2 uses a related but different control: logit soft-capping, soft_cap * tanh(logits / soft_cap), on attention logits and the final layer, together with pre+post RMSNorm. That clamps logit range, but it doesn't rescale Q and K.[10]Reference 10Gemma 2: Improving Open Language Models at a Practical Sizehttps://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf

How does QK-Norm reduce attention instability?

Answer

QK-Norm normalizes query and key vectors before their dot product, limiting the contribution from their vector magnitude to attention logits. Rybakov et al. study a LayerNorm-based form and report that combining QK-Norm with softmax capping increases the stable learning-rate range in their evaluated setup. The normalizer and its exact placement relative to positional operations remain architecture choices.

Test the boundary directly. Multiplying a query by twenty changes its raw dot-product logit. RMS-rescaling both query and key should remove that scale-only change while leaving their direction intact. Predict the two pairs of logits before running the cell.

qk-norm-controls-dot-scale.py
1from math import isclose, sqrt 2 3def rms_rescale(values): 4 rms = sqrt(sum(value * value for value in values) / len(values)) 5 return [value / rms for value in values] 6 7def attention_logit(query, key): 8 return sum(q * k for q, k in zip(query, key)) / sqrt(len(query)) 9 10query = [6.0, -3.0, 2.0, 1.0] 11key = [4.0, -2.0, 1.0, 3.0] 12large_query = [20.0 * value for value in query] 13 14raw = attention_logit(query, key) 15raw_large = attention_logit(large_query, key) 16controlled = attention_logit(rms_rescale(query), rms_rescale(key)) 17controlled_large = attention_logit(rms_rescale(large_query), rms_rescale(key)) 18 19print("raw logits:", round(raw, 3), round(raw_large, 3)) 20print("RMS-rescaled logits:", round(controlled, 3), round(controlled_large, 3)) 21assert raw_large == 20.0 * raw 22assert isclose(controlled, controlled_large, rel_tol=1e-12)
Output
1raw logits: 17.5 350.0 2RMS-rescaled logits: 1.807 1.807

The raw logit grows from 17.5 to 350.0; RMS rescaling holds both cases at 1.807. That controls query-key magnitude only. It doesn't repair a residual stream whose norm or update share is drifting.

Why do hybrid layouts such as Peri-LN or Pre + Post norm exist if Pre-LN is already stable?

Answer

Pre-LN includes a direct identity gradient route, but residual-stream scale can still become a concern. Hybrid layouts add output-side normalization while preserving an input-side norm, then test whether that changes activation and optimization behavior.

Peri-LN and hybrids

Return to Peri-LN, the input-and-output layout in the table. Kim et al. report more balanced variance growth and steadier gradients than their compared layouts in experiments up to 3.2B parameters.[12]Reference 12Peri-LN: Revisiting Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2502.02732 Gemma 2's technical report describes the same input-and-output pattern with RMSNorm without using the Peri-LN name.[10]Reference 10Gemma 2: Improving Open Language Models at a Practical Sizehttps://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf These controls belong around sublayers; QK-Norm belongs inside attention's dot-product path.


Debug normalization failures

When a real stack fails, start with the path that is misbehaving. Log residual-stream norms, attention-logit ranges, layerwise gradient norms, the final normalization, and the learning-rate schedule before changing several knobs at once.

SymptomLikely causeFix
Model diverges at step 1 with loss = NaNFor a Post-LN run, an overly large early update is one hypothesis; Xiong et al. connect Post-LN initialization gradients to warmup sensitivity.Inspect layerwise gradient norms and learning-rate schedule. Compare a warmer start, smaller initial rate, or a Pre-LN variant under the same setup.
One normalization placement underperforms anotherPlacement interacts with depth, initialization, optimizer, data, and objective.Treat placement as an ablation, then compare training loss, gradient profile, and downstream quality under matched conditions.
A Pre-LN implementation differs sharply from a reference modelA final normalization used by the target Pre-LN architecture may be absent.Check the architecture definition and restore its final LayerNorm or RMSNorm before the output projection when applicable.[1]Reference 1On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745
RMSNorm is described as a faster spelling of LayerNormMean-centering and RMS rescaling were collapsed into one operation.Compare the formulas: LayerNorm centers and rescales; RMSNorm rescales without subtracting the mean.
Attention becomes nearly one-hot during unstable trainingQuery-key dot products may have grown enough to saturate softmax.Inspect attention-logit statistics; QK-Norm or softmax capping are targeted candidates to evaluate in such runs.
Attention logits are controlled but another instability remainsQK-Norm was treated as a replacement for residual normalization.Measure query-key scale and residual-stream dynamics separately, then target each failure directly.

Practice: build a toggle-switch transformer block

Make the placement choice executable. One block can switch between Post-LN and Pre-LN, and a small stack can report the gradient norm of each sublayer at initialization. The toy uses a linear layer and ReLU instead of attention, so it tests the residual route rather than a full language model.

Predict whether the six gradient norms in each row will be identical before running the cell. Then vary depth and seed.

practice-build-a-toggle-switch-transformer.py
1import torch 2import torch.nn as nn 3 4class ToggleBlock(nn.Module): 5 """Transformer block with style='pre' or style='post'.""" 6 def __init__(self, d_model: int, style: str = "pre"): 7 super().__init__() 8 assert style in ("pre", "post") 9 self.style = style 10 self.norm = nn.LayerNorm(d_model) 11 # Simplified sublayer: single linear + ReLU 12 self.sublayer = nn.Sequential( 13 nn.Linear(d_model, d_model * 4), 14 nn.ReLU(), 15 nn.Linear(d_model * 4, d_model), 16 ) 17 18 def forward(self, x: torch.Tensor) -> torch.Tensor: 19 if self.style == "pre": 20 # x_l + Sublayer(LayerNorm(x_l)) 21 return x + self.sublayer(self.norm(x)) 22 else: 23 # LayerNorm(x_l + Sublayer(x_l)) 24 return self.norm(x + self.sublayer(x)) 25 26class TinyStack(nn.Module): 27 def __init__(self, d_model: int = 64, depth: int = 6, style: str = "pre"): 28 super().__init__() 29 self.blocks = nn.ModuleList([ToggleBlock(d_model, style) for _ in range(depth)]) 30 self.final_norm = nn.LayerNorm(d_model) if style == "pre" else nn.Identity() 31 self.head = nn.Linear(d_model, 16) 32 33 def forward(self, x: torch.Tensor) -> torch.Tensor: 34 for block in self.blocks: 35 x = block(x) 36 return self.head(self.final_norm(x)) 37 38def inspect_grad(style: str, d_model: int = 64, depth: int = 6, seed: int = 0): 39 torch.manual_seed(seed) 40 model = TinyStack(d_model=d_model, depth=depth, style=style) 41 x = torch.randn(2, 8, d_model) 42 target = torch.randn(2, 8, 16) 43 out = model(x) 44 loss = nn.functional.mse_loss(out, target) 45 loss.backward() 46 grad_norms = [block.sublayer[0].weight.grad.norm().item() for block in model.blocks] 47 rounded_norms = [round(value, 4) for value in grad_norms] 48 print(f"{style:4} | loss={loss.item():.4f} | block grad norms={rounded_norms}") 49 return grad_norms 50 51pre_grad_norms = inspect_grad("pre") 52post_grad_norms = inspect_grad("post") 53 54assert len(pre_grad_norms) == len(post_grad_norms) == 6 55assert all(value > 0 for value in pre_grad_norms + post_grad_norms)
Output
1pre | loss=1.3247 | block grad norms=[0.2664, 0.2565, 0.2494, 0.2392, 0.2472, 0.2415] 2post | loss=1.3183 | block grad norms=[0.2663, 0.2614, 0.2619, 0.2555, 0.2684, 0.2698]

The profiles aren't identical, even in a six-block toy. That variation is a measurement to inspect, not proof that one layout wins: Xiong et al.'s result concerns their analyzed initialization and evaluated tasks. Use the same logging pattern when a deeper run diverges, then change one architectural or optimization variable at a time.

Mini exercise

  1. Change depth from 6 to 12. How do the per-block gradient profiles change?
  2. Replace final_norm with nn.Identity() in Pre-LN mode and inspect output scale and loss under several seeds.
  3. Try replacing nn.LayerNorm with the RMSNorm class from earlier. Does the gradient norm change meaningfully?

For feedback, compare ranges rather than a single number. Depth changes how many Jacobians are composed; removing the final norm changes the scale reaching the head; switching to RMSNorm changes the statistic but not the residual placement.


Mastery check

  • Compute LayerNorm by hand, including mean, population variance, epsilon, learned scale, and learned shift.
  • Explain why LayerNorm uses one token's feature values while BatchNorm uses statistics across examples or positions.
  • Identify Pre-LN and Post-LN from code, then derive I+JFJLNI + J_FJ_{LN}I+JF​JLN​ versus JLN(I+JF)J_{LN}(I + J_F)JLN​(I+JF​).
  • Connect Post-LN's initialization-time gradient profile to learning-rate warmup without claiming the result applies to every training setup.
  • Explain the reported deep Pre-LN representation-collapse risk and name measurements that would reveal it.
  • Implement RMSNorm and state exactly which LayerNorm operations it removes.
  • Separate residual-stream normalization, QK-Norm, DeepNorm, and input-output hybrid layouts by the mechanism each changes.

Practice checkpoints

Your 48-layer decoder diverges during first few hundred steps, but the same optimizer works on a 12-layer version. What should you inspect first?

Answer

Start with norm placement and residual-path statistics. If the stack is Post-LN, inspect initialization-time layerwise gradients and check whether warmup or initial learning rate differs from a stable setup. If the model follows a Pre-LN architecture with a final norm, confirm that norm remains present; then inspect attention-logit statistics and candidate controls such as QK-Norm or softmax capping.

Why should you avoid treating Pre-LN as a universal replacement for Post-LN?

Answer

Published comparisons are tied to architecture, depth, initialization, schedule, and task. Xiong et al. establish a Pre-LN initialization and warmup advantage in evaluated setups, while DeepNorm demonstrates a stable Post-LN-derived alternative under a different design. Measure the candidate layout under matched conditions.

When does RMSNorm change the model's behavior less than replacing Pre-LN with Post-LN?

Answer

RMSNorm changes the normalization formula inside a broadly similar residual layout. Switching Pre-LN to Post-LN changes where normalization sits relative to the skip path, which directly changes gradient flow and optimization stability. So RMSNorm is a smaller architectural change than moving the norm across the residual addition.

Complete the lesson

Mastery Check

Answer every question, then check your score. Score 75% or higher to mark this lesson complete.

1.For x = [3, 1, -1, 5], with gamma = [1, 1, 1, 1], beta = [0, 0, 0, 0], and epsilon ignored, what does LayerNorm output?

Correct answer: [0.45, -0.45, -1.34, 1.34], because subtracting mean 2 and dividing by sqrt(5) centers and rescales the token.

The token mean is 2, so the centered vector is [1, -1, -3, 3]. The variance is (1 + 1 + 9 + 9) / 4 = 5, so dividing by sqrt(5) gives approximately [0.45, -0.45, -1.34, 1.34]. Gamma of 1 and beta of 0 leave that normalized vector unchanged.

2.Why is BatchNorm awkward for autoregressive decoding while LayerNorm is natural?

Correct answer: BatchNorm uses mini-batch or running population statistics, so batch-size-one and variable-length decoding can differ from training; LayerNorm uses each token's feature values.

BatchNorm's statistics depend on other examples in a mini-batch or on stored running estimates, which creates train/eval and batch-size sensitivity. LayerNorm computes statistics within each token's feature vector, so it behaves naturally for variable-length sequences and autoregressive generation with batch size one.

3.Which pair of block equations correctly distinguishes Post-LN from Pre-LN and its gradient route?

Correct answer: Post-LN: x_{l+1} = LayerNorm(x_l + F(x_l)); Pre-LN: x_{l+1} = x_l + F(LayerNorm(x_l)). Pre-LN keeps an identity contribution on the residual path.

Post-LN normalizes after the residual update, so the skip path passes through a LayerNorm Jacobian. Pre-LN normalizes only the sublayer input and then adds the result back to the original x_l, giving the backward Jacobian a direct identity term plus the sublayer-through-norm term.

4.A 72-layer Post-LN transformer diverges in the first few updates. Gradient logging at initialization shows much larger norms in top blocks than lower blocks. Which intervention directly addresses this failure mode while preserving diagnostics?

Correct answer: Start with learning-rate warmup or a smaller early rate, then keep checking layerwise gradient norms.

Post-LN places LayerNorm on the residual highway, so repeated block Jacobians can make gradient scale uneven across depth. Xiong et al. connect Post-LN initialization to especially large gradients near the output layers, and warmup reduces the size of early updates while diagnostics confirm whether that is the observed failure mode.

5.An engineer ports a GPT-2-style Pre-LN decoder but deletes the final normalization before the vocabulary projection, arguing that each block already has LayerNorm on its sublayer input. What should the review flag?

Correct answer: Restore the final norm, because Pre-LN block outputs bypass block-output normalization and the accumulated stream reaches the head.

Inside a Pre-LN block, the residual stream is added after the sublayer input normalization, so the block output itself isn't normalized. GPT-2-style Pre-LN architectures and the Pre-LN setup discussed by Xiong et al. include a final normalization before prediction; deleting it changes the architecture and can change output scale.

6.In a 96-layer Pre-LN ablation, late blocks still run, but their updates are small compared with the accumulated residual stream. Which interpretation and follow-up fit this observation?

Correct answer: It matches the reported representation-collapse risk; measure relative update sizes and compare matched norm layouts.

Representation collapse is a reported Pre-LN risk where the accumulated residual stream can dominate newer sublayer updates, making later blocks' relative changes small. It is a behavior to measure and compare against alternatives, not proof that Pre-LN always fails or that Post-LN is universally better.

7.For x = [3, 1, -1, 5], with gamma all ones and epsilon ignored, RMSNorm divides by sqrt((3^2 + 1^2 + (-1)^2 + 5^2) / 4) = 3. Which output is correct?

Correct answer: [1.00, 0.33, -0.33, 1.67], with RMS 1 and a nonzero mean because RMSNorm doesn't subtract the mean.

RMSNorm uses the root mean square of the original feature values. Here the RMS is 3, so each feature is divided by 3. Unlike LayerNorm, RMSNorm doesn't subtract the vector mean, so the normalized vector has controlled magnitude but isn't centered at zero.

8.A preprocessing change adds the same constant offset to every feature in a token vector. Gamma, beta, and epsilon stay fixed. How do LayerNorm and RMSNorm respond?

Correct answer: LayerNorm produces the same output because mean-centering removes the shared offset, while RMSNorm generally changes because it doesn't center the vector.

Adding a constant shifts LayerNorm's input mean by that same amount, so subtracting the new mean cancels the offset and leaves the centered deviations unchanged. RMSNorm divides the uncentered vector by its root mean square, so an additive offset generally changes both that scale and the resulting feature ratios.

9.A team wants to scale a very deep transformer and asks whether Pre-LN has made all other normalization layouts obsolete. Which comparison is technically sound?

Correct answer: Pre-LN improves the identity gradient route, but DeepNorm and input-output norm layouts are distinct designs to evaluate under matched initialization and schedules.

Pre-LN gives a direct identity gradient route, but it's not a universal winner. DeepNorm is a Post-LN-derived design with residual scaling and matched initialization, while Peri-LN-style or Pre + Post norm layouts normalize around sublayers to manage activation and gradient behavior. These recipes must be compared under matched conditions.

10.DeepNorm writes a Post-LN-shaped update as x_{l+1} = LN(alpha * x_l + G_l(x_l)). What does alpha do, and what else must be matched?

Correct answer: Alpha scales the skip path by depth; match it with smaller residual-branch initialization.

DeepNorm keeps Post-LN shape but inserts a depth-dependent residual scale alpha on the skip path and pairs it with smaller residual-branch init gains. That combination, not LayerNorm alone, is what the DeepNet recipe uses to train extremely deep stacks.

10 questions remaining.

Next Step
Continue to Mechanistic Interpretability

You can now read a Pre-LN residual stream as an accumulating sum of block writes, with a final norm before the head. The next chapter opens that stream and shows how sparse autoencoders turn opaque activation vectors into features you can inspect, test, and sometimes steer.

PreviousPositional Encoding: RoPE & ALiBi
Share this article
XFacebookLinkedInBlueskyRedditHacker NewsEmail
References

On Layer Normalization in the Transformer Architecture.

Xiong, R., et al. · 2020 · ICML 2020

https://arxiv.org/abs/2002.04745

Layer Normalization.

Ba, J. L., Kiros, J. R., & Hinton, G. E. · 2016

https://arxiv.org/abs/1607.06450

Attention Is All You Need.

Vaswani, A., et al. · 2017

https://arxiv.org/abs/1706.03762

BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.

Devlin, J., et al. · 2019 · NAACL 2019

https://arxiv.org/abs/1810.04805

Language Models are Unsupervised Multitask Learners.

Radford, A., et al. · 2019

https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf

ResiDual: Transformer with Dual Residual Connections.

Xie, Z., et al. · 2023 · ACL 2023

https://arxiv.org/abs/2304.14802

DeepNet: Scaling Transformers to 1,000 Layers.

Wang, H., Ma, S., et al. · 2022

https://arxiv.org/abs/2203.00555

Root Mean Square Layer Normalization.

Zhang, B. & Sennrich, R. · 2019 · NeurIPS 2019

https://arxiv.org/abs/1910.07467

Qwen2.5 Technical Report

Qwen Team · 2024

https://arxiv.org/abs/2412.15115

Gemma 2: Improving Open Language Models at a Practical Size

Gemma Team, Google DeepMind · 2024

https://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf

QK-Norm: Improving Transformer Training Stability.

Rybakov, O., et al. · 2024

https://arxiv.org/abs/2410.16682

Peri-LN: Revisiting Layer Normalization in the Transformer Architecture.

Kim, J., et al. · 2025 · ICML 2025

https://arxiv.org/abs/2502.02732

Discussion

Questions and insights from fellow learners.

Discussion loads when you reach this section.