LLM Under the Hood — Part 5: The Transformer Encoder

Disclaimer: The opinions expressed in this article are my own and do not represent the views of Google. This content is based solely on publicly available information.

TL;DR of Parts 1-4. Embeddings turn words into vectors (Part 1). LSTMs read them sequentially but can’t be parallelised and struggle with long-range dependencies (Part 2). Scaled dot-product attention reads every position in parallel using Q, K, V projections (Part 3). Multi-head attention runs several attention patterns in parallel for richer relations, and sine/cosine positional encoding injects word order back into the otherwise position-blind mechanism (Part 4). See Part 4: Multi-Head Attention & Positional Encoding for the building block we now need to extend.

Series navigation

Part 4 left us with multi-head attention and sine/cosine positional encoding — a mechanism that reads every position in parallel and has a notion of word order. What we still lack is the scaffolding that makes the whole thing trainable at depth: a way to preserve information across layers and a way to keep activations numerically stable. In this part we add exactly that, wrapping multi-head attention with residual connections and layer normalization, then appending a small position-wise feed-forward network to produce the complete Transformer Encoder Layer — the smallest reusable unit in the architecture.

What an Encoder Layer does

Before we write any code it helps to have the end goal in mind: a layer that takes a sequence of vectors and returns a sequence of the same shape, allowing arbitrary stacking. The four operations below give us exactly that.

The encoder layer takes a sequence of d_model-dimensional vectors and returns a sequence of d_model-dimensional vectors of the same length and same dimension. That property is what allows us to stack many of these on top of each other.

Inside, four operations run in sequence:

1. Multi-head self-attention — every token mixes information from every other token.
2. Add & Norm — add the attention output to the original input (residual), then normalize.
3. Position-wise feed-forward network (FFN) — a per-position multi-layer perceptron (MLP) that lets each token transform its representation independently.
4. Add & Norm — same trick, applied around the FFN.

flowchart TD
    A["Input\n(batch, seq, d_model)"] --> B["Multi-Head Self-Attention"]
    B --> C["Add & Norm\n(residual #1)"]
    C --> D["Feed-Forward Network\n(d_model → dff → d_model)"]
    D --> E["Add & Norm\n(residual #2)"]
    E --> F["Output\n(batch, seq, d_model)"]

    A -- "residual" --> C
    C -- "residual" --> E

Figure 1: A single Transformer Encoder Layer.

The same shape flows in (batch, seq, d_model) and out, which is precisely what makes stacking work. The block contains two sub-layers — multi-head self-attention and a position-wise FFN — each wrapped with Add & Norm: a residual connection that re-adds the sub-layer’s input to its output, followed by a layer normalization. The two residual arrows in the diagram are doing the heavy lifting; without them, the rest of the block would barely train at depth.

Let’s look at each piece.

Residual connections

Now that we know what the encoder layer is supposed to do, the natural question is how to make it trainable when we stack dozens of them. The answer is the residual connection — arguably the most impactful single trick in deep learning.

A residual connection simply adds a layer’s input to its output:

import torch
import torch.nn as nn

def with_residual(layer, x):
    return layer(x) + x   # or: x + layer(x)

This single trick, popularized by ResNet (He et al., 2015) in computer vision, makes very deep networks trainable. The intuition: at the start of training, when layer produces near-random output, the residual ensures that information from x still reaches the next layer unchanged. The model can incrementally add useful transformations on top of the identity, rather than having to learn each layer from scratch through random initialization.

In a Transformer, residuals are not optional — they are essential. A 12-layer Transformer without residuals would be nearly impossible to train. With them, even 96-layer models train cleanly.

You can verify the gradient benefit of residuals directly:

import torch
import torch.nn as nn

torch.manual_seed(0)
deep_no_residual = nn.Sequential(*[nn.Linear(64, 64) for _ in range(20)])
deep_residual    = nn.ModuleList([nn.Linear(64, 64) for _ in range(20)])

x = torch.randn(1, 64, requires_grad=True)

# No residual: chain of linear layers
y = x.clone()
for layer in deep_no_residual:
    y = torch.tanh(layer(y))
loss_no_res = y.sum()
loss_no_res.backward()
grad_no_res = x.grad.norm().item()

