The architecture

The Transformer Architecture

Transformer

2017’s Attention is All You Need threw out RNNs entirely and built sequence models from self-attention, positionwise MLPs, residuals, and layer norm.

The Transformer is now the architecture for language, vision, speech, and beyond. Same code, 6 layers (original) to 96+ layers (frontier LLMs).

What this deck builds

Bottom-up assembly:

• positionwise feed-forward network,
• residual connection + layer norm (“Add & Norm”),
• encoder block, encoder stack,
• decoder block (masked self-attention + cross-attention),
• training and attention visualization on En→Fr.

from d2l import jax as d2l
from flax import linen as nn
from jax import numpy as jnp
import jax
import math
import pandas as pd

The Transformer architecture.

Encoder: N identical blocks (self-attention → Add & Norm → FFN → Add & Norm). Decoder: same, plus masked self-attention and encoder-decoder cross-attention. Embedding + positional encoding before the first block.

Encoder vs decoder attention

The same multi-head attention operator is used in three different roles:

• Encoder self-attention: queries, keys, and values all come from the source sequence. Every source position can read every other non-padding source position.
• Decoder masked self-attention: queries, keys, and values come from the target prefix. The causal mask hides future target tokens.
• Cross-attention: decoder states provide the queries; encoder outputs provide keys and values. This is where the target prefix looks back at the source sentence.

Positionwise FFN

A two-layer MLP applied independently at every sequence position — same weights everywhere. Lets each position process its attention output through a nonlinear feature mixer:

class PositionWiseFFN(nn.Module):
    """The positionwise feed-forward network."""
    ffn_num_hiddens: int
    ffn_num_outputs: int

    def setup(self):
        self.dense1 = nn.Dense(self.ffn_num_hiddens)
        self.dense2 = nn.Dense(self.ffn_num_outputs)

    def __call__(self, X):
        return self.dense2(nn.relu(self.dense1(X)))

Shape check: rank-3 input, only the last dim changes:

ffn = PositionWiseFFN(4, 8)
ffn.init_with_output(d2l.get_key(), jnp.ones((2, 3, 4)))[0][0]

Array([[ 0.27652267,  0.4836862 , -0.32886025,  0.06538588,  0.24364541,
        -0.39205405,  0.61960614, -0.23260641],
       [ 0.27652267,  0.4836862 , -0.32886025,  0.06538588,  0.24364541,
        -0.39205405,  0.61960614, -0.23260641],
       [ 0.27652267,  0.4836862 , -0.32886025,  0.06538588,  0.24364541,
        -0.39205405,  0.61960614, -0.23260641]], dtype=float32)

LayerNorm vs BatchNorm

BatchNorm normalizes across the batch — fragile with variable-length sequences and small batches. LayerNorm normalizes across the feature dimension of one example — batch-size and length independent. That’s why NLP picked it.

ln = nn.LayerNorm()
bn = nn.BatchNorm()
X = d2l.tensor([[1, 2], [2, 3]], dtype=d2l.float32)
# Compute mean and variance from X in the training mode
print('layer norm:', ln.init_with_output(d2l.get_key(), X)[0],
      '\nbatch norm:', bn.init_with_output(d2l.get_key(), X,
                                           use_running_average=False)[0])

layer norm: [[-0.9999979  0.9999979]
 [-0.9999979  0.9999979]] 
batch norm: [[-0.9999799 -0.9999799]
 [ 0.9999799  0.9999799]]

AddNorm

The repeating motif: residual connection (X + sublayer(X)), dropout, then LayerNorm. Both inputs must have the same shape:

class AddNorm(nn.Module):
    """The residual connection followed by layer normalization."""
    dropout: float

    @nn.compact
    def __call__(self, X, Y, training=False):
        return nn.LayerNorm()(
            nn.Dropout(self.dropout)(Y, deterministic=not training) + X)

add_norm = AddNorm(0.5)
shape = (2, 3, 4)
output, _ = add_norm.init_with_output(d2l.get_key(), d2l.ones(shape),
                                      d2l.ones(shape))
d2l.check_shape(output, shape)

Encoder block

One block = MultiHead self-attention → AddNorm → FFN → AddNorm. Shape in = shape out, so blocks stack without any projection in between:

