AI Computer Institute
Expert-curated CS & AI curriculum aligned to CBSE standards. A bharath.ai initiative. About Us

Transformer Architecture: Layer Design and Stacking

📚 Deep Learning⏱️ 20 min read🎓 Grade 12
✍️ AI Computer Institute Editorial Team Updated: September 2026 CBSE-aligned · Peer-reviewed · 20 min read
Content curated by subject matter experts with IIT/NIT backgrounds. All chapters are fact-checked against official CBSE/NCERT syllabi.

Mumbai's dabbawalas move about 200,000 lunchboxes a day through a relay of five or six handlers between a home kitchen and an office desk, and the system famously runs at near-Six-Sigma reliability. The reason isn't that every handler is infallible. It's that the tiffin itself, the physical box, passes through the chain untouched by default. Each handler reads the coded markings already on the lid, adds or corrects a routing mark if needed, and hands the box on. If one handler in the middle has a bad day and adds nothing useful, the box still arrives, just missing one refinement. Nothing downstream depends on any single handler having done real work. This is precisely the design decision that makes a 96-layer neural network trainable: a "residual stream" that carries the representation straight through the whole stack, with each layer allowed to add a correction to it rather than being forced to compute it from scratch. Get that wrapping wrong and a stack of good sublayers still fails to train past a dozen or so layers. This chapter is about that wrapping: how a single transformer layer is built from its two sublayers, exactly where the normalization and the addition go, and what happens numerically when you stack many of these on top of each other.

The residual stream: the one object every layer touches

Strip away attention and feed-forward math for a moment and look at what actually flows between layers. Let x_l be the sequence of vectors (one per token) entering layer l. A transformer layer does not replace x_l with something new. It computes a correction and adds it: x_{l+1} = x_l + F_l(x_l), where F_l is whatever the sublayer computes (attention, or the feed-forward network). Stack L of these and the final representation is literally a sum: the original embedding plus the accumulated corrections of every layer, x_L = x_0 + F_0(x_0) + F_1(x_1) + ... + F_{L-1}(x_{L-1}). This additive structure, introduced for very deep convolutional networks by He, Zhang, Ren and Sun (2016, "Deep Residual Learning for Image Recognition"), is what Vaswani et al. (2017, "Attention Is All You Need") built the transformer layer around. Its main practical consequence is about gradients, not activations: because each layer adds to the stream rather than overwriting it, the gradient of the loss with respect to x_0 has a direct path through every addition, unmodified by any layer's Jacobian. A layer whose weights start near zero (which is exactly how they are initialized) does not block gradient flow the way a plain feed-forward stack without residuals would. That single property is what makes it possible to stack dozens of layers and still train layer 3 and layer 90 with the same optimizer and the same learning rate.

Two sublayers, one wrapping pattern: pre-LN vs post-LN

Every transformer layer has exactly two sublayers, multi-head self-attention and a position-wise feed-forward network, and both are wrapped the same way. The question that actually distinguishes architectures is where LayerNorm sits relative to the addition. The original paper wrapped it as x_{l+1} = LN(x_l + F_l(x_l)), normalizing after the residual add. This is called post-LN, and it is literally what the "Add & Norm" box in the famous diagram means. Almost no large language model trained after about 2019 uses this. GPT-2, GPT-3, and the LLaMA family instead use pre-LN: x_{l+1} = x_l + F_l(LN(x_l)), where the sublayer reads a normalized copy of the stream but the addition happens on the raw, unnormalized stream. The difference sounds cosmetic. It isn't. Xiong et al. (2020, "On Layer Normalization in the Transformer Architecture") traced the gradient at initialization for both variants and showed that in post-LN, the gradient magnitude at the output layer is large and grows with depth unless training is preceded by a slow learning-rate warm-up; skip the warm-up on a deep post-LN stack and training diverges in the first few hundred steps. In pre-LN, the gradient reaching any layer is already well-scaled at initialization because the normalization inside each branch bounds how much any single sublayer's output can perturb the stream, so warm-up becomes optional rather than load-bearing. That is the actual reason production LLM codebases moved the LayerNorm inside the residual branch: it is a training-stability fix for depth, not a stylistic preference.

