Encoder-Decoder Transformer From Scratch

Why Build This

Karpathy's "Let's build GPT" makes decoder-only transformers make sense. But the 2017 paper that started all of this is a different architecture: an encoder reads the full source sentence (both directions), and a decoder generates the translation one token at a time, pulling info from the encoder via cross-attention.

I wanted to build that myself. encoder-decoder-transformer is the result. Mostly I was curious about cross-attention and how encoder/decoder blocks actually differ once you write them out.

Architecture Overview

Pure PyTorch, no HuggingFace. Pre-norm residuals like the paper:

Source (DE) → src_emb + pos_emb → [EncoderBlock × N] → memory
Target (EN) → tgt_emb + pos_emb → [DecoderBlock × N] → LayerNorm → lm_head → logits

Defaults: 6 encoder layers, 6 decoder layers, 512-dim embeddings, 8 heads, dropout 0.1, block size 64. All tweakable from the CLI.

Source and target share one embedding table (GPT-2 BPE vocab + PAD/BOS/EOS). Positional embeddings are learned up to block_size.

Encoder Block

Two sub-layers per block:

1. Multi-head self-attention, fully bidirectional (is_causal=False). Every token sees every other token in the German sentence.
2. Feed-forward, two linear layers with ReLU: 512 → 2048 → 512.

x = ln1(x + mhsa(x, is_causal=False))
x = ln2(x + ff(x))

Output is a tensor (B, T_src, C) called memory. Computed once per forward pass, then fed to every decoder layer.

Decoder Block

Three sub-layers here (one extra vs the encoder):

1. Masked self-attention, causal (is_causal=True). Token i only sees tokens 0..i. Lower-triangular mask before softmax.
2. Cross-attention. Queries from the decoder, keys/values from encoder memory. No causal mask, decoder can look at the full source at every step.
3. Feed-forward, same 4× expansion as encoder.

x = ln1(x + mhsa(x, is_causal=True))
x = ln2(x + mhca(x, memory, memory))
x = ln3(x + ff(x))

Cross-attention is the whole point of the decoder side. Q comes from the partial English output, K and V come from the fixed German encoding.

Unified Attention Head

One Head class handles both self-attention and cross-attention instead of duplicating code:

def forward(self, query, key, value, is_causal=False):
    q = self.query(query)   # (B, Tq, hs)
    k = self.key(key)       # (B, Tk, hs)
    v = self.value(value)   # (B, Tk, hs)
    att = (q @ k.transpose(-2, -1)) / sqrt(d_k)
    if is_causal:
        att = att.masked_fill(~lower_triangular_mask, -inf)
    out = softmax(att) @ v

Self-attention: pass the same tensor for Q, K, V. Cross-attention: decoder states as Q, encoder memory as K/V. Same math, just different inputs.

Tokenization and Data Pipeline

GPT-2 BPE via tiktoken, plus three special tokens:

For each German-English pair from WMT 2014, you get three sequences:

• src - encoded German, padded to batch max length
• tgt_in - [BOS] + english_tokens (what the decoder sees during training)
• tgt_out - english_tokens + [EOS] (what the model predicts)

Standard teacher forcing: during training the decoder gets the real prefix, not its own guesses. It learns to predict token t given the source and tokens 0..t-1.

Batches pad to the local max length in each batch, not the global max across the whole dataset. Saves a lot of wasted compute.

Loss Function

Plain cross-entropy over the target vocab:

loss = F.cross_entropy(
    logits.reshape(B * T, V),
    targets.reshape(B * T),
    ignore_index=pad_id,
)

Padding tokens are ignored in the loss. Everything runs in parallel during training, no autoregressive unrolling needed.

Training

WMT 2014 English-German CSVs, AdamW at 3e-4. Example run:

python train.py \
  --block-size 128 \
  --n-embd 256 \
  --n-layers 4 \
  --n-heads 4 \
  --dropout 0.2 \
  --batch-size 32 \
  --max-iters 100000 \
  --eval-interval 500

Evaluates every 500 steps over 100 batches, logs train/val loss, saves a checkpoint at the end with model weights, optimizer state, tokenizer info, and hyperparams. The CLI reads all of that back so you don't have to hardcode architecture details.

Inference: Greedy Decoding CLI

Load a checkpoint, type German, get English:

src = encode("Ich liebe maschinelles Lernen.")
tgt = [BOS_ID]

for step in range(max_new_tokens):
    logits, _ = model(src, tgt)
    next_token = argmax(logits[:, -1, :])
    tgt.append(next_token)
    if next_token == EOS_ID: break

Greedy decoding (argmax only, no sampling or beam search). Encoder runs once, decoder grows one token at a time. Rejects inputs longer than block_size.

$ python cli.py
>> Ich liebe maschinelles Lernen.
EN: I love machine learning.

Translation quality depends on how long you train and how big the model is. This is a learning project, not something you'd ship. But the full loop works: download data, train, load checkpoint, translate interactively.

Decoder-Only vs Encoder-Decoder

What I'd Do Next

• Beam search or nucleus sampling (greedy is pretty rough)
• LR warmup + cosine decay (I have a scheduler in my tiny-transformer repo I could port over)
• Actual BLEU/chrF eval on the val set
• Separate source/target embedding tables