class TransformerEncoderBlock(nn.Module):
    """The Transformer encoder block."""
    num_hiddens: int
    ffn_num_hiddens: int
    num_heads: int
    dropout: float
    use_bias: bool = False

    def setup(self):
        self.attention = d2l.MultiHeadAttention(self.num_hiddens, self.num_heads,
                                                self.dropout, self.use_bias)
        self.addnorm1 = AddNorm(self.dropout)
        self.ffn = PositionWiseFFN(self.ffn_num_hiddens, self.num_hiddens)
        self.addnorm2 = AddNorm(self.dropout)

    def __call__(self, X, valid_lens, training=False):
        output, attention_weights = self.attention(X, X, X, valid_lens,
                                                   training=training)
        Y = self.addnorm1(X, output, training=training)
        return self.addnorm2(Y, self.ffn(Y), training=training), attention_weights

Shape check:

X = jnp.ones((2, 100, 24))
valid_lens = jnp.array([3, 2])
encoder_blk = TransformerEncoderBlock(24, 48, 8, 0.5)
(output, _), _ = encoder_blk.init_with_output(d2l.get_key(), X, valid_lens,
                                              training=False)
d2l.check_shape(output, X.shape)

Encoder stack

Embed tokens, scale by \sqrt{d} to balance against the positional encoding, add positions, then run N blocks. Save attention weights per block for later visualization:

class TransformerEncoder(d2l.Encoder):
    """The Transformer encoder."""
    vocab_size: int
    num_hiddens:int
    ffn_num_hiddens: int
    num_heads: int
    num_blks: int
    dropout: float
    use_bias: bool = False

    def setup(self):
        self.embedding = nn.Embed(self.vocab_size, self.num_hiddens)
        self.pos_encoding = d2l.PositionalEncoding(self.num_hiddens, self.dropout)
        self.blks = [TransformerEncoderBlock(self.num_hiddens,
                                             self.ffn_num_hiddens,
                                             self.num_heads,
                                             self.dropout, self.use_bias)
                     for _ in range(self.num_blks)]

    def __call__(self, X, valid_lens, training=False):
        # Since positional encoding values are between -1 and 1, the embedding
        # values are multiplied by the square root of the embedding dimension
        # to rescale before they are summed up
        X = self.embedding(X) * math.sqrt(self.num_hiddens)
        X = self.pos_encoding(X, training=training)
        attention_weights = [None] * len(self.blks)
        for i, blk in enumerate(self.blks):
            X, attention_w = blk(X, valid_lens, training=training)
            attention_weights[i] = attention_w
        # Flax sow API is used to capture intermediate variables
        self.sow('intermediates', 'enc_attention_weights', attention_weights)
        return X

encoder = TransformerEncoder(200, 24, 48, 8, 2, 0.5)
d2l.check_shape(encoder.init_with_output(d2l.get_key(),
                                         jnp.ones((2, 100), dtype=jnp.int32),
                                         valid_lens)[0],
                (2, 100, 24))

Decoder block

Three sublayers, each wrapped in AddNorm:

1. Masked self-attention — every position only sees positions \le t (no peeking at the answer).
2. Cross-attention — queries from the decoder, keys/values from the encoder output.
3. FFN.

class TransformerDecoderBlock(nn.Module):
    # The i-th block in the Transformer decoder
    num_hiddens: int
    ffn_num_hiddens: int
    num_heads: int
    dropout: float
    i: int

    def setup(self):
        self.attention1 = d2l.MultiHeadAttention(self.num_hiddens,
                                                 self.num_heads,
                                                 self.dropout)
        self.addnorm1 = AddNorm(self.dropout)
        self.attention2 = d2l.MultiHeadAttention(self.num_hiddens,
                                                 self.num_heads,
                                                 self.dropout)
        self.addnorm2 = AddNorm(self.dropout)
        self.ffn = PositionWiseFFN(self.ffn_num_hiddens, self.num_hiddens)
        self.addnorm3 = AddNorm(self.dropout)

    def __call__(self, X, state, training=False):
        enc_outputs, enc_valid_lens = state[0], state[1]
        # During training, all the tokens of any output sequence are processed
        # at the same time, so state[2][self.i] is None as initialized. When
        # decoding any output sequence token by token during prediction,
        # state[2][self.i] contains representations of the decoded output at
        # the i-th block up to the current time step
        if state[2][self.i] is None:
            key_values = X
        else:
            key_values = jnp.concatenate((state[2][self.i], X), axis=1)
        state[2][self.i] = key_values
        if training:
            batch_size, num_steps, _ = X.shape
            # Shape of dec_valid_lens: (batch_size, num_steps), where every
            # row is [1, 2, ..., num_steps]
            dec_valid_lens = jnp.tile(jnp.arange(1, num_steps + 1),
                                      (batch_size, 1))
        else:
            dec_valid_lens = None
        # Self-attention
        X2, attention_w1 = self.attention1(X, key_values, key_values,
                                           dec_valid_lens, training=training)
        Y = self.addnorm1(X, X2, training=training)
        # Encoder-decoder attention. Shape of enc_outputs:
        # (batch_size, num_steps, num_hiddens)
        Y2, attention_w2 = self.attention2(Y, enc_outputs, enc_outputs,
                                           enc_valid_lens, training=training)
        Z = self.addnorm2(Y, Y2, training=training)
        return self.addnorm3(Z, self.ffn(Z), training=training), state, attention_w1, attention_w2

