Full Encoder-Decoder Transformer

Introduction

Over the last seven chapters we built every component of the Transformer piece by piece. Chapter 7 finished the encoder — a stack of $N$ identical self-attention blocks that turns a source sentence into a bank of contextualized vectors we will call the memory. Section 5 of this chapter finished the decoder — a stack of $N$ identical blocks that does masked self-attention, cross-attention against the memory, and feed-forward projection.

This section glues them together. We wrap the two stacks, plus the source & target embedding tables and the final vocabulary projection, inside a single Transformer module whose $\text{forward}(src, tgt)$ runs end-to-end in one call. This is exactly the Vaswani et al. (2017) architecture — no extras, no missing pieces.

What changed in 2017: before the Transformer, sequence-to-sequence models were encoder-decoder RNNs. Even attention-augmented variants (Bahdanau 2014, Luong 2015) still stepped through the sequence one token at a time. The Transformer replaces recurrence with pure attention + FFN, so encoder and decoder are batched matrix multiplies — trivially parallelizable on GPUs and the reason we can now train trillion-parameter models.

Architecture Overview

At the highest level, the full Transformer is a three-stage pipeline. Source tokens flow through the encoder once; the resulting memory is then reused by every decoder step.

1Source token ids (B, T_src)
2                            |
3                  +---------v---------+
4                  | src_embed * sqrt  |
5                  |  + positional enc |
6                  +---------+---------+
7                            |
8                            v
9                  +---------------------+
10                  |   Encoder (N x)     |     self-attn + FFN
11                  |  - MHSA  - FFN      |     Add & Norm inside
12                  |  stack: L1 -> ... L6|
13                  +---------+-----------+
14                            |
15                            v          memory: (B, T_src, d_model)
16  Target ids (B, T_tgt) ----+-------------------------+
17         |                                            |
18         v                                            |
19+---------------------+                               |
20| tgt_embed * sqrt    |                               |
21|  + positional enc   |                               |
22+---------+-----------+                               |
23          |                                           |
24          v                                           v
25+-----------------------------+          (K, V come from memory)
26|     Decoder (N x)           |
27|  - Masked MHSA              |
28|  - Cross-Attn (Q from self, |
29|    K/V from memory)         |
30|  - FFN                      |
31|  stack: L1 -> ... L6        |
32+-------------+---------------+
33              |
34              v  hidden: (B, T_tgt, d_model)
35   +----------+-----------+
36   |  output_proj (Linear)|   d_model -> tgt_vocab
37   +----------+-----------+
38              |
39              v
40         logits (B, T_tgt, tgt_vocab)

Three things are worth staring at:

1. The encoder runs exactly once per input sentence. During inference, its output (the memory) is computed and then reused across every decoder step.
2. The decoder's cross-attention gets its $K, V$ from memory and its $Q$ from the decoder's own self-attention output. That is how target tokens "look at" the source.
3. The final projection has shape $(d_{\text{model}}, V_{\text{tgt}})$, producing one logit per target vocab id at every position.

Forward Pass Equations

Written compactly, the end-to-end forward pass is just two lines. Let $\text{src} \in \mathbb{Z}^{B \times T_{src}}$ be source token ids and $\text{tgt} \in \mathbb{Z}^{B \times T_{tgt}}$ target token ids. Then

$\text{memory} = \text{Encoder}(\text{src}, \text{src\_mask})$, and $\text{logits} = \text{Decoder}(\text{tgt}, \text{memory}, \text{tgt\_mask}, \text{src\_mask})$.

Here $\text{memory} \in \mathbb{R}^{B \times T_{src} \times d_{model}}$ and $\text{logits} \in \mathbb{R}^{B \times T_{tgt} \times V_{tgt}}$. The source mask $\text{src\_mask}$ (a key-padding mask) is reused in both the encoder's self-attention and the decoder's cross-attention. The target mask $\text{tgt\_mask}$ is the $T_{tgt} \times T_{tgt}$ causal mask from §2.

