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 tracks word order. What we still lack is the scaffolding that makes the network trainable at depth by preserving information across layers and keeping activations numerically stable. In this part we wrap multi-head attention with residual connections and layer normalization, then append a position-wise feed-forward network to produce the complete Transformer Encoder Layer, the smallest reusable unit in the architecture.

What an Encoder Layer does

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 mixes information from every token across the sequence.
2. Add & Norm adds the attention output to the original input (residual), then normalizes.
3. Position-wise feed-forward network (FFN) runs a per-position multi-layer perceptron (MLP) that lets each token transform its representation independently.
4. Add & Norm applies the same residual and normalization step 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 layer normalization. The two residual arrows in the diagram do the heavy lifting; without them, the rest of the block would barely train at depth.

Residual connections

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. 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, instead of having to learn each layer from scratch through random initialization.

In a Transformer, residuals are essential. A 12-layer Transformer without residuals is 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⁻⁴, effectively 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 whose values grow or shrink across layers and destabilize training. Layer normalization keeps activation scales under control. 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 use layer normalization instead of 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, giving the same behaviour at train and test time 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

Attention alone only computes weighted averages of values and cannot learn arbitrary per-token functions. 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, with 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

Here is the entire Transformer Encoder Layer in standard PyTorch, matching PyTorch’s own TransformerEncoderLayer:

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 account for 32×64 + 64 = 2,112 and 64×32 + 32 = 2,080 → 4,192 weights + biases
• multi-head attention (MHA) projections use four 32×32 matrices inside MultiheadAttention → 4 × 1,024 = 4,096 weights
• LayerNorm γ/β pairs add 2 × 2 × 32 = 128 parameters
• MHA biases add 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 original “Attention Is All You Need” paper applied LayerNorm after the residual addition (the post-norm layout I showed above), while modern models moved normalization before the sub-layer to stabilize deep training:

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) allows gradients to grow large in deep stacks and requires careful learning-rate warmup.
• Pre-norm (modern) keeps gradients stable for faster convergence with 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 with more layers, a larger d_model, and a 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 using the exact 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

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, since 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, which works well for understanding. Generating text (translating from one language to another, completing a sentence, or writing code) requires 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.