class TransformerEncoder(nn.Module):
"""Bidirectional encoder: embeddings plus unmasked pre-norm blocks."""
def __init__(self, vocab_size, num_hiddens=128, num_heads=4, num_blks=4,
max_len=64):
super().__init__()
self.token_emb = nn.Embedding(vocab_size, num_hiddens)
self.pos_emb = nn.Embedding(max_len, num_hiddens)
for emb in (self.token_emb, self.pos_emb):
nn.init.normal_(emb.weight, std=0.02)
self.blks = nn.ModuleList([
d2l.TransformerBlock(num_hiddens, num_heads)
for _ in range(num_blks)])
self.norm = nn.RMSNorm(num_hiddens)
def forward(self, X):
H = self.token_emb(X) + self.pos_emb(
torch.arange(X.shape[1], device=X.device))
for blk in self.blks:
H = blk(H) # no valid_lens: every token attends everywhere
return self.norm(H)Encoders, Decoders, and Cross-Attention
Dive into Deep Learning · §11.4
Encoders, decoders, and cross-attention
three wirings of one block · predicting from both sides · an alignment you can verify · the latent bottleneck
Three wirings of one block
The block leaves two questions open: which positions may attend to which, and where keys and values come from.
- Encoder-only: no mask — no left-to-right generation, built to represent (BERT; the ViT is this wiring over patches).
- Decoder-only: the GPT of the previous sections — generation is the objective run forward.
- Encoder–decoder: causal decoder cross-attends into a bidirectional encoder (the 2017 original).
A bidirectional encoder in a dozen lines
One difference from CharLM: no valid_lens — nothing is masked.
The masked-token objective
Next-token prediction is trivial when the future is visible. Hide what you ask for:
\max \; \sum_{t \in \mathcal{M}} \log p\left(x_t \mid \mathbf{x}_{\setminus \mathcal{M}}\right)
masked loss at step 500/1000/2000: 1.92/1.35/1.08What the second side is worth
Window edges are one-sided by construction — position 63 is where a causal LM lives permanently. Bin the loss by position:
masked accuracy 0.65
loss at position 0: 2.06, at position 63: 2.08, interior mean: 1.10Flat interior at ~1.1 nats, roughly double at the one-sided edges.
Predictions you can read
'the time _raveller for so it will be con'... -> 't' 0.94, 'g' 0.02, 'f' 0.01
'the time trave_ler for so it will be con'... -> 'l' 0.99, ' ' 0.01, 'r' 0.00
'the time traveller for so it w_ll be con'... -> 'i' 0.63, 'e' 0.29, 'a' 0.06Only the right context pins down “trave_ler”.
Scaled up — subwords, sentence pairs, gigabytes — this is BERT (2018). ModernBERT (2024) refits the same wiring with this chapter’s block: RoPE, gated FFN, local–global attention, 8k context.
A task whose alignment we know
Machine translation gives no ground truth for its attention. Reversal does: target t must copy source n-1-t, and content only flows through cross-attention.
source 'ioaofaoebgaj' -> target 'jagbeoafoaoi'Random strings = unlimited fresh data — no memorization possible.
The decoder block: one more sublayer
Causal self-attention → cross-attention (queries: target stream; keys/values: encoder output) → FFN. Self-attention masked, cross-attention not — the source is fully known.
class DecoderBlock(nn.Module):
"""Pre-norm decoder block: causal self-attention, cross-attention,
FFN."""
def __init__(self, num_hiddens, num_heads):
super().__init__()
self.norm1 = nn.RMSNorm(num_hiddens)
self.norm2 = nn.RMSNorm(num_hiddens)
self.norm3 = nn.RMSNorm(num_hiddens)
self.self_attention = d2l.MultiHeadAttention(num_hiddens, num_heads,
dropout=0)
self.cross_attention = d2l.MultiHeadAttention(num_hiddens,
num_heads, dropout=0)
self.ffn = d2l.FeedForward(num_hiddens)
def forward(self, X, enc):
B, T = X.shape[:2]
causal = torch.arange(1, T + 1, device=X.device).repeat(B, 1)
Y = self.norm1(X)
X = X + self.self_attention(Y, Y, Y, causal)
X = X + self.cross_attention(self.norm2(X), enc, enc, None)
return X + self.ffn(self.norm3(X))Three masks, two primitives
A ragged real batch composes every mask from a padding mask and the causal triangle (§10.2): source padding for encoder self-attention, causal ∧ target padding for decoder self-attention, source padding again for cross-attention:
src_len, tgt_len = torch.tensor([3, 5]), torch.tensor([3, 4])
S, T_dec = int(src_len.max()), int(tgt_len.max())
src_valid = (torch.arange(S)[None, :] < src_len[:, None])[:, None, :]
tgt_valid = (torch.arange(T_dec)[None, :] < tgt_len[:, None])[:, None, :]
causal = (torch.arange(T_dec)[None, :]
<= torch.arange(T_dec)[:, None])[None] # (1, T, T)
enc_self = src_valid.expand(-1, S, -1) # (B, S, S): source padding
dec_self = causal & tgt_valid # (B, T, T): causal AND padding
cross = src_valid.expand(-1, T_dec, -1) # (B, T, S): source padding
for name, m in (('encoder self', enc_self), ('decoder self', dec_self),
('cross', cross)):
print(f'{name}-attention mask, sequence 0 (1 = may attend):')
print(m[0].int().numpy())encoder self-attention mask, sequence 0 (1 = may attend):
[[1 1 1 0 0]
[1 1 1 0 0]
[1 1 1 0 0]
[1 1 1 0 0]
[1 1 1 0 0]]
...
[1 1 1 0]]
cross-attention mask, sequence 0 (1 = may attend):
[[1 1 1 0 0]
[1 1 1 0 0]
[1 1 1 0 0]
[1 1 1 0 0]]Train it, decode it
One encoder block, one decoder block, 600 steps of on-line batches:
loss at step 100/200/400/600: 0.005/0.001/0.001/0.000exact match on 1000 fresh strings: 1.000
source 'acppaolbfbak' -> predicted 'kabfbloappca'Reading the alignment
head-averaged argmax hits the true source position on 100% of rows; mean weight there 0.88The anti-diagonal, learned — and verifiable, because the task fixes the answer. Deeper, more-headed models solve the task with delocalized maps (exercise): readable attention is the exception.
Cross-attention as interface
A query is just a vector — it can be a learned parameter. M latents read a length-N input:
O(MN) to read, O(M^2) to think — the N^2 term never appears (Perceiver, 2021).
A minimal Perceiver
class PerceiverEncoder(nn.Module):
"""M learned latents cross-attend into the input, then process among
themselves."""
def __init__(self, num_latents, num_hiddens, num_heads=4, num_blks=2):
super().__init__()
self.latents = nn.Parameter(
0.02 * torch.randn(num_latents, num_hiddens))
self.norm_q = nn.RMSNorm(num_hiddens)
self.cross_attention = d2l.MultiHeadAttention(num_hiddens,
num_heads, dropout=0)
self.blks = nn.ModuleList([
d2l.TransformerBlock(num_hiddens, num_heads)
for _ in range(num_blks)])
self.norm = nn.RMSNorm(num_hiddens)
def forward(self, X):
Z = self.latents.expand(X.shape[0], -1, -1)
Z = Z + self.cross_attention(self.norm_q(Z), X, X, None) # O(MN)
for blk in self.blks:
Z = blk(Z) # O(M^2)
return self.norm(Z)
perceiver = PerceiverEncoder(num_latents=64, num_hiddens=256).to(device)
X = torch.randn(1, 4096, 256, device=device)
print('input', tuple(X.shape), '-> latent summary',
tuple(perceiver(X).shape))input (1, 4096, 256) -> latent summary (1, 64, 256)The cost curve
N= 1024: self-attention 0.99 ms, perceiver 1.31 ms
N= 2048: self-attention 2.80 ms, perceiver 1.40 ms
N= 4096: self-attention 9.46 ms, perceiver 1.35 ms
N= 8192: self-attention 34.79 ms, perceiver 1.72 msSelf-attention: ~4× per doubling of N. Perceiver: barely moves; over an order of magnitude ahead by N = 8192. At short inputs the bottleneck buys nothing.
Which wiring when
- Encoder-only — the product is the representation: embeddings, retrieval, classification (BERT descendants, ModernBERT).
- Encoder–decoder — input fully known, own tower worth it: T5 text-to-text, Whisper speech recognition.
- Decoder-only — everything else: one stack, one objective, generation free, in-context learning included.
Perceiver descendants live on inside multimodal models: Flamingo’s resampler, BLIP-2’s Q-Former, DETR’s object queries.
Recap
- Three wirings, one block: bidirectional / causal / both + cross.
- Masked prediction trains encoders; both-sides context cuts the loss roughly in half vs. one side.
- Cross-attention verified: the reversal task’s anti-diagonal appears in the learned map.
- Learned queries turn attention into an interface: O(MN), flat cost curve, Perceiver → resampler → Q-Former.
- Encoders represent, encoder–decoders translate and transcribe, decoders do the rest.