Internally, Decoder expands to$h_0 = \text{Dropout}(\text{Embed}(\text{tgt})\sqrt{d_{model}} + PE)$,$h_\ell = \text{DecoderLayer}_\ell(h_{\ell-1}, \text{memory}, \text{tgt\_mask}, \text{src\_mask})$for $\ell = 1, \ldots, N$, and finally$\text{logits} = h_N W_{\text{out}}^{\top}$. The Encoder expands analogously (§5 of ch07).

Training vs Inference

Training: teacher forcing, one forward pass per batch

During training we have the full target sequence $y = (y_1, \ldots, y_{T_{tgt}})$ in hand. We feed $\text{tgt\_in} = [\langle sos \rangle, y_1, \ldots, y_{T_{tgt}-1}]$ into the decoder and ask it to predict $\text{tgt\_out} = [y_1, \ldots, y_{T_{tgt}-1}, \langle eos \rangle]$. The causal mask ensures that position $t$ cannot peek at positions $> t$. One forward pass computes logits for ALL target positions simultaneously, and the loss is the mean cross-entropy between logits and $\text{tgt\_out}$.

Why teacher forcing works: the causal mask turns what is morally $T_{tgt}$ separate autoregressive problems into one parallel matrix computation. Every target position sees only its valid left context, exactly as it would during inference — but training throughput is multiplied by $T_{tgt}$.

Inference: autoregressive loop

At test time we don't have $y$. We start with $\langle sos \rangle$ and generate one token per step until we hit $\langle eos \rangle$ or $t = T_{\max}$. Each step runs the FULL decoder on the partial sequence, takes the last position's logits, picks a next token, and appends it.

This naive loop is $O(T_{tgt}^3)$ in work and $O(T_{tgt}^2)$ in peak memory, because step $t$ recomputes attention for all previous positions. Chapter 9 introduces beam search, temperature / top-k / top-p sampling, and — critically — KV-caching, which stores the $K, V$ of previous positions so step $t$ only pays $O(T_{tgt})$ work.

Shared vs Separate Vocabularies

The Transformer has two embedding tables: $E_{src} \in \mathbb{R}^{V_{src} \times d_{model}}$ and $E_{tgt} \in \mathbb{R}^{V_{tgt} \times d_{model}}$, plus an output projection $W_{\text{out}} \in \mathbb{R}^{V_{tgt} \times d_{model}}$. You have two design axes:

• Shared source & target embedding ($E_{src} = E_{tgt}$): sensible when the two languages share a script and you train a joint BPE vocabulary over both corpora. Common for En↔De, En↔Fr in the original paper, and for code-completion models where input and output come from the same token space.
• Weight tying ($W_{\text{out}} = E_{tgt}$): ties the output projection to the target-embedding table (Press & Wolf 2017). Saves $V_{tgt} \cdot d_{model}$ parameters and usually improves perplexity, because both matrices encode the same "which direction in hidden space represents token $v$?" relation.
• Separate vocabularies: required when scripts or tokenizers differ (En↔Zh, En↔Ar). Here a shared BPE would waste capacity on tokens that never co-occur.

For the Multi30k translation project in chapter 13 we use separate source (de) and target (en) BPE vocabularies — $V_{src} \approx 7000$ and $V_{tgt} \approx 5500$. That is the setting the code below targets.

Plain-Python End-to-End

Before introducing the Transformer module, let's watch the shapes move by hand. We use the shared chapter config — $B=1, T_{src}=4, T_{tgt}=3, d_{model}=8, H=2, d_{ff}=16, N=2$, $V_{src}=12, V_{tgt}=10$, and $\text{torch.manual\_seed}(0)$. The actual printed values below were computed by running this code.

tensor([[11, 8, 2, 4]]) shape (1, 4)

tensor([[8, 7, 8]]) shape (1, 3)

(1, 4, 8). First row: src_x[0,0,:4] = [2.008, -3.342, -1.167, 3.733] (with seed=0).

(1, 4, 8). memory[0,0,:4] = [0.730, 0.329, -0.952, 1.499] (seed=0).

(1, 3, 8) — target-side token+position vectors.

[[0, -inf, -inf],
 [0,   0 , -inf],
 [0,   0 ,   0 ]]