x.grad = None
y = x.clone()
# residual around a sub-layer that includes a tanh non-linearity, matching a real Transformer block
for layer in deep_residual:
    y = y + torch.tanh(layer(y))    # residual
loss_res = y.sum()
loss_res.backward()
grad_res = x.grad.norm().item()

print(f"Gradient norm WITHOUT residuals: {grad_no_res:.6f}")
print(f"Gradient norm WITH    residuals: {grad_res:.6f}")

Output:

Gradient norm WITHOUT residuals: 0.000123
Gradient norm WITH    residuals: 41.508270

Without residuals in a 20-layer network, the gradient reaching the input is on the order of 10⁻⁴ — vanished through the layers. With residuals, the gradient propagates cleanly, five orders of magnitude larger.

Layer normalization

Residual connections solve the vanishing-gradient problem, but deep networks still suffer from numerically unstable activations — values that grow or shrink across layers and destabilize training. Layer normalization is the complement to residuals that keeps the scale under control.

The other ingredient is layer normalization. For each token vector, we subtract the mean and divide by the standard deviation across the model dimension, then apply a learned scale and bias:

$$\text{LN}(x) = \gamma \odot \frac{x - \mu}{\sigma + \epsilon} + \beta$$

γ and β are learned per-feature parameters (size d_model). μ and σ are computed per-token, across the model dimension.

Why this rather than batch normalization? Batch norm normalizes across the batch dimension, which interacts badly with variable-length sequences (each sequence in a batch may have a different length) and behaves differently at train vs. inference time when batch size is small. Layer norm depends only on the single token’s vector — same behaviour at train and at test, regardless of batch size or sequence length.

import torch
import torch.nn as nn

x = torch.tensor([[1.0, 2.0, 3.0, 4.0],
                  [10.0, 20.0, 30.0, 40.0]])

ln = nn.LayerNorm(normalized_shape=4)     # γ=1, β=0 at init
x_normed = ln(x)

print("Input:")
print(x)
print("\nAfter LayerNorm:")
print(x_normed.round(decimals=4))
print("\nMean per token (should be ~0):", x_normed.mean(dim=-1).tolist())
print("Std  per token (should be ~1):", x_normed.std(dim=-1, unbiased=False).tolist())

Output:

Input:
tensor([[ 1.,  2.,  3.,  4.],
        [10., 20., 30., 40.]])

After LayerNorm:
tensor([[-1.3416, -0.4472,  0.4472,  1.3416],
        [-1.3416, -0.4472,  0.4472,  1.3416]], grad_fn=<RoundBackward1>)

Mean per token (should be ~0): [0.0, 0.0]
Std  per token (should be ~1): [0.9999960064888, 1.0]

Both tokens are normalized to the same distribution, regardless of their original scale (1-4 vs. 10-40).

The feed-forward network

We have a stable, trainable mechanism for mixing information across token positions. But attention alone can only compute weighted averages of values — it cannot learn arbitrary per-token functions. That is where the feed-forward network comes in, giving each position its own independent transformation after the cross-position mixing is done.

After attention has mixed information across positions, we run a small two-layer MLP independently at every position:

$$\text{FFN}(x) = \max(0,, x W_1 + b_1), W_2 + b_2$$

Standard expand-and-contract structure: project up to a higher dimension dff (typically 4 × d_model), apply a non-linearity (ReLU in the original paper, GELU in most modern models), project back down to d_model. The same FFN weights are applied to every position — there is no cross-position interaction here. That happens only in the attention.

You can think of attention as the layer that does information mixing across tokens and the FFN as the layer that does per-token transformation. Both are essential.

Empirical evidence: models with attention but no FFN fail to learn many factual associations. The FFN appears to act as a kind of key-value memory — storing facts in the weight matrices that can be “looked up” when the right query pattern activates the relevant neurons (Geva et al., 2021).

Putting it all together in PyTorch

With residuals, layer norm, and the FFN individually understood, assembling the full encoder layer is a matter of wiring them in the right order. The code below is exactly what PyTorch’s own TransformerEncoderLayer implements under the hood.

Here is the entire Transformer Encoder Layer in standard PyTorch:

import torch
import torch.nn as nn