Decoder shape check

Run the decoder with fake encoder outputs and source valid_lens. The output is target-position logits; the state carries encoder outputs plus per-block caches used during autoregressive prediction:

decoder_blk = TransformerDecoderBlock(24, 48, 8, 0.5, 0)
X = d2l.ones((2, 100, 24))
state = [encoder_blk.init_with_output(d2l.get_key(), X, valid_lens)[0][0],
         valid_lens, [None]]
d2l.check_shape(decoder_blk.init_with_output(d2l.get_key(), X, state)[0][0],
                X.shape)

Decoder stack

Token embedding + positional encoding -> N decoder blocks -> vocab projection. During training, causal masks are built from target positions; during prediction, the cache grows one token at a time.

class TransformerDecoder(nn.Module):
    vocab_size: int
    num_hiddens: int
    ffn_num_hiddens: int
    num_heads: int
    num_blks: int
    dropout: float

    def setup(self):
        self.embedding = nn.Embed(self.vocab_size, self.num_hiddens)
        self.pos_encoding = d2l.PositionalEncoding(self.num_hiddens,
                                                   self.dropout)
        self.blks = [TransformerDecoderBlock(self.num_hiddens,
                                             self.ffn_num_hiddens,
                                             self.num_heads, self.dropout, i)
                     for i in range(self.num_blks)]
        self.dense = nn.Dense(self.vocab_size)

    def init_state(self, enc_outputs, enc_valid_lens):
        return [enc_outputs, enc_valid_lens, [None] * self.num_blks]

    def __call__(self, X, state, training=False):
        # During step-by-step prediction, position-encode the new token using
        # its true offset (the number of tokens already decoded), rather than
        # always re-applying P[0:1]. This matches the pos encoding seen at
        # training time and is critical for stable autoregressive decoding.
        pos_offset = 0 if state[2][0] is None else state[2][0].shape[1]
        X = self.embedding(X) * jnp.sqrt(jnp.float32(self.num_hiddens))
        X = self.pos_encoding(X, training=training, offset=pos_offset)
        attention_weights = [[None] * len(self.blks) for _ in range(2)]
        for i, blk in enumerate(self.blks):
            X, state, attention_w1, attention_w2 = blk(X, state,
                                                       training=training)
            # Decoder self-attention weights
            attention_weights[0][i] = attention_w1
            # Encoder-decoder attention weights
            attention_weights[1][i] = attention_w2
        # Flax sow API is used to capture intermediate variables
        self.sow('intermediates', 'dec_attention_weights', attention_weights)
        return self.dense(X), state

Training: tiny config

Same MTFraEng dataset as the seq2seq chapter. 2 layers, 256 hidden, 4 heads, dropout 0.2. Adam lr=0.001, gradient clip 1, 30 epochs:

data = d2l.MTFraEng(batch_size=128)
num_hiddens, num_blks, dropout = 256, 2, 0.2
ffn_num_hiddens, num_heads = 64, 4
encoder = TransformerEncoder(
    len(data.src_vocab), num_hiddens, ffn_num_hiddens, num_heads,
    num_blks, dropout)
decoder = TransformerDecoder(
    len(data.tgt_vocab), num_hiddens, ffn_num_hiddens, num_heads,
    num_blks, dropout)
model = d2l.Seq2Seq(encoder, decoder, tgt_pad=data.tgt_vocab['<pad>'],
                    lr=0.001)
