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 tensorflow as d2l
import numpy as np
import pandas as pd
import tensorflow as tf

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(tf.keras.layers.Layer):
    """The positionwise feed-forward network."""
    def __init__(self, ffn_num_hiddens, ffn_num_outputs):
        super().__init__()
        self.dense1 = tf.keras.layers.Dense(ffn_num_hiddens)
        self.relu = tf.keras.layers.ReLU()
        self.dense2 = tf.keras.layers.Dense(ffn_num_outputs)

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

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

ffn = PositionWiseFFN(4, 8)
ffn(tf.ones((2, 3, 4)))[0]

<tf.Tensor: shape=(3, 8), dtype=float32, numpy=
array([[-0.13792595,  0.1292612 ,  0.08510286, -0.20555764, -0.19384152,
        -0.07714051,  0.11963478,  0.13712181],
       [-0.13792595,  0.1292612 ,  0.08510286, -0.20555764, -0.19384152,
        -0.07714051,  0.11963478,  0.13712181],
       [-0.13792595,  0.1292612 ,  0.08510286, -0.20555764, -0.19384152,
        -0.07714051,  0.11963478,  0.13712181]], 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 = tf.keras.layers.LayerNormalization()
bn = tf.keras.layers.BatchNormalization()
X = tf.constant([[1, 2], [2, 3]], dtype=tf.float32)
print('layer norm:', ln(X), '\nbatch norm:', bn(X, training=True))

layer norm: tf.Tensor(
[[-0.998006    0.9980061 ]
 [-0.9980061   0.99800587]], shape=(2, 2), dtype=float32) 
batch norm: tf.Tensor(
[[-0.998006   -0.9980061 ]
 [ 0.9980061   0.99800587]], shape=(2, 2), dtype=float32)

AddNorm

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

class AddNorm(tf.keras.layers.Layer):
    """The residual connection followed by layer normalization."""
    def __init__(self, norm_shape, dropout):
        super().__init__()
        self.dropout = tf.keras.layers.Dropout(dropout)
        # `norm_shape` mirrors PyTorch's `nn.LayerNorm` convention: it gives
        # the shape of the trailing dims to normalize over. Convert that to
        # Keras's `axis` argument (negative axis indices counting from the end).
        self.ln = tf.keras.layers.LayerNormalization(
            axis=list(range(-len(norm_shape), 0)))

    def call(self, X, Y, training=False, **kwargs):
        return self.ln(self.dropout(Y, training=training) + X)

# Match the 1-axis LayerNorm used by the trained model (norm_shape=[2]
# normalizes the trailing feature axis), keeping demo and training consistent.
add_norm = AddNorm([2], 0.5)
shape = (2, 3, 4)
d2l.check_shape(add_norm(tf.ones(shape), tf.ones(shape), training=False),
                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(tf.keras.layers.Layer):
    """The Transformer encoder block."""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 norm_shape, ffn_num_hiddens, num_heads, dropout, bias=False):
        super().__init__()
        self.attention = d2l.MultiHeadAttention(
            key_size, query_size, value_size, num_hiddens, num_heads, dropout,
            bias)
        self.addnorm1 = AddNorm(norm_shape, dropout)
        self.ffn = PositionWiseFFN(ffn_num_hiddens, num_hiddens)
        self.addnorm2 = AddNorm(norm_shape, dropout)

    def call(self, X, valid_lens, training=False, **kwargs):
        Y = self.addnorm1(X, self.attention(X, X, X, valid_lens,
                          training=training), training=training)
        return self.addnorm2(Y, self.ffn(Y), training=training)

Shape check:

X = tf.ones((2, 100, 24))
valid_lens = tf.constant([3, 2])
norm_shape = [i for i in range(len(X.shape))][1:]
encoder_blk = TransformerEncoderBlock(24, 24, 24, 24, norm_shape, 48, 8, 0.5)
d2l.check_shape(encoder_blk(X, valid_lens, training=False), 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."""
    def __init__(self, vocab_size, key_size, query_size, value_size,
                 num_hiddens, norm_shape, ffn_num_hiddens, num_heads,
                 num_blks, dropout, bias=False):
        super().__init__()
        self.num_hiddens = num_hiddens
        self.embedding = tf.keras.layers.Embedding(vocab_size, num_hiddens)
        self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)
        self.blks = [TransformerEncoderBlock(
            key_size, query_size, value_size, num_hiddens, norm_shape,
            ffn_num_hiddens, num_heads, dropout, bias) for _ in range(
            num_blks)]

    def call(self, X, valid_lens, training=False, **kwargs):
        # 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.pos_encoding(self.embedding(X) * tf.math.sqrt(
            tf.cast(self.num_hiddens, dtype=tf.float32)), training=training)
        self.attention_weights = [None] * len(self.blks)
        for i, blk in enumerate(self.blks):
            X = blk(X, valid_lens, training=training)
            self.attention_weights[
                i] = blk.attention.attention.attention_weights
        return X

encoder = TransformerEncoder(200, 24, 24, 24, 24, [1, 2], 48, 8, 2, 0.5)
d2l.check_shape(encoder(tf.ones((2, 100)), valid_lens, training=False),
                (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(tf.keras.layers.Layer):
    # The i-th block in the Transformer decoder
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 norm_shape, ffn_num_hiddens, num_heads, dropout, i):
        super().__init__()
        self.i = i
        self.attention1 = d2l.MultiHeadAttention(
            key_size, query_size, value_size, num_hiddens, num_heads, dropout)
        self.addnorm1 = AddNorm(norm_shape, dropout)
        self.attention2 = d2l.MultiHeadAttention(
            key_size, query_size, value_size, num_hiddens, num_heads, dropout)
        self.addnorm2 = AddNorm(norm_shape, dropout)
        self.ffn = PositionWiseFFN(ffn_num_hiddens, num_hiddens)
        self.addnorm3 = AddNorm(norm_shape, dropout)

    def call(self, X, state, training=False, **kwargs):
        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 = tf.concat((state[2][self.i], X), axis=1)
        state[2][self.i] = key_values
        if training:
            x_shape = tf.shape(X)
            batch_size, num_steps = x_shape[0], x_shape[1]
            # Shape of dec_valid_lens: (batch_size, num_steps), where every
            # row is [1, 2, ..., num_steps]
            dec_valid_lens = tf.repeat(
                tf.reshape(tf.range(1, num_steps + 1),
                           shape=(1, -1)), repeats=batch_size, axis=0)
        else:
            dec_valid_lens = None
        # Self-attention
        X2 = 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 = 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

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, 24, 24, 24, [1, 2], 48, 8, 0.5, 0)
X = tf.ones((2, 100, 24))
state = [encoder_blk(X, valid_lens), valid_lens, [None]]
d2l.check_shape(decoder_blk(X, state, training=False)[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(d2l.AttentionDecoder):
    def __init__(self, vocab_size, key_size, query_size, value_size,
                 num_hiddens, norm_shape, ffn_num_hiddens, num_heads,
                 num_blks, dropout):
        super().__init__()
        self.num_hiddens = num_hiddens
        self.num_blks = num_blks
        self.embedding = tf.keras.layers.Embedding(vocab_size, num_hiddens)
        self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)
        self.blks = [TransformerDecoderBlock(
            key_size, query_size, value_size, num_hiddens, norm_shape,
            ffn_num_hiddens, num_heads, dropout, i)
                     for i in range(num_blks)]
        self.dense = tf.keras.layers.Dense(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, **kwargs):
        # 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]
        seq_len = X.shape[1]
        X = self.embedding(X) * tf.math.sqrt(
            tf.cast(self.num_hiddens, dtype=tf.float32))
        X = X + tf.cast(self.pos_encoding.P[:, pos_offset:pos_offset+seq_len, :],
                        dtype=X.dtype)
        X = self.pos_encoding.dropout(X, training=training)
        # 2 attention layers in decoder
        self._attention_weights = [[None] * len(self.blks) for _ in range(2)]
        for i, blk in enumerate(self.blks):
            X, state = blk(X, state, training=training)
            # Decoder self-attention weights
            self._attention_weights[0][i] = (
                blk.attention1.attention.attention_weights)
            # Encoder-decoder attention weights
            self._attention_weights[1][i] = (
                blk.attention2.attention.attention_weights)
        return self.dense(X), state

    @property
    def attention_weights(self):
        return self._attention_weights

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
key_size, query_size, value_size = 256, 256, 256
norm_shape = [2]
with d2l.try_gpu():
    encoder = TransformerEncoder(
        len(data.src_vocab), key_size, query_size, value_size, num_hiddens,
        norm_shape, ffn_num_hiddens, num_heads, num_blks, dropout)
    decoder = TransformerDecoder(
        len(data.tgt_vocab), key_size, query_size, value_size, num_hiddens,
        norm_shape, 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)
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(
    data.build(engs, fras), d2l.try_gpu(), 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 . => ["j'ai", 'perdu', '.'], bleu,1.000
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 = model.predict_step(
    data.build([engs[-1]], [fras[-1]]), d2l.try_gpu(), data.num_steps, True)
enc_attention_weights = d2l.concat(model.encoder.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] for step in dec_attention_weights
                            for attn in step
                            for blk in attn for head in blk]
dec_attention_weights_filled = tf.convert_to_tensor(
    np.asarray(pd.DataFrame(dec_attention_weights_2d).fillna(
        0.0).values).astype(np.float32))
dec_attention_weights = tf.reshape(dec_attention_weights_filled, shape=(
    -1, 2, num_blks, num_heads, data.num_steps))
dec_self_attention_weights, dec_inter_attention_weights = tf.transpose(
    dec_attention_weights, perm=(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.