Worked example: tracing variance through one block

Pre-LN buys stability, but it does not buy something for free: it lets the magnitude of the residual stream drift upward, layer after layer, because LayerNorm only normalizes the branch going into a sublayer, never the stream itself. Trace this by hand with a toy 4-dimensional stream and a toy sublayer.

Let the incoming vector be x0 = [1, -1, 2, -2]. Its mean is 0 and its (population) variance is (1+1+4+4)/4 = 2.5, so its standard deviation is std0 = sqrt(2.5) ≈ 1.5811. LayerNorm with no learned scale or shift (gamma=1, beta=0) divides by that standard deviation: LN(x0) = x0 / 1.5811 = [0.6325, -0.6325, 1.2649, -1.2649]. This vector has variance exactly 1 by construction, that's what LayerNorm guarantees regardless of the input.

Now stand in for the sublayer (attention plus its output projection, in reality) with a simple affine map, f(z) = 0.5·z + [1,1,1,1], chosen only so the arithmetic is checkable by hand. Applying it:

import numpy as np

def layer_norm(x, gamma=1.0, beta=0.0, eps=1e-8):
    mean = x.mean()
    var = x.var()  # population variance, ddof=0, matches LayerNorm
    return gamma * (x - mean) / np.sqrt(var + eps) + beta

def sublayer(z):
    return 0.5 * z + np.ones_like(z)

x0 = np.array([1.0, -1.0, 2.0, -2.0])
x1 = x0 + sublayer(layer_norm(x0))

print(x1)
print(x1.var())

Following this by hand: f(LN(x0)) = 0.5·[0.6325, -0.6325, 1.2649, -1.2649] + [1,1,1,1] = [1.3163, 0.6838, 1.6325, 0.3675]. Adding this to x0 gives x1 = [2.3163, -0.3163, 3.6325, -1.6325], so print(x1) prints that array (to four decimals), and its mean is 1.0. The deviations from that mean are [1.3163, -1.3163, 2.6325, -2.6325]; squaring and averaging gives x1.var() ≈ 4.3311.

Here is the point of the exercise. A naive guess would be that variance is additive across independent contributions: Var(x0) + Var(f) = 2.5 + 0.25 = 2.75, since Var(f) = 0.5² · Var(LN(x0)) = 0.25 · 1 = 0.25. The actual answer, 4.3311, is well above that. The gap is the covariance term the naive guess dropped: Var(x0 + f) = Var(x0) + Var(f) + 2·Cov(x0, f). Because f here is literally a rescaled copy of x0 (LayerNorm just divides by a constant when the mean is already zero), x0 and f are strongly correlated, not independent, and Cov(x0, f) = 0.5 · Cov(x0, LN(x0)) = 0.5 · std0 = 0.7906. Plugging in: 2.5 + 0.25 + 2(0.7906) = 4.3311, matching the computed value exactly. In fact, because the toy sublayer is linear, x1 is just an affine rescaling of x0, and the whole thing collapses to a clean closed form: Var(x1) = (std0 + a)² where a=0.5 is the sublayer's coefficient. Check it: (1.5811 + 0.5)² = 2.0811² = 4.3310, matching to rounding.

Real sublayers are not linear rescalings of their input (attention has a softmax, the FFN has a GELU nonlinearity, and both mix across the whole 768- or 12288-dimensional vector rather than acting elementwise), so this exact closed form does not hold at scale. But the qualitative finding survives every empirical check: because each pre-LN block adds a positively-correlated contribution back onto the stream it just read from, the stream's norm grows monotonically with depth, and the growth is faster than the naive "sum of independent variances" estimate would suggest. That is the real cost of the pre-LN fix, and it is exactly what the next section's architecture had to be designed around.

How deep, how wide: real numbers