trainer = d2l.Trainer(max_epochs=30, gradient_clip_val=1, num_gpus=1)
trainer.fit(model, data)

Translate four sentences

This is a tiny model on a tiny dataset. Look for good short translations and BLEU differences across examples; errors are usually data/model-size limits, not a change in the architecture.

engs = ['i lost .', 'i\'m calm .', 'i\'m home .']
fras = ['j\'ai perdu .', 'je suis calme .', 'je suis chez moi .']
preds, _ = model.predict_step(
    trainer.state.params, data.build(engs, fras), data.num_steps)
for en, fr, p in zip(engs, fras, preds):
    translation = []
    for token in data.tgt_vocab.to_tokens(p):
        if token == '<eos>':
            break
        translation.append(token)
    print(f'{en} => {translation}, bleu,'
          f'{d2l.bleu(" ".join(translation), fr, k=2):.3f}')

i lost . => ['je', 'suis', 'perdu', '.'], bleu,0.537
i'm calm . => ['je', 'suis', 'calme', '.'], bleu,1.000
i'm home . => ['je', 'suis', 'chez', 'moi', '.'], bleu,1.000

Encoder self-attention

Pull the encoder’s stored attention weights, reshape into (layers × heads × queries × keys), heatmap them. Different heads attend to different patterns:

_, (dec_attention_weights, enc_attention_weights) = model.predict_step(
    trainer.state.params, data.build([engs[-1]], [fras[-1]]),
    data.num_steps, True)
enc_attention_weights = d2l.concat(enc_attention_weights, 0)
shape = (num_blks, num_heads, -1, data.num_steps)
enc_attention_weights = d2l.reshape(enc_attention_weights, shape)
d2l.check_shape(enc_attention_weights,
                (num_blks, num_heads, data.num_steps, data.num_steps))

d2l.show_heatmaps(
    enc_attention_weights, xlabel='Key positions', ylabel='Query positions',
    titles=['Head %d' % i for i in range(1, 5)], figsize=(7, 3.5))

Decoder attention tensors

The decoder has two attention sublayers per block — masked self-attention and encoder-decoder cross-attention. Pull both from the prediction trace and reshape them into (blocks × heads × queries × keys):

dec_attention_weights_2d = [head[0].tolist() for step in dec_attention_weights
                            for attn in step
                            for blk in attn for head in blk]
dec_attention_weights_filled = d2l.tensor(
    pd.DataFrame(dec_attention_weights_2d).fillna(0.0).values)
dec_attention_weights = dec_attention_weights_filled.reshape(
    (-1, 2, num_blks, num_heads, data.num_steps))
dec_self_attention_weights, dec_inter_attention_weights = \
    dec_attention_weights.transpose(1, 2, 3, 0, 4)

d2l.check_shape(dec_self_attention_weights,
                (num_blks, num_heads, data.num_steps, data.num_steps))
d2l.check_shape(dec_inter_attention_weights,
                (num_blks, num_heads, data.num_steps, data.num_steps))

Decoder causal mask

The self-attention heatmap must be lower triangular: query position t can attend only to keys at positions \le t. That is what makes the decoder a language model during generation.

d2l.show_heatmaps(
    dec_self_attention_weights[:, :, :, :],
    xlabel='Key positions', ylabel='Query positions',
    titles=['Head %d' % i for i in range(1, 5)], figsize=(7, 3.5))

Cross-attention masking

Cross-attention from decoder queries to encoder keys: notice zero weight on source padding tokens. Masking with valid_lens during attention is what enforces this:

d2l.show_heatmaps(
    dec_inter_attention_weights, xlabel='Key positions',
    ylabel='Query positions', titles=['Head %d' % i for i in range(1, 5)],
    figsize=(7, 3.5))

Recap

• Transformer = encoder of N identical blocks (self-attn + FFN, each with Add & Norm), decoder of N identical blocks (masked self-attn + cross-attn + FFN).
• LayerNorm beats BatchNorm for variable-length sequences.
• Positionwise FFN is just an MLP applied per position; it gives the model nonlinear feature mixing on top of the linear-in-values attention.
• Cross-attention is the seq2seq glue: decoder queries encoder outputs at every layer.
• Same architecture scales: 6 layers (original) → 96+ layers (frontier LLMs), same code.