(1, 3, 8). h[0,0,:4] = [1.348, -0.131, -1.344, -0.364] (seed=0).

(1, 3, 10). logits[0,0] = [0.361, -0.603, 0.058, 0.903, 0.147, -0.001, -0.121, -0.979, 0.720, 0.566].

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4import math
5
6# ---- Shared tiny config (see §4 of this chapter) ----
7torch.manual_seed(0)
8B, T_src, T_tgt = 1, 4, 3
9d_model, H, d_ff = 8, 2, 16
10N = 2
11src_vocab, tgt_vocab = 12, 10
12
13# Random source / target token ids
14src = torch.randint(1, src_vocab, (B, T_src))   # [[11, 8, 2, 4]]
15tgt = torch.randint(1, tgt_vocab, (B, T_tgt))   # [[8, 7, 8]]
16
17# Separate embedding tables (de vs en vocabularies)
18src_embed = nn.Embedding(src_vocab, d_model)
19tgt_embed = nn.Embedding(tgt_vocab, d_model)
20
21# Sinusoidal positional encoding (fixed, not learned)
22def make_pe(max_len, d):
23    pe = torch.zeros(max_len, d)
24    pos = torch.arange(0, max_len).float().unsqueeze(1)
25    div = torch.exp(torch.arange(0, d, 2).float() * (-math.log(10000.0) / d))
26    pe[:, 0::2] = torch.sin(pos * div)
27    pe[:, 1::2] = torch.cos(pos * div)
28    return pe
29
30pe = make_pe(16, d_model)
31
32# Stage 1: source side -> encoder
33src_x = src_embed(src) * math.sqrt(d_model) + pe[:T_src].unsqueeze(0)
34
35enc_layer = nn.TransformerEncoderLayer(
36    d_model=d_model, nhead=H, dim_feedforward=d_ff,
37    batch_first=True, activation='relu')
38encoder = nn.TransformerEncoder(enc_layer, num_layers=N)
39memory = encoder(src_x)                         # [1, 4, 8]
40
41# Stage 2: target side -> decoder (queries memory via cross-attn)
42tgt_x = tgt_embed(tgt) * math.sqrt(d_model) + pe[:T_tgt].unsqueeze(0)
43
44dec_layer = nn.TransformerDecoderLayer(
45    d_model=d_model, nhead=H, dim_feedforward=d_ff,
46    batch_first=True, activation='relu')
47decoder = nn.TransformerDecoder(dec_layer, num_layers=N)
48
49causal = torch.triu(torch.ones(T_tgt, T_tgt) * float('-inf'), diagonal=1)
50h = decoder(tgt_x, memory, tgt_mask=causal)     # [1, 3, 8]
51
52# Stage 3: project decoder hidden states to vocabulary logits
53W_out = nn.Linear(d_model, tgt_vocab, bias=False)
54logits = W_out(h)                               # [1, 3, 10]
55
56print("memory[0,0,:4] =", memory[0, 0, :4].tolist())
57print("logits[0,0]    =", [round(x, 3) for x in logits[0, 0].tolist()])

Running this prints:

1memory[0,0,:4] = [0.7302, 0.3295, -0.9525, 1.4986]
2logits[0,0]    = [0.361, -0.603, 0.058, 0.903, 0.147, -0.001, -0.121, -0.979, 0.720, 0.566]
3argmax logits[0,0] = 3     # predicted vocab id for the first target position

The shape chain is exactly what the architecture diagram promised: source ids $(1, 4)$ $\to$ memory $(1, 4, 8)$ $\to$ decoder hidden $(1, 3, 8)$ $\to$ logits $(1, 3, 10)$. Every sublayer we wrote from scratch in the previous chapters is hidden inside those two stack calls.

PyTorch Implementation

The Transformer module

Now the real thing — a single $\text{nn.Module}$ that owns every learnable piece.