class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, dff, dropout=0.1):
        super().__init__()
        self.mha   = nn.MultiheadAttention(d_model, num_heads, dropout=dropout,
                                            batch_first=True)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, dff),
            nn.ReLU(),
            nn.Linear(dff, d_model),
        )
        self.drop1 = nn.Dropout(dropout)
        self.drop2 = nn.Dropout(dropout)

    def forward(self, x):
        attn, _ = self.mha(x, x, x, need_weights=False)
        x = self.norm1(x + self.drop1(attn))   # residual #1
        ff = self.ffn(x)
        return self.norm2(x + self.drop2(ff))  # residual #2

Notice we pass x three times to self.mha(...). That is what makes it self-attention: queries, keys, and values all come from the same input. (When we get to the decoder in Part 6, we will see cross-attention, where the query comes from one source and key+value from another.)

I built this layer with d_model=32, num_heads=4, dff=64 and called print on it. PyTorch gives us the actual structure of every sub-module:

import torch
import torch.nn as nn

layer = EncoderLayer(d_model=32, num_heads=4, dff=64)
print(layer)
n_params = sum(p.numel() for p in layer.parameters())
print(f"\nTotal parameters: {n_params:,}")

EncoderLayer(
  (mha): MultiheadAttention(
    (out_proj): NonDynamicallyQuantizableLinear(in_features=32, out_features=32, bias=True)
  )
  (norm1): LayerNorm((32,), eps=1e-05, elementwise_affine=True)
  (norm2): LayerNorm((32,), eps=1e-05, elementwise_affine=True)
  (ffn): Sequential(
    (0): Linear(in_features=32, out_features=64, bias=True)
    (1): ReLU()
    (2): Linear(in_features=64, out_features=32, bias=True)
  )
  (drop1): Dropout(p=0.1, inplace=False)
  (drop2): Dropout(p=0.1, inplace=False)
)

Total parameters: 8,544

Every Linear, every LayerNorm, every Dropout is exactly where you would expect. The parameter count is dominated by:

• FFN linear layers: 32×64 + 64 = 2,112 and 64×32 + 32 = 2,080 → 4,192 weights + biases
• multi-head attention (MHA) projections: four 32×32 matrices inside MultiheadAttention → 4 × 1,024 = 4,096 weights
• LayerNorm γ/β pairs: 2 × 2 × 32 = 128 parameters
• MHA biases: 3×32 (in_proj) + 32 (out_proj) = 128 parameters

Total: 8,544 parameters.

Let us run a real forward pass and check the shapes at each stage:

import torch

torch.manual_seed(0)
layer = EncoderLayer(d_model=32, num_heads=4, dff=64)
layer.eval()

x = torch.randn(2, 10, 32)     # batch=2, seq=10, d_model=32

with torch.no_grad():
    out = layer(x)

print(f"Input  shape : {x.shape}")     # (2, 10, 32)
print(f"Output shape : {out.shape}")   # (2, 10, 32)
print(f"Shape unchanged: {x.shape == out.shape}")
print(f"\nInput  mean/std  : {x.mean():.4f} / {x.std():.4f}")
print(f"Output mean/std  : {out.mean():.4f} / {out.std():.4f}")

Output:

Input  shape : torch.Size([2, 10, 32])
Output shape : torch.Size([2, 10, 32])
Shape unchanged: True

Input  mean/std  : -0.0232 / 1.0324
Output mean/std  : 0.0000 / 1.0008

Same shape in, same shape out — the fundamental invariant that makes these layers stackable. The output mean and std sit very close to 0 and 1 because the final norm2 re-normalizes each token vector.

Pre-norm vs post-norm

The code above matches the original “Attention Is All You Need” paper exactly — but modern models quietly changed one detail that significantly affects training stability. It is worth understanding which variant you are looking at when you read open-source model code.

One subtle detail: the original Transformer paper applied LayerNorm after the residual addition (the post-norm layout I showed above):

x = norm(x + sublayer(x))

Most modern implementations use pre-norm, applying LayerNorm to the input before the sub-layer:

import torch
import torch.nn as nn

class PreNormEncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, dff, dropout=0.1):
        super().__init__()
        self.mha   = nn.MultiheadAttention(d_model, num_heads, dropout=dropout,
                                            batch_first=True)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn   = nn.Sequential(
            nn.Linear(d_model, dff), nn.GELU(), nn.Linear(dff, d_model)
        )
        self.drop  = nn.Dropout(dropout)

    def forward(self, x):
        # Pre-norm: normalize BEFORE the sub-layer
        normed = self.norm1(x)
        attn, _ = self.mha(normed, normed, normed, need_weights=False)
        x = x + self.drop(attn)

        normed = self.norm2(x)
        return x + self.drop(self.ffn(normed))

Pre-norm gives more stable gradients in deep stacks and has become the default in models like Llama, Qwen, and GPT-style architectures. The difference:

• Post-norm (original): gradients can grow large in deep stacks; requires careful learning-rate warmup.
• Pre-norm (modern): gradients are more stable; faster convergence, fewer warmup steps needed.

For the conceptual story, both work — and both produce the same building block: same shape in, same shape out.

Stacking encoders to depth

The key payoff of the shape-preserving design — same (batch, seq, d_model) in and out — is that we can stack as many encoder layers as we want by simply repeating them. Depth is how we get from a tiny 8 K-parameter proof-of-concept to the 110 M–400 M encoder models used in production natural language processing (NLP).

A Transformer Encoder Stack is just a list of N independent EncoderLayers, applied in sequence:

import torch
import torch.nn as nn

class TransformerEncoder(nn.Module):
    def __init__(self, num_layers, d_model, num_heads, dff, dropout=0.1):
        super().__init__()
        self.layers = nn.ModuleList(
            [EncoderLayer(d_model, num_heads, dff, dropout)
             for _ in range(num_layers)]
        )
        self.norm = nn.LayerNorm(d_model)    # final layer norm

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return self.norm(x)

# A 6-layer encoder, d_model=512, 8 heads, dff=2048
encoder = TransformerEncoder(num_layers=6, d_model=512, num_heads=8, dff=2048)
n = sum(p.numel() for p in encoder.parameters())
print(f"6-layer encoder parameters: {n:,}")    # 18,915,328

Output:

6-layer encoder parameters: 18,915,328

Because each layer has the same input and output shape, we can stack 6, 12, 24, or 96 of them. Bigger models do exactly this — more layers, larger d_model, larger dff. The architecture stays identical.

Real model sizes for comparison

Note: ModernBERT uses GeGLU in its FFN, so its effective FFN hidden size and dff are not directly comparable to BERT’s standard 4 × d_model = 4096.

All of these are encoder-only models — exactly the architecture we just built, at larger scale. They excel at understanding tasks (classification, named entity recognition, question answering) but cannot generate text.

Encoder-only Transformers: BERT

Stacking encoder layers at scale gives us the encoder-only architecture, and BERT is its most influential instantiation. Understanding BERT’s training objective rounds out the picture of what an encoder-only model is good for — and why it cannot generate text.

Encoder-only Transformers like Bidirectional Encoder Representations from Transformers (BERT) (Devlin et al., 2018) are exactly: embedding layer + positional encoding + a stack of encoder layers + a task-specific head. They excel at understanding (classification, named-entity recognition, question answering) but not at generation — there is no mechanism to produce one token at a time.

BERT’s pre-training objective is Masked Language Modeling (MLM): randomly mask 15% of tokens and ask the model to predict the masked tokens. Because the model sees bidirectional context (every token can attend to every other token), it builds rich contextual representations. BERT-base achieves state-of-the-art on the General Language Understanding Evaluation (GLUE) benchmark in 2018 with 12 encoder layers.

To generate text, we need the other half of the original architecture.

Up next — Part 6: The Decoder & Full Transformer

The encoder layer we just built reads in a sequence and produces a sequence of contextualized representations — perfect for understanding. But to generate — translate from one language to another, complete a sentence, write code — we need a second half.

In Part 6: The Decoder & Full Transformer we will build the DecoderLayer, which adds two new things on top of what we have here: masked self-attention (so the decoder cannot peek at future tokens during training) and cross-attention (so it can read the encoder’s output). We will then assemble both halves into a full nn.Transformer and look at the real print(model) of an architecture you can train end to end.