Depth and width are both hyperparameters of the stack, and production models use surprisingly stable ratios. Vaswani et al.'s original transformer used 6 encoder layers and 6 decoder layers, d_model = 512, and a feed-forward inner dimension d_ff = 2048, a 4x expansion ratio that has stuck almost everywhere since. BERT-base (Devlin et al., 2019) uses 12 layers at d_model = 768, which under the same 4x rule gives d_ff = 3072; each feed-forward sublayer then has two weight matrices, 768×3072 and 3072×768, for 2 × 768 × 3072 = 4,718,592 parameters, not counting biases, and that single sublayer's parameter count alone is larger than the entire attention sublayer at the same width. GPT-3 (Brown et al., 2020), at the far end, uses 96 layers with d_model = 12288 and 96 attention heads. Applying the residual-stream reasoning above to a stack that deep is not academic: if each of GPT-3's 96 layers contributed even a roughly unit-variance, roughly-independent addition to the stream, the stream's standard deviation at the output would be on the order of sqrt(1 + 96) ≈ 9.85 times its starting value, and the worked example above shows the true growth (with realistic positive correlation) runs faster than that independence assumption. This is exactly why every pre-LN model, without exception, appends one final LayerNorm after the last block and before the output projection (GPT-2's checkpoint calls this layer ln_f): the stream that reaches the last block has drifted far from unit scale, and the unembedding matrix was never trained to see anything but a normalized input.

When depth breaks the recipe: DeepNorm

Pre-LN keeps training stable up to the depths just discussed, but it does not scale indefinitely: past a few hundred layers, the residual stream's growth becomes large enough that gradients through the final blocks become vanishingly informative relative to gradients through the early blocks, an effect researchers call representation collapse. Wang et al. (2022, "DeepNet: Scaling Transformers to 1,000 Layers") addressed this directly, motivated by exactly the variance-accumulation argument traced above. Their fix, DeepNorm, keeps the post-LN structure — LayerNorm applied after the residual addition — but rewrites the wrapping as x_{l+1} = LN(α·x_l + G_l(x_l,θ_l)): it scales up the incoming stream by a constant α>1 before adding the sublayer's output, then normalizes that sum, and simultaneously scales down the sublayer's own weight initialization by a matching depth-dependent factor β. The net effect is to hold the relative contribution of each new layer roughly constant regardless of how many layers came before it, which is the property the naive "independent variance" model implicitly assumed and the worked example showed pre-LN violates on its own. With this modification the authors demonstrated stable training of transformers with up to 1,000 layers, two orders of magnitude beyond what unmodified pre-LN reliably supports. The broader lesson for layer design is that "stack more layers" is never architecture-neutral: every depth regime past the one a wrapping was validated at can require revisiting how the residual branch is scaled.

Diagram: the pre-LN stack as a residual highway

Pre-LN Transformer Stack: the Residual Stream View Each sublayer reads a normalized copy of the stream and adds its output back onto the unnormalized stream x₀ = embeddings LN Multi-Head Attention + Var: 2.50 → 4.33 (worked example) LN Feed-Forward, 4×, GELU + = layer l output stack repeats: L−2 more identical layers (L=96 for GPT-3) LN Multi-Head Attention + LN Feed-Forward, 4×, GELU + stream norm grows ~monotonically with depth LN_f → unembedding → logits LayerNorm Self-Attention Feed-Forward Residual add (+) Residual stream x_l

A misconception worth retiring

Looking at the original "Attention Is All You Need" diagram, most students conclude that "Add & Norm" means the addition and the normalization are interchangeable in order, or that the exact placement is a minor implementation detail. It is not. Post-LN, exactly as drawn in that figure, normalizes the sum x_l + F_l(x_l) after the fact. Pre-LN, used in GPT-2 onward, normalizes only the copy of x_l that feeds into F_l, leaving the addition itself unnormalized. These are not two notations for the same computation; they have measurably different gradient behavior at initialization, as the Xiong et al. analysis showed, and that difference is the actual reason virtually every LLM you can name (GPT-3, LLaMA, GPT-4-class systems as far as is publicly known) diverged from the textbook diagram. If you see production model code with LayerNorm as the very first operation inside a block rather than wrapped around the residual sum, that is not a bug relative to the paper, it is the paper's own follow-up correction, made necessary by exactly the depth these systems needed to reach.

Active recall

Q1. In one sentence, why does a post-LN transformer typically need learning-rate warm-up to train, while a pre-LN transformer usually does not?

Q2. Using the worked example's setup (x0 = [1,-1,2,-2], f(z) = a·z + [1,1,1,1] applied to LN(x0)), recompute Var(f), Cov(x0,f), and Var(x1) if the sublayer's coefficient is changed from a=0.5 to a=1.0.

Q3. If each of GPT-3's 96 layers contributed a roughly independent, unit-variance addition to a residual stream that started at variance 1, what would the standard deviation of the final stream be relative to the starting standard deviation?

Q4. BERT-base uses d_model = 768 under the same 4x feed-forward expansion ratio as the original transformer. What is d_ff, and how many weight parameters (ignoring biases) does one feed-forward sublayer have?

Q5. Why can't you simply drop LayerNorm from a pre-LN block and keep only the residual connections, if residual connections are what actually protects gradient flow?

Q6. Pre-LN stacks always apply one extra LayerNorm after the last block, before the output projection. Given the worked example, explain why this final normalization is necessary rather than optional.

Worked answers

A1. In post-LN, the gradient at the output layer is large at initialization and grows with depth because the normalization comes after the residual sum, so nothing bounds how much a single early, poorly-scaled sublayer's output can distort the sum before it is corrected; a small initial learning rate (warm-up) is needed to avoid divergence in the first steps. In pre-LN, the sublayer only ever sees a normalized (unit-variance) input, so its output scale is well-behaved from the very first step and no warm-up is needed for stability, per Xiong et al. (2020).

A2. With a=1.0, Var(f) = a² · Var(LN(x0)) = 1² · 1 = 1 (LayerNorm output always has unit variance). Cov(x0,f) = a · Cov(x0, LN(x0)) = a · std0 = 1 · 1.5811 = 1.5811. Then Var(x1) = Var(x0) + Var(f) + 2·Cov(x0,f) = 2.5 + 1 + 2(1.5811) = 6.6623. Using the closed form for this linear toy case, (std0 + a)² = (1.5811 + 1)² = 2.5811² ≈ 6.6621, matching to rounding. Doubling the sublayer's coefficient nearly doubled the resulting stream variance, showing that even a "local" hyperparameter of one sublayer has a compounding, non-obvious effect on the global scale of the residual stream, not just on that one layer's own output.

A3. Variance would be 1 + 96 = 97, so the standard deviation ratio is sqrt(97) ≈ 9.85. The worked example showed real (positively correlated) growth exceeds this independence estimate, so 9.85x is a lower bound on how much the stream's scale can drift by the final layer, reinforcing why a final LayerNorm before unembedding is not optional at this depth.

A4. d_ff = 4 × 768 = 3072. One feed-forward sublayer has two linear maps, 768→3072 and 3072→768, giving 768×3072 + 3072×768 = 2 × 768 × 3072 = 4,718,592 weight parameters, before biases.

A5. The residual connection guarantees a gradient path through addition, but it says nothing about the scale of the activations flowing through that path. Without LayerNorm re-centering and re-scaling the input to each sublayer, the sublayer would receive an input whose variance keeps growing across the stack (exactly as the worked example traced), pushing its internal computations (dot products in attention, GELU inputs in the FFN) into regimes far from where the weights were initialized, which causes vanishing or exploding activations independent of the gradient-flow benefit the residual connection provides.

A6. The worked example showed that even a single pre-LN block can push the stream's variance from 2.5 to 4.33, a growth that compounds over dozens or hundreds of layers (empirically closer to sqrt(depth) or faster). The unembedding matrix that turns the final hidden state into vocabulary logits is trained expecting inputs at a particular, consistent scale. Without a final LayerNorm, that matrix would have to be trained against a residual stream whose scale has drifted far from its initial distribution, purely as an artifact of stacking, rather than against a well-conditioned, depth-independent representation.

Think About It

Think about this: How would you explain transformer architecture: layer design and stacking to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind transformer architecture: layer design and stacking, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.

← Attention Mechanisms: The Foundation of TransformersScaling Laws: Understanding Model and Data Relationships →

Found this useful? Share it!

📱 WhatsApp 🐦 Twitter 💼 LinkedIn