1import math
2import torch
3import torch.nn as nn
4from typing import Optional
5
6
7class Transformer(nn.Module):
8    """Full Vaswani 2017 encoder-decoder Transformer.
9
10    Wires:
11      - source token embedding + positional encoding
12      - N encoder layers  -> memory
13      - target token embedding + positional encoding
14      - N decoder layers (queries memory)       -> hidden
15      - output projection to target vocab logits
16    """
17
18    def __init__(
19        self,
20        src_vocab: int,
21        tgt_vocab: int,
22        d_model: int = 512,
23        num_heads: int = 8,
24        num_layers: int = 6,
25        d_ff: int = 2048,
26        max_len: int = 5000,
27        dropout: float = 0.1,
28    ):
29        super().__init__()
30        self.d_model = d_model
31
32        # Separate embedding tables: src_vocab != tgt_vocab (translation case)
33        self.src_embed = nn.Embedding(src_vocab, d_model)
34        self.tgt_embed = nn.Embedding(tgt_vocab, d_model)
35
36        # Fixed sinusoidal positional encoding (shared across source & target)
37        self.register_buffer("pe", self._build_pe(max_len, d_model))
38        self.dropout = nn.Dropout(dropout)
39
40        # Encoder (reuses the block from ch07)
41        enc_layer = nn.TransformerEncoderLayer(
42            d_model, num_heads, d_ff, dropout, batch_first=True)
43        self.encoder = nn.TransformerEncoder(enc_layer, num_layers=num_layers)
44
45        # Decoder (reuses the block from §4/§5)
46        dec_layer = nn.TransformerDecoderLayer(
47            d_model, num_heads, d_ff, dropout, batch_first=True)
48        self.decoder = nn.TransformerDecoder(dec_layer, num_layers=num_layers)
49
50        # Output projection to target vocabulary
51        self.output_proj = nn.Linear(d_model, tgt_vocab, bias=False)
52
53    @staticmethod
54    def _build_pe(max_len: int, d_model: int) -> torch.Tensor:
55        pe = torch.zeros(max_len, d_model)
56        pos = torch.arange(0, max_len).float().unsqueeze(1)
57        div = torch.exp(torch.arange(0, d_model, 2).float()
58                        * (-math.log(10000.0) / d_model))
59        pe[:, 0::2] = torch.sin(pos * div)
60        pe[:, 1::2] = torch.cos(pos * div)
61        return pe.unsqueeze(0)  # (1, max_len, d_model)
62
63    def encode(self, src: torch.Tensor,
64               src_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
65        T = src.size(1)
66        x = self.src_embed(src) * math.sqrt(self.d_model) + self.pe[:, :T]
67        x = self.dropout(x)
68        return self.encoder(x, src_key_padding_mask=src_mask)
69
70    def decode(self, tgt: torch.Tensor, memory: torch.Tensor,
71               tgt_mask: Optional[torch.Tensor] = None,
72               memory_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
73        T = tgt.size(1)
74        x = self.tgt_embed(tgt) * math.sqrt(self.d_model) + self.pe[:, :T]
75        x = self.dropout(x)
76        return self.decoder(
77            x, memory,
78            tgt_mask=tgt_mask,
79            memory_key_padding_mask=memory_mask,
80        )
81
82    def forward(
83        self,
84        src: torch.Tensor,
85        tgt: torch.Tensor,
86        src_mask: Optional[torch.Tensor] = None,
87        tgt_mask: Optional[torch.Tensor] = None,
88    ) -> torch.Tensor:
89        memory = self.encode(src, src_mask=src_mask)
90        hidden = self.decode(tgt, memory,
91                             tgt_mask=tgt_mask,
92                             memory_mask=src_mask)
93        logits = self.output_proj(hidden)  # (B, T_tgt, tgt_vocab)
94        return logits

A couple of design choices worth calling out:

• PE as a buffer, not a Parameter. We want PE to travel with the module via $\text{.to(device)}$ and $\text{state\_dict}$, but we do NOT want the optimizer to update it.
• src_mask is fed twice: once as $\text{src\_key\_padding\_mask}$ for the encoder, again as $\text{memory\_key\_padding\_mask}$ for the decoder's cross-attention. Both places need to know which source positions are padding.
• No softmax in forward. We return raw logits so that $F.\text{cross\_entropy}$ can call $\log\text{-softmax}$ internally with better numerics.

Greedy generation

For the simplest possible inference, we pair $\text{forward}$ with a $\text{generate}$ method that loops over steps. This belongs on the same module so it can reuse $\text{self.encode}$ / $\text{self.decode}$.

1@torch.no_grad()
2    def generate(
3        self,
4        src: torch.Tensor,
5        max_len: int,
6        sos_id: int,
7        eos_id: int,
8        src_mask: Optional[torch.Tensor] = None,
9    ) -> torch.Tensor:
10        """Greedy autoregressive decoding. O(max_len * forward).
11
12        For beam search, top-k / top-p sampling, and KV-cache acceleration,
13        see chapter 9.
14        """
15        self.eval()
16        device = src.device
17        B = src.size(0)
18
19        # 1) Encode source ONCE. Memory is reused for every step.
20        memory = self.encode(src, src_mask=src_mask)
21
22        # 2) Start with <sos> for every batch element.
23        ys = torch.full((B, 1), sos_id, dtype=torch.long, device=device)
24
25        for _ in range(max_len - 1):
26            T = ys.size(1)
27            causal = torch.triu(
28                torch.ones(T, T, device=device) * float("-inf"),
29                diagonal=1,
30            )
31            hidden = self.decode(ys, memory,
32                                 tgt_mask=causal,
33                                 memory_mask=src_mask)
34            step_logits = self.output_proj(hidden[:, -1])   # (B, tgt_vocab)
35            next_id = step_logits.argmax(dim=-1, keepdim=True)  # (B, 1)
36            ys = torch.cat([ys, next_id], dim=1)
37
38            if (next_id == eos_id).all():
39                break
40
41        return ys

Forward reference: greedy decoding is the cheapest, least-diverse strategy. Chapter 9 covers beam search (keep top-$k$ partial hypotheses), sampling (temperature, top-$k$, top-$p$, typical decoding), and KV-caching (reuse past $K, V$ so each step is $O(T)$ instead of $O(T^2)$).

Parameter Count

Let's sanity-check the model size for the ch13 Multi30k config: $d_{model}=512, H=8, d_{ff}=2048, N=6, V_{src}=7000, V_{tgt}=5500$. We count each learnable piece.

About 53M parameters — in the same ballpark as the original Vaswani base model (65M), a bit smaller because our vocabularies are smaller than WMT's. For reference, Multi30k is a low-resource dataset and you'll usually train a smaller variant ($d_{model}=256, N=3$, closer to 10–15M parameters) to avoid overfitting.

Two memory-saving tricks cut this further:

• Weight tying $W_{\text{out}} = E_{tgt}$: removes the 2.82M output-projection parameters (the two matrices are shape $(V_{tgt}, d_{model})$ and $(d_{model}, V_{tgt})$ — each other's transpose).
• Shared source+target embedding (when vocabularies are merged via joint BPE): further saves $\min(V_{src}, V_{tgt}) \cdot d_{model}$ parameters.

Interactive Visualization

The visualization below animates the full pipeline on a toy German→English pair, \text{"Der Hund läuft"} \to \text{"The dog runs"}. Watch the source tokens flow through the encoder stack, deposit as memory blocks, and then get queried by each decoder step through glowing cross-attention lines. Hover a line to see its (illustrative) attention weight.

What to look for: (1) the encoder colors intensify as all four source tokens flow through; (2) memory blocks crystallize once the encoder finishes; (3) each target token's cross-attention lines concentrate on its aligned source word — "The"→"Der", "dog"→"Hund", "runs"→"läuft". These alignment patterns are the hallmark of a trained translation model.

In Modern Systems

Encoder-decoder: T5, BART, mBART, NLLB

The original Vaswani architecture is still the right default for sequence-to-sequence tasks where the full input is available up front and the output is a separate sequence. T5 (Raffel et al. 2019) reframes every NLP task as text-to-text and trains a shared encoder-decoder on a denoising objective. BART (Lewis et al. 2019) uses a similar shape for summarization and infilling. NLLB-200 (Meta 2022) scales mBART to 200-language translation — still an encoder-decoder, because the decoder wants to attend to the FULL source regardless of where it is in generation.

Decoder-only: GPT, Llama, Mistral, PaLM

Modern general-purpose LLMs drop the encoder entirely. They are single decoder stacks with masked self-attention only — no cross-attention, no separate vocabulary. Input and output are concatenated into one token stream, and training is plain next-token prediction on web-scale corpora. Inference splits into prefill (process the prompt in parallel) and decode (generate tokens one at a time with KV-caching). This architecture won for general LLMs because (1) training data is unlabeled text, no source/target split needed; (2) one stack is cheaper than two; (3) the prefill/decode split maps cleanly onto GPU systems.

Encoder-only: BERT, RoBERTa

Drop the decoder instead, keep the encoder. Useful for classification, retrieval, and sentence embeddings — tasks where you want a rich representation, not a generated sequence.

When each wins

Recent research has revived encoder-decoder models for their efficiency on supervised seq2seq tasks — Flan-T5, UL2 (Google 2022), and PaLM-2's mixture-of-encoders suggest the split is alive, even as decoder-only dominates general chat.

Summary

1. The full Transformer is three stages: source → memory (encoder, once), target + memory → hidden (decoder, one step per output token at inference), hidden → logits (linear).
2. Training is one parallel forward pass over the whole target using teacher forcing + causal mask. Inference is an autoregressive loop; KV-caching (ch09) makes it affordable.
3. Source and target vocabularies can be separate (our Multi30k setup) or shared (joint BPE, common for same-script pairs). Weight tying $W_{\text{out}} = E_{tgt}$ removes a full embedding-sized matrix of parameters for free.
4. Modern taxonomy: encoder-decoder for seq2seq, decoder-only for general LLMs, encoder-only for classification/embeddings. Each shape exists because it fits a specific family of tasks.

Chapter Summary

Over the six sections of this chapter, we built the decoder from scratch:

1. Section 1 — Decoder Architecture Overview. The three-sublayer pattern (masked self-attn + cross-attn + FFN) and why the decoder exists.
2. Section 2 — Causal Masking. The triangular $-\infty$ mask that lets teacher forcing pretend to be autoregressive.
3. Section 3 — Cross-Attention. Queries from the target side, keys and values from the encoder memory; how the decoder "looks at" the source.
4. Section 4 — Implementing TransformerDecoderLayer. The three sublayers wired together with Add&Norm, dropout placement, and the pre-norm vs post-norm debate.
5. Section 5 — Complete Transformer Decoder. Stacking $N$ decoder layers with target embedding, positional encoding, and output projection. Weight tying and depth-scaling.
6. Section 6 — Full Encoder-Decoder Transformer (this section). Glue encoder + decoder into one module with a unified forward pass and a greedy generate loop.

You now have every piece of the Vaswani 2017 architecture implemented end-to-end. The rest of the book builds on this foundation: chapter 9 adds smarter decoding, chapter 10 adds training infrastructure, and chapter 13 walks through training this exact model on Multi30k until it translates real German sentences.

Exercises

Easy

1. Modify the $\text{Transformer}$ constructor to accept separate $\text{num\_enc\_layers}$ and $\text{num\_dec\_layers}$. Re-count the parameters for the Multi30k config with 4 encoder layers and 6 decoder layers.

Medium

2. Implement weight tying by setting $\text{self.output\_proj.weight} = \text{self.tgt\_embed.weight}$. Verify that training still runs (both tensors must have the same shape and dtype). Measure the parameter-count reduction and sanity-check the initial logits distribution.

3. Extend $\text{generate}$ to return attention weights from the LAST decoder layer's cross-attention so you can visualize source↔target alignments after translation.

Hard

4. Replace the encoder-decoder with a decoder-only variant: concatenate $\text{src}$ and $\text{tgt}$ with a separator token and train next-token prediction on the concatenation. Measure translation quality on Multi30k — does the encoder pull its weight, or can a decoder-only beat it at this scale?

Next Section Preview

Chapter 9 picks up exactly where $\text{generate}$ left off. We look at greedy decoding, beam search, sampling (temperature, top-$k$, top-$p$), and KV-caching — the collection of tricks that turn the naive $O(T^3)$ generation loop above into the real-time autoregressive systems powering modern LLMs.