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 mxnet as d2l
import math
from mxnet import autograd, init, np, npx
from mxnet.gluon import nn
import pandas as pd
npx.set_np()

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.Block):
    """The positionwise feed-forward network."""
    def __init__(self, ffn_num_hiddens, ffn_num_outputs):
        super().__init__()
        self.dense1 = nn.Dense(ffn_num_hiddens, flatten=False,
                               activation='relu')
        self.dense2 = nn.Dense(ffn_num_outputs, flatten=False)

    def forward(self, X):
        return self.dense2(self.dense1(X))

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

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

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()
ln.initialize()
bn = nn.BatchNorm()
bn.initialize()
X = d2l.tensor([[1, 2], [2, 3]])
# Compute mean and variance from X in the training mode
with autograd.record():
    print('layer norm:', ln(X), '\nbatch norm:', bn(X))

AddNorm

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

class AddNorm(nn.Block):
    """The residual connection followed by layer normalization."""
    def __init__(self, dropout):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.ln = nn.LayerNorm()

    def forward(self, X, Y):
        return self.ln(self.dropout(Y) + X)
add_norm = AddNorm(0.5)
add_norm.initialize()
shape = (2, 3, 4)
d2l.check_shape(add_norm(d2l.ones(shape), d2l.ones(shape)), 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.Block):
    """The Transformer encoder block."""
    def __init__(self, num_hiddens, ffn_num_hiddens, num_heads, dropout,
                 use_bias=False):
        super().__init__()
        self.attention = d2l.MultiHeadAttention(
            num_hiddens, num_heads, dropout, use_bias)
        self.addnorm1 = AddNorm(dropout)
        self.ffn = PositionWiseFFN(ffn_num_hiddens, num_hiddens)
        self.addnorm2 = AddNorm(dropout)

    def forward(self, X, valid_lens):
        Y = self.addnorm1(X, self.attention(X, X, X, valid_lens))
        return self.addnorm2(Y, self.ffn(Y))

Shape check:

X = d2l.ones((2, 100, 24))
valid_lens = d2l.tensor([3, 2])
encoder_blk = TransformerEncoderBlock(24, 48, 8, 0.5)
encoder_blk.initialize()
d2l.check_shape(encoder_blk(X, valid_lens), 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, num_hiddens, ffn_num_hiddens,
                 num_heads, num_blks, dropout, use_bias=False):
        super().__init__()
        self.num_hiddens = num_hiddens
        self.embedding = nn.Embedding(vocab_size, num_hiddens)
        self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)
        self.blks = nn.Sequential()
        for _ in range(num_blks):
            self.blks.add(TransformerEncoderBlock(
                num_hiddens, ffn_num_hiddens, num_heads, dropout, use_bias))
        self.initialize()

    def forward(self, X, valid_lens):
        # 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) * math.sqrt(self.num_hiddens))
        self.attention_weights = [None] * len(self.blks)
        for i, blk in enumerate(self.blks):
            X = blk(X, valid_lens)
            self.attention_weights[
                i] = blk.attention.attention.attention_weights
        return X
encoder = TransformerEncoder(200, 24, 48, 8, 2, 0.5)
d2l.check_shape(encoder(np.ones((2, 100)), valid_lens), (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.Block):
    # The i-th block in the Transformer decoder
    def __init__(self, num_hiddens, ffn_num_hiddens, num_heads, dropout, i):
        super().__init__()
        self.i = i
        self.attention1 = d2l.MultiHeadAttention(num_hiddens, num_heads,
                                                 dropout)
        self.addnorm1 = AddNorm(dropout)
        self.attention2 = d2l.MultiHeadAttention(num_hiddens, num_heads,
                                                 dropout)
        self.addnorm2 = AddNorm(dropout)
        self.ffn = PositionWiseFFN(ffn_num_hiddens, num_hiddens)
        self.addnorm3 = AddNorm(dropout)

    def forward(self, X, state):
        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 = np.concatenate((state[2][self.i], X), axis=1)
        state[2][self.i] = key_values

        if autograd.is_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 = np.tile(np.arange(1, num_steps + 1, ctx=X.ctx),
                                     (batch_size, 1))
        else:
            dec_valid_lens = None
        # Self-attention
        X2 = self.attention1(X, key_values, key_values, dec_valid_lens)
        Y = self.addnorm1(X, X2)
        # Encoder-decoder attention. Shape of enc_outputs:
        # (batch_size, num_steps, num_hiddens)
        Y2 = self.attention2(Y, enc_outputs, enc_outputs, enc_valid_lens)
        Z = self.addnorm2(Y, Y2)
        return self.addnorm3(Z, self.ffn(Z)), 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, 48, 8, 0.5, 0)
decoder_blk.initialize()
X = np.ones((2, 100, 24))
state = [encoder_blk(X, valid_lens), valid_lens, [None]]
d2l.check_shape(decoder_blk(X, state)[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, num_hiddens, ffn_num_hiddens, num_heads,
                 num_blks, dropout):
        super().__init__()
        self.num_hiddens = num_hiddens
        self.num_blks = num_blks
        self.embedding = nn.Embedding(vocab_size, num_hiddens)
        self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)
        self.blks = nn.Sequential()
        for i in range(num_blks):
            self.blks.add(TransformerDecoderBlock(
                num_hiddens, ffn_num_hiddens, num_heads, dropout, i))
        self.dense = nn.Dense(vocab_size, flatten=False)
        self.initialize()

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

    def forward(self, X, state):
        # 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) * math.sqrt(self.num_hiddens)
        seq_len = X.shape[1]
        X = X + self.pos_encoding.P[:, pos_offset:pos_offset+seq_len, :].as_in_ctx(X.ctx)
        X = self.pos_encoding.dropout(X)
        self._attention_weights = [[None] * len(self.blks) for _ in range (2)]
        for i, blk in enumerate(self.blks):
            X, state = blk(X, state)
            # 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
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(
    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}')

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 = [d2l.tensor(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 = d2l.reshape(dec_attention_weights_filled, (
    -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.