class TransformerEncoder(nnx.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, rngs=None):
rngs = nnx.Rngs(0) if rngs is None else rngs
init = nnx.initializers.normal(0.02)
self.token_emb = nnx.Embed(vocab_size, num_hiddens,
embedding_init=init, rngs=rngs)
self.pos_emb = nnx.Embed(max_len, num_hiddens, embedding_init=init,
rngs=rngs)
self.blks = nnx.List([
d2l.TransformerBlock(num_hiddens, num_heads, rngs=rngs)
for _ in range(num_blks)])
self.norm = nnx.RMSNorm(num_hiddens, rngs=rngs)
def __call__(self, X):
H = self.token_emb(X) + self.pos_emb(jnp.arange(X.shape[1]))
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.97/1.54/1.21What 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.61
loss at position 0: 2.11, at position 63: 2.20, interior mean: 1.20Flat 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.72, 'g' 0.05, 'f' 0.05
'the time trave_ler for so it will be con'... -> 'l' 0.96, ' ' 0.02, 'r' 0.01
'the time traveller for so it w_ll be con'... -> 'e' 0.46, 'i' 0.40, 'o' 0.09Only 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 'fainddapflek' -> target 'kelfpaddniaf'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(nnx.Module):
"""Pre-norm decoder block: causal self-attention, cross-attention,
FFN."""
def __init__(self, num_hiddens, num_heads, rngs=None):
rngs = nnx.Rngs(0) if rngs is None else rngs
self.norm1 = nnx.RMSNorm(num_hiddens, rngs=rngs)
self.norm2 = nnx.RMSNorm(num_hiddens, rngs=rngs)
self.norm3 = nnx.RMSNorm(num_hiddens, rngs=rngs)
self.self_attention = d2l.MultiHeadAttention(num_hiddens, num_heads,
dropout=0, rngs=rngs)
self.cross_attention = d2l.MultiHeadAttention(
num_hiddens, num_heads, dropout=0, rngs=rngs)
self.ffn = d2l.FeedForward(num_hiddens, rngs=rngs)
def __call__(self, X, enc):
B, T = X.shape[:2]
causal = jnp.tile(jnp.arange(1, T + 1), (B, 1))
Y = self.norm1(X)
X = X + self.self_attention(Y, Y, Y, causal)[0]
Y2, cross_weights = self.cross_attention(self.norm2(X), enc, enc,
None)
X = X + Y2
return X + self.ffn(self.norm3(X)), cross_weightsThree 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 = jnp.array([3, 5]), jnp.array([3, 4])
S, T_dec = int(src_len.max()), int(tgt_len.max())
src_valid = (jnp.arange(S)[None, :] < src_len[:, None])[:, None, :]
tgt_valid = (jnp.arange(T_dec)[None, :] < tgt_len[:, None])[:, None, :]
causal = (jnp.arange(T_dec)[None, :]
<= jnp.arange(T_dec)[:, None])[None] # (1, T, T)
B = len(src_len)
enc_self = jnp.broadcast_to(src_valid, (B, S, S)) # source padding
dec_self = causal & tgt_valid # (B, T, T): causal AND padding
cross = jnp.broadcast_to(src_valid, (B, T_dec, 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].astype(int))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.002/0.001/0.000/0.000exact match on 1000 fresh strings: 1.000
source 'mpajamgfccab' -> predicted 'baccfgmajapm'Reading the alignment
head-averaged argmax hits the true source position on 92% of rows; mean weight there 0.87The 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(nnx.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,
rngs=None):
rngs = nnx.Rngs(0) if rngs is None else rngs
self.latents = nnx.Param(
0.02 * jax.random.normal(rngs.params(),
(num_latents, num_hiddens)))
self.norm_q = nnx.RMSNorm(num_hiddens, rngs=rngs)
self.cross_attention = d2l.MultiHeadAttention(
num_hiddens, num_heads, dropout=0, rngs=rngs)
self.blks = nnx.List([
d2l.TransformerBlock(num_hiddens, num_heads, rngs=rngs)
for _ in range(num_blks)])
self.norm = nnx.RMSNorm(num_hiddens, rngs=rngs)
def __call__(self, X):
Z = jnp.broadcast_to(self.latents[...],
(X.shape[0],) + self.latents.shape)
Z = Z + self.cross_attention(self.norm_q(Z), X, X, None)[0] # 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)
perceiver.eval()
X = jax.random.normal(jax.random.key(0), (1, 4096, 256))
print('input', X.shape, '-> latent summary', perceiver(X).shape)input (1, 4096, 256) -> latent summary (1, 64, 256)The cost curve
N= 1024: self-attention 0.33 ms, perceiver 0.23 ms
N= 2048: self-attention 0.62 ms, perceiver 0.34 ms
N= 4096: self-attention 2.69 ms, perceiver 0.16 ms
N= 8192: self-attention 10.00 ms, perceiver 0.27 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.