CourseTransformers: From Attention to Trained, Served and Interpreted Models · Module 1: The Road to Attention · part 9 of 111
Part 9 · Module 1: The Road to Attention

Lesson 3: Attention before transformers, rebuilt on one task

20 min read·13 Oct 2026
Module 1 · Lesson 3 · About 50 minutes · Needs: Lesson 1 (softmax, shapes) and Lesson 2 (the GRU, the copy task)

The situation

Sol 400, evening. The lab's first report is written: on the copy task the transformer reached 1.000 and the GRU 0.568. Kofi Mensah reads it and pushes back. "Attention is older than the transformer. Recurrent translators used it for years. So what did the transformer actually change?"

FOUR DESIGNS, ONE TASK 2014 fixed vector one state 2014 additive Bahdanau 2015 multiplicative Luong 2017 transformer attention only this lesson trains all four on the reversal task
Attention arrived in 2014 as a patch on a recurrent encoder-decoder; the 2017 transformer kept the attention and removed the recurrence (red marks the flawed design, blue the attention fixes).

You answer the lab's way: rebuild the history on one task, the shared reversal data, and measure each step.

A real case, first

In September 2014 Dzmitry Bahdanau, Kyunghyun Cho and Yoshua Bengio posted "Neural Machine Translation by Jointly Learning to Align and Translate" (arXiv:1409.0473, published at ICLR 2015). Translation models then squeezed a whole source sentence into one vector; the authors "conjecture that the use of a fixed-length vector is a bottleneck", and let the decoder "(soft-)search" the source instead.

BLEU, ENGLISH TO FRENCH (CITED) fixed vector 17.82 attention 26.75 attention, longer 28.45 Bahdanau, Cho, Bengio 2014, Table 1, models trained on sentences up to 50 words
In Bahdanau and colleagues' Table 1 the attention model scores 26.75 BLEU against 17.82 for the fixed-vector model, and 28.45 when trained longer (cited).

In their Figure 2 the basic encoder-decoder's score "dramatically drops as the length of the sentences increases"; the attention model "shows no performance deterioration even with sentences of length 50 or more". Table 1: 26.75 BLEU (a translation-overlap score) against 17.82.

ONE LONG SENTENCE, THE END OF IT (CITED) source: based on his status as a health care worker at a hospital fixed vector based on his state of health attention his status as a health care worker
The fixed-vector model was right up to "a medical center", then wrote "his state of health" for "his status as a health care worker"; the attention model kept the detail (cited, translations shown in English).

On one long sentence the fixed-vector model was correct up to "a medical center", then replaced "based on his status as a health care worker at a hospital" with "based on his state of health". It kept the gist and lost the detail.

What you will be able to do

  • Measure the encoder-decoder bottleneck as accuracy against length.
  • Implement Bahdanau's additive and Luong's multiplicative attention and compare them.
  • Explain attention as a soft, differentiable lookup, and show why a hard one cannot learn.
  • Read an alignment matrix, and test what it does and does not show.
  • State what the 2017 transformer removed and kept, and tell the three framings apart.

1. The encoder-decoder bottleneck (Essential)

Everything in this lesson fixes one design flaw, so measure the flaw first. Start with the torch setup cell and shared blocks 1 and 9, as in Lesson 2.

python
import torch
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"   # every lesson runs on "cpu"; a GPU is optional
torch.manual_seed(11)                                     # the course seed (see course-guide/06)
print("torch", torch.__version__, "| device", DEVICE, "| CPU threads", torch.get_num_threads())

Code walkthrough

  • The torch setup cell; the author's Mac printed torch 2.14.1 | device cpu | CPU threads 5.
python
import random
import torch

COURSE_DATE = "2026-10-09"   # versions, model names and prices checked on this day
SEED = 11                    # every generator in the course starts from this seed
SNAPSHOT_SOL = 400           # the corpus is the station record up to sol 399

def set_seed(seed=SEED):
    """Seed Python and PyTorch so a cell gives the same numbers every run."""
    random.seed(seed)
    torch.manual_seed(seed)

set_seed()
print("torch", torch.__version__, "| threads", torch.get_num_threads())

Code walkthrough

  • Shared block 1, as written: SEED = 11 and set_seed().
python
from typing import NamedTuple

class TaskData(NamedTuple):
    name: str
    vocab_size: int
    inputs: torch.Tensor    # (n, T) token ids
    targets: torch.Tensor   # (n, T) the next token at every position
    mask: torch.Tensor      # (n, T) 1 where the answer is scored, 0 elsewhere

def _shift(name, vocab, seq, scored_from):
    inputs, targets = seq[:, :-1], seq[:, 1:]
    mask = torch.zeros_like(inputs)
    mask[:, scored_from:] = 1
    return TaskData(name, vocab, inputs, targets, mask)

def make_task(name, n=4096, length=16, vocab=32, seed=SEED):
    """copy, reverse or induction. Token 0 is the separator; content tokens are 1 to vocab - 1."""
    g = torch.Generator().manual_seed(seed)
    body = torch.randint(1, vocab, (n, length), generator=g)
    sep = torch.zeros(n, 1, dtype=torch.long)
    if name == "copy":
        seq = torch.cat([body, sep, body], dim=1)
    elif name == "reverse":
        seq = torch.cat([body, sep, body.flip(1)], dim=1)
    elif name == "induction":                       # random tokens, then the same random tokens again
        seq = torch.cat([body, body], dim=1)
    return _shift(name, vocab, seq, scored_from=length)
TASKS = {name: make_task(name) for name in ["copy", "reverse", "induction"]}
for t in TASKS.values():
    print(t.name, tuple(t.inputs.shape), "scored positions per row:", int(t.mask[0].sum()))

Code walkthrough

  • Shared block 9, as written. This lesson uses reverse: 16 tokens, separator 0, the same tokens backwards.
ROW 0 OF THE SHARED REVERSE TASK source 7 19 19 22 13 17 11 29 10 1 24 16 9 16 20 4 decoder in 0 4 20 16 9 16 24 1 10 29 11 17 13 22 19 19 target 4 20 16 9 16 24 1 10 29 11 17 13 22 19 19 7 pos 0 pos 15 the decoder reads the target shifted right; 0 is the separator
The running example in encoder-decoder form: 16 source tokens (blue), the decoder's input (grey, the separator then the target shifted) and the target, the source reversed (green); token 7 travels from first to last (amber).

Lesson 2 fed each row to one model as a single stream. An encoder-decoder splits it: an encoder reads the 16 source tokens, and a decoder writes the 16 target tokens one at a time. The decoder's input is the target moved one place right, starting with the separator; this is called teacher forcing.

python
def pairs(task, L):
    """Split a shared task row into encoder-decoder form: source, decoder input, decoder target."""
    return task.inputs[:, :L], task.inputs[:, L:], task.targets[:, L:]

src, dec_in, tgt = pairs(TASKS["reverse"], 16)
print("source     :", src[0].tolist())
print("decoder in :", dec_in[0].tolist())
print("target     :", tgt[0].tolist())
print("target is the source reversed, every row:", torch.equal(tgt, src.flip(1)))

Code walkthrough

  • pairs slices the shared tensors: the first 16 inputs are the source, the rest the decoder's input, the matching targets what it must write.
  • Row 0, the running example, holds the 16 tokens Lesson 2 copied. The last line prints True.
ENCODER: A STATE AFTER EVERY TOKEN h1 h2 h3 h4 ... h16 reads 7 reads 19 reads 19 reads 22 reads 4 64 numbers all that passes decoder writes 16 tokens h1 to h15 thrown away
In the basic encoder-decoder only the encoder's last state, 64 numbers here, reaches the decoder (red: the bottleneck); the other 15 states are discarded.

In the 2014 design the encoder's GRU leaves a state after every token, but only the last one, a context vector of 64 numbers here, is handed to the decoder. Whatever the decoder needs about token 7 at position 0 must survive 15 more updates inside that vector. That is the bottleneck.

python
import torch.nn as nn
import torch.nn.functional as F

class Encoder(nn.Module):
    """Embed the source and run a GRU over it. Returns every state H and the last state h."""
    def __init__(self, vocab=32, d=64):
        super().__init__()
        self.embed, self.gru = nn.Embedding(vocab, d), nn.GRU(d, d, batch_first=True)

    def forward(self, src):
        H, h = self.gru(self.embed(src))          # H (B, S, d): one state per source token; h (1, B, d)
        return H, h

class Fixed(nn.Module):
    """The basic encoder-decoder: the decoder starts from the encoder's last state, then never looks back."""
    def __init__(self, vocab=32, d=64):
        super().__init__()
        self.encoder, self.embed = Encoder(vocab, d), nn.Embedding(vocab, d)
        self.gru, self.head = nn.GRU(d, d, batch_first=True), nn.Linear(d, vocab)

    def forward(self, src, dec_in):
        H, h = self.encoder(src)                  # H is computed and thrown away: only h goes on
        out, _ = self.gru(self.embed(dec_in), h)  # (B, T, d)
        return self.head(out)                     # (B, T, vocab)

Code walkthrough

  • Encoder returns H, all 16 states, and h, the last one; every model here shares it.
  • Fixed starts the decoder GRU from h and ignores H. head scores all 32 tokens at each position.
python
import time
N_VALID = 512                                     # the last 512 rows are held out, as in Lesson 2

def score(model, task, L):
    src, dec_in, tgt = (t[-N_VALID:] for t in pairs(task, L))
    model.eval()
    with torch.no_grad():
        right = (model(src, dec_in).argmax(-1) == tgt).float().mean().item()
    model.train()
    return right

def train(model, task, L, steps=600, lr=3e-3):
    set_seed(); g = torch.Generator().manual_seed(SEED)
    src, dec_in, tgt = pairs(task, L)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    t0 = time.perf_counter()
    for _ in range(steps):
        i = torch.randint(0, len(src) - N_VALID, (64,), generator=g)
        loss = F.cross_entropy(model(src[i], dec_in[i]).flatten(0, 1), tgt[i].flatten())
        opt.zero_grad(); loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step()
    return score(model, task, L), time.perf_counter() - t0

Code walkthrough

  • score is token accuracy on the 512 held-out rows.
  • train is Lesson 2's recipe (AdamW, seeded batches of 64, clipping at 1.0) and returns accuracy and seconds.
python
LENGTHS = [8, 16, 32]
REV = {L: TASKS["reverse"] if L == 16 else make_task("reverse", length=L) for L in LENGTHS}  # 8 and 32: added here
count = lambda m: sum(p.numel() for p in m.parameters())
RESULTS, MODELS = {}, {}

def run(name, make):
    for L in LENGTHS:
        set_seed(); model = make()
        acc, secs = train(model, REV[L], L)
        RESULTS[name, L], MODELS[name, L] = acc, model
        print(f"{name:11} length {L:2}  accuracy {acc:.3f}  ({secs:4.1f} s, {count(model):,} parameters)")

run("fixed", Fixed)
set_seed(); wide = Fixed(d=88)
print(f"fixed, d=88 length 32  accuracy {train(wide, REV[32], 32)[0]:.3f}  ({count(wide):,} parameters)")

Code walkthrough

  • REV holds reversal at three lengths, made with the shared make_task; TASKS stay unchanged.
  • run trains a fresh model per length from the same seed and keeps results and models for later.
  • The last lines train a wider fixed model: 102,464 parameters, more than any attention model below.
FIXED MODEL, TOKEN ACCURACY (MEASURED HERE) length 8 0.733 length 16 0.456 length 32 0.270 32, wider model 0.344 chance 0.032; wider model has 102,464 parameters
Measured here: the fixed model falls from 0.733 to 0.270 as the source grows from 8 to 32 tokens, and a wider one (grey) reaches only 0.344.

Measured here: 0.733 at length 8, 0.456 at 16, 0.270 at 32 (chance is 0.032). The wider model reaches only 0.344: parameters do not remove the bottleneck. This is Bahdanau's Figure 2 in miniature. Each run took 1 to 4 seconds on the author's Mac; Colab is slower.

2. Bahdanau additive attention (Recommended)

Bahdanau's fix keeps every encoder state and lets the decoder choose among them at each step. It is the first attention mechanism in this course.

BAHDANAU'S SCORE FOR ONE SOURCE POSITION decoder state s source state h_j W s U h_j add tanh dot with v one score the two inputs meet in a sum: additive
Bahdanau's alignment model scores each source state: project it and the decoder state, add, apply tanh, dot with a learned v (amber: a learned, uncertain score).

Before writing each token, the decoder compares its current state s with every source state h_j. Bahdanau's alignment model projects both, adds them, applies tanh and takes a dot product with a learned vector v: score = v · tanh(W s + U h_j). Because the two inputs meet in a sum, this is called additive attention.

ONE DECODER STEP 16 scores softmax 16 weights weighted sum context GRU step previous token + context next state s loops 16 times
Each step softmaxes 16 scores into weights, sums the source states with them into a fresh context, and feeds it to the next recurrent step, so the loop runs once per output token.

Softmax turns the 16 scores into attention weights that sum to 1. The weighted sum of the source states is that step's context, and it goes into the GRU together with the previous token. Every step gets a fresh context, so no single vector has to hold everything.

python
class Bahdanau(nn.Module):
    """Additive attention: score = v . tanh(W s + U h_j), recomputed before every output step."""
    def __init__(self, vocab=32, d=64):
        super().__init__()
        self.encoder, self.embed = Encoder(vocab, d), nn.Embedding(vocab, d)
        self.W, self.U = nn.Linear(d, d, bias=False), nn.Linear(d, d, bias=False)
        self.v = nn.Linear(d, 1, bias=False)
        self.cell, self.head = nn.GRUCell(2 * d, d), nn.Linear(d, vocab)

    def forward(self, src, dec_in):
        H, h = self.encoder(src)
        s, UH, y = h[0], self.U(H), self.embed(dec_in)   # U h_j does not change between steps: computed once
        outs, weights = [], []
        for t in range(y.shape[1]):                      # one output step at a time
            e = self.v(torch.tanh(self.W(s).unsqueeze(1) + UH)).squeeze(-1)   # (B, S) scores
            a = e.softmax(-1)                            # (B, S) weights that sum to 1
            c = (a.unsqueeze(1) @ H).squeeze(1)          # (B, d) context: weighted sum of source states
            s = self.cell(torch.cat([y[:, t], c], -1), s)   # the context goes into the recurrence
            outs.append(s); weights.append(a)
        self.align = torch.stack(weights, 1)             # (B, T, S), kept for reading later
        return self.head(torch.stack(outs, 1))

run("bahdanau", Bahdanau)

Code walkthrough

  • W, U and v are the alignment model. UH is computed once, as the paper's appendix suggests.
  • W(s).unsqueeze(1) + UH adds the decoder state to all 16 source states; tanh and v make each sum a score.
  • nn.GRUCell runs one recurrent step on the previous token and the context. self.align keeps the weights.
  • The paper's encoder also read backwards (bidirectional); this one reads forwards only.
TOKEN ACCURACY (MEASURED HERE) 0.73 0.97 length 8 0.46 0.96 length 16 0.27 0.88 length 32 fixed Bahdanau
Measured here: with additive attention accuracy stays high as the source grows (0.97, 0.96, 0.88) while the fixed model falls to 0.27.

Measured here: 0.967, 0.959 and 0.882 at lengths 8, 16 and 32, against the fixed model's 0.733, 0.456 and 0.270. The extra 20,544 parameters are not the reason: the wider fixed model had more. Access to every state is.

3. Luong multiplicative attention (Recommended)

A year later Luong, Pham and Manning simplified the score and moved it. The move is what made attention cheap.

BAHDANAU: ATTEND, THEN STEP LUONG: STEP, THEN ATTEND previous state weights, context GRU step output GRU step current state weights, context combine, output context inside the loop loop finishes first
Bahdanau feeds the context into each recurrent step, so attention sits inside the loop (red); Luong attends after the step, so the whole decoder can run first (green).

In "Effective Approaches to Attention-based Neural Machine Translation" (2015) they score with the decoder's state after its recurrent step, not before, and mix the context in afterwards. They say their path "is simpler": state, then weights, then context, then output. Attention no longer feeds the recurrence.

MULTIPLICATIVE: ALL SCORES IN ONE PRODUCT S (16, 64) x Wa (64, 64) x H transposed (64, 16) = scores (16, 16) ADDITIVE: A TENSOR FIRST (16, 16, 64) sums then tanh on each 16,384 tanh calls per row of the batch
Luong's general score for every step and source position is one chain of matrix products; additive attention must build a (16, 16, 64) tensor and apply tanh to all 16,384 entries (arithmetic).

Their scores are multiplicative: "dot" is s · h_j, and "general" is s · W h_j, with no tanh. Stack all 16 decoder states and every score is one matrix product, (16, 64) by (64, 64) by (64, 16). Additive attention must build a (16, 16, 64) tensor and apply tanh to all 16,384 entries (arithmetic).

python
class Luong(nn.Module):
    """Multiplicative ('general') attention: score = s_t . W h_j. Run the decoder first, then attend once."""
    def __init__(self, vocab=32, d=64):
        super().__init__()
        self.encoder, self.embed = Encoder(vocab, d), nn.Embedding(vocab, d)
        self.gru, self.Wa = nn.GRU(d, d, batch_first=True), nn.Linear(d, d, bias=False)
        self.combine, self.head = nn.Linear(2 * d, d), nn.Linear(d, vocab)

    def forward(self, src, dec_in, force=None):
        H, h = self.encoder(src)
        S, _ = self.gru(self.embed(dec_in), h)          # (B, T, d): every decoder state, no attention inside
        scores = S @ self.Wa(H).transpose(1, 2)         # (B, T, S): every score in one matrix product
        a = scores.softmax(-1) if force is None else force   # force: used in section 5
        self.align = a
        c = a @ H                                       # (B, T, d): every context in one matrix product
        return self.head(torch.tanh(self.combine(torch.cat([S, c], -1))))

run("luong", Luong)
print("length  " + "".join(f"{n:>10}" for n in ["fixed", "bahdanau", "luong"]))
for L in LENGTHS:
    print(f"{L:6}  " + "".join(f"{RESULTS[n, L]:10.3f}" for n in ["fixed", "bahdanau", "luong"]))

Code walkthrough

  • The decoder GRU runs over the whole target in one call: attention no longer feeds it.
  • S @ self.Wa(H).transpose(1, 2) is every general score at once, shape (B, 16, 16).
  • combine mixes state and context after the recurrence. force lets section 5 overwrite the weights.
TOKEN ACCURACY AND TIME (MEASURED HERE) fixed Bahdanau Luong length 8 0.733 0.967 0.990 length 16 0.456 0.959 0.931 length 32 0.270 0.882 0.780 seconds at 32 3.7 to 3.8 12.5 to 14.6 4.2 to 6.1 author's Mac, 600 steps, one seed; seconds over three runs
Measured here: both attention models beat the bottleneck at every length, and at length 32 Luong trains in 4.2 to 6.1 seconds against Bahdanau's 12.5 to 14.6, because its scores leave the loop.

Measured here: Luong scored 0.990, 0.931 and 0.780. At length 32 it trained in 4.2 to 6.1 seconds against Bahdanau's 12.5 to 14.6 (three runs). Bahdanau was more accurate at 32 on this one seed; Module 5 adds the noise floor needed to call that a difference. The speed gap is structural: one product against a 32-step loop.

4. Attention as soft differentiable lookup (Essential)

Strip away translation and attention is a lookup table that can be trained. Module 2 builds on this view, so check it on small numbers.

HARD LOOKUP SOFT LOOKUP query must equal a key one key matches exactly its value: 10 value 10 weight 0.500 value 20 weight 0.111 value 30 weight 0.389 blend: 18.90
A dictionary returns one value for an exact key; soft lookup scores every key, softmaxes the scores into weights (amber) and returns their blend of the values, 18.90.

A Python dictionary is a hard lookup: an exact key returns one value. Attention scores the query against every key, turns the scores into weights with softmax, and returns the weighted blend of the values. Here the weights are 0.5, 0.111 and 0.389, and the answer is 18.90: mostly the first value, partly the third. This is soft lookup.

python
keys = torch.tensor([[1.0, 0.0], [0.0, 1.0], [0.7, 0.7]])   # three keys
values = torch.tensor([10.0, 20.0, 30.0])                    # what each key hands back
query = torch.tensor([2.0, 0.5], requires_grad=True)         # what we are looking for

scores = keys @ query                                        # (3,) one score per key
weights = scores.softmax(-1)
soft = weights @ values                                      # a blend of every value
hard = values[scores.argmax()]                               # the single best key's value
print("scores", [round(s, 2) for s in scores.tolist()], " weights", [round(w, 3) for w in weights.tolist()])
print(f"soft lookup {soft.item():.2f}   hard lookup {hard.item():.0f}")
soft.backward()
print("gradient reaching the query (soft):", [round(g, 3) for g in query.grad.tolist()])
print("can the hard lookup pass a gradient back?", hard.requires_grad)

Code walkthrough

  • keys @ query gives scores 2.0, 0.5 and 1.75; hard takes the top key's value, 10, as a dictionary would.
  • soft.backward() sends the query a gradient, [-1.419, 3.147]. hard.requires_grad is False: argmax picks an index, and an index has no slope.
FOLLOWING THE GRADIENT BACK TO THE QUERY soft output softmax weights query query gets -1.419, 3.147 hard output argmax index query query gets nothing arrows point backwards: the path the gradient takes
Softmax passes a gradient back to the query, here -1.419 and 3.147 (green); argmax returns an index, which has no slope, so the query gets nothing (red).

That gradient is the point. A soft lookup is differentiable: a small change to the query changes the output a little, so training can learn what the query, keys and values should be. In Bahdanau's words, the model "directly computes a soft alignment, which allows the gradient of the cost function to be backpropagated through".

Break it: replace softmax with argmax

python
class HardLuong(Luong):
    """Luong, but each step takes only the best-scoring source state: argmax instead of softmax."""
    def forward(self, src, dec_in):
        H, h = self.encoder(src)
        S, _ = self.gru(self.embed(dec_in), h)
        scores = S @ self.Wa(H).transpose(1, 2)
        a = F.one_hot(scores.argmax(-1), H.shape[1]).float()   # 1 for the best position, 0 elsewhere
        return self.head(torch.tanh(self.combine(torch.cat([S, a @ H], -1))))

set_seed(); hard_model = HardLuong()
acc, _ = train(hard_model, REV[16], 16)
print(f"hard attention, length 16: {acc:.3f}   soft: {RESULTS['luong', 16]:.3f}   fixed: {RESULTS['fixed', 16]:.3f}")
print("gradient that reached the scoring matrix Wa:", hard_model.Wa.weight.grad)

Code walkthrough

  • HardLuong changes one line of Luong: a one-hot of the argmax replaces softmax.
  • After training, Wa.weight.grad prints None: no gradient ever reached the scoring weights.
LENGTH 16, TOKEN ACCURACY (MEASURED HERE) soft attention 0.931 hard attention 0.486 fixed, no attention 0.456
Measured here: swapping softmax for argmax drops accuracy from 0.931 to 0.486, barely above having no attention at all (0.456).

Measured here: 0.486, near the fixed model's 0.456 and far below 0.931. The scoring weights kept their random start, so each step looked somewhere arbitrary. Softness is what lets a model learn where to look.

5. Reading an alignment matrix, and its limits (Recommended)

Attention weights are tempting to read as an explanation, as Bahdanau's Figure 3 did. Read one yourself, then test it.

LUONG ALIGNMENT, RUNNING EXAMPLE 7 19 19 22 13 17 11 29 10 1 24 16 9 16 20 4 source tokens, positions 0 to 15 step 0 writes 4 step 1 writes 20 step 2 writes 16 step 3 writes 9 step 4 writes 16 step 5 writes 24 step 6 writes 1 step 7 writes 10 step 8 writes 29 step 9 writes 11 step 10 writes 17 step 11 writes 13 step 12 writes 22 step 13 writes 19 step 14 writes 19 step 15 writes 7 green outline: where the answer sits
The learned alignment matrix for the running example is a clean anti-diagonal: reversal, learned from data. Darker blue is more weight; the green outlines mark where each answer actually sits.

Stack each step's weights as a row and you have an alignment matrix: rows are output steps, columns source positions. For the running example the Luong model draws a near-perfect anti-diagonal: reversal, learned from data. Bahdanau's paper shows the same kind of picture for French, where "zone" aligns with "Area" by jumping over two words.

python
model = MODELS["luong", 16].eval()
src, dec_in, tgt = pairs(REV[16], 16)
with torch.no_grad():
    model(src[:1], dec_in[:1])
A = model.align[0]                                   # (16 output steps, 16 source positions)
print("source:", src[0].tolist())
for t in range(16):
    print(f"step {t:2} writes {tgt[0, t].item():2} (source position {15 - t:2})"
          f"   peak at position {A[t].argmax().item():2}, weight {A[t].max().item():.2f}")

Code walkthrough

  • model.align[0] is row 0's (16, 16) matrix. Each line shows what a step writes, where that token sits, and where the weight peaks.
SOURCE POSITIONS 0 TO 4 pos 0 7 pos 1 19 pos 2 19 pos 3 22 pos 4 13 step 14 reads 19, writes 19 0.00 0.26 0.74 0.00 0.00 step 15 reads 19, writes 7 0.00 0.76 0.24 0.00 0.00 weight follows the 19 just read, not the answer (green)
Each step's weight (amber) lands on the token it just read, not on the answer it writes (green outline); the source holds two 19s, so the weight splits 0.74 and 0.26 between them.

Look closely. Step 15 writes 7, from position 0, but its peak is at position 1. Every step from 1 on peaks one place to the right of its answer: on the token the decoder just read. When that token appears twice, the weight splits: 0.74 and 0.26 across the two 19s.

python
src, dec_in, tgt = (t[-N_VALID:] for t in pairs(REV[16], 16))
answer_pos = torch.arange(15, -1, -1)                # step t writes source position 15 - t
read_pos = (answer_pos + 1).clamp(max=15)            # the position of the token the decoder just read
with torch.no_grad():
    model(src, dec_in)
    peaks = model.align.argmax(-1)                   # (512, 16) peak position per step
    print("peak on the answer's position :", round((peaks == answer_pos).float().mean().item(), 3))
    print("peak on the token just read   :", round((peaks == read_pos).float().mean().item(), 3))
    for name, pos in [("as learned", None), ("forced onto the answer", answer_pos), ("forced onto the peak", read_pos)]:
        force = None if pos is None else F.one_hot(pos, 16).float().expand(len(src), -1, -1)
        acc = (model(src, dec_in, force=force).argmax(-1) == tgt).float().mean().item()
        print(f"attention {name:23} accuracy {acc:.3f}")

Code walkthrough

  • answer_pos and read_pos are two readings of the matrix; the first prints count how often the peak matches each, over 512 held-out rows.
  • The loop overrides the weights with force, all on the answer or all on the peak: an intervention. Change the attention, measure the output.
LUONG, LENGTH 16, ACCURACY (MEASURED HERE) as learned 0.931 forced onto answer 0.685 forced onto peak 1.000
Measured here: forcing all attention onto the answer's position drops accuracy to 0.685, while forcing it onto the learned peak gives 1.000, so the peak is what the model uses.

Measured here: the peak sits on the answer 6.4% of the time and on the token just read 92.8%. Forcing attention onto the answer drops accuracy to 0.685; onto the peak, it rises to 1.000. The model finds the token it just read, and the GRU state there still carries its left neighbour.

An alignment shows where the query matched, not where the answer came from. Jain and Wallace ("Attention is not Explanation", 2019) found across many language tasks that learned weights "are frequently uncorrelated with gradient-based measures of feature importance". Module 14 returns with tools that test cause.

6. What Attention Is All You Need removed (Essential)

You now have the 2016 standard: recurrent encoder, recurrent decoder, attention between. The 2017 change is clearest against it.

REMOVED KEPT ADDED recurrent encoder recurrent decoder convolutions encoder-decoder cross-attention dot-product scores self-attention positions causal mask Vaswani et al., 2017
The 2017 transformer removed the recurrent parts (red), kept the encoder-decoder shape and its attention (blue), and had to add self-attention, positions and a causal mask (green).

Vaswani and colleagues' abstract describes a model "based solely on attention mechanisms, dispensing with recurrence and convolutions entirely". They kept the encoder-decoder shape, and their encoder-decoder attention "mimics the typical encoder-decoder attention mechanisms" of the models above. They kept multiplicative scores, adding a scaling factor (Module 2), and replaced both GRUs with self-attention: a sequence attending to itself.

python
class TinyTransformer(nn.Module):
    """Vaswani et al.'s shape, one layer each: self-attention replaces both GRUs; cross-attention stays."""
    def __init__(self, vocab=32, d=48, heads=4, context=64):
        super().__init__()
        self.src_embed, self.tgt_embed = nn.Embedding(vocab, d), nn.Embedding(vocab, d)
        self.pos = nn.Embedding(context, d)            # order must now be added by hand
        enc = nn.TransformerEncoderLayer(d, heads, 4 * d, dropout=0.0, batch_first=True, norm_first=True)
        dec = nn.TransformerDecoderLayer(d, heads, 4 * d, dropout=0.0, batch_first=True, norm_first=True)
        self.encoder = nn.TransformerEncoder(enc, 1, enable_nested_tensor=False)
        self.decoder = nn.TransformerDecoder(dec, 1)
        self.head = nn.Linear(d, vocab)

    def encode(self, src):                             # self-attention over the source, no mask
        return self.encoder(self.src_embed(src) + self.pos(torch.arange(src.shape[1])))

    def forward(self, src, dec_in):
        T = dec_in.shape[1]
        causal = nn.Transformer.generate_square_subsequent_mask(T)   # no looking ahead in the target
        y = self.tgt_embed(dec_in) + self.pos(torch.arange(T))
        return self.head(self.decoder(y, self.encode(src), tgt_mask=causal, tgt_is_causal=True))

run("transformer", TinyTransformer)

Code walkthrough

  • self.pos adds learned positions: without recurrence nothing else gives order (Module 4).
  • The encoder layer is self-attention plus a feed-forward network. The decoder layer adds masked self-attention and cross-attention to the encoder's output, the heir of Bahdanau's context.
  • Width 48 gives 73,760 parameters, between Luong and Bahdanau. Modules 2 and 3 build these layers by hand.
LENGTH 32, TOKEN ACCURACY (MEASURED HERE) fixed 0.270 Bahdanau 0.882 Luong 0.780 transformer 1.000 56,096 / 76,640 / 68,448 / 73,760 parameters
Measured here at length 32: the one-layer transformer reaches 1.000 with 73,760 parameters, where every recurrent design falls short.

Measured here: 1.000 at lengths 8, 16 and 32, in 1.6 to 6.1 seconds. No recurrent step is left: Lesson 2's sequential dependency and long gradient paths went with the GRUs. Three new jobs arrived, self-attention, the causal mask and positions, and Modules 2 to 4 are those jobs.

7. Encoder, decoder and encoder-decoder framings (Essential)

Every transformer you will meet uses one of three framings. They share blocks and differ in who may attend to whom, so the framing tells you what a model can do.

encoder-only every position sees all decoder-only earlier positions only encoder-decoder source block, causal target rows attend to columns; amber: cross-attention to the source
The three framings differ only in who may attend to whom: everything (encoder-only), earlier positions (decoder-only), or a full source plus a causal target that reads the source through cross-attention (amber).

The reversal task fits all three. Encoder-only: every position sees every source token at once, so position i can output token 15 - i in one pass. Decoder-only: Lesson 2's single stream, each position seeing only earlier ones. Encoder-decoder: this lesson's split, with cross-attention between the halves.

python
model = MODELS["transformer", 16].eval()
s, d_in = (t[:1] for t in pairs(REV[16], 16)[:2])    # the running example

def change_last(x):
    x = x.clone(); x[0, -1] = 1 + x[0, -1] % 31      # a different token in the last position
    return x

with torch.no_grad():
    enc = (model.encode(change_last(s)) - model.encode(s)).abs().amax(-1)[0]
    dec = (model(s, change_last(d_in)) - model(s, d_in)).abs().amax(-1)[0]
    cross = (model(change_last(s), d_in) - model(s, d_in)).abs().amax(-1)[0]
print(f"encoder: last source token changed -> source position 0 moved {enc[0]:.3f}")
print(f"decoder: last decoder input changed -> steps 0 to 14 moved {dec[:-1].max():.3f}, step 15 moved {dec[-1]:.3f}")
print(f"cross  : last source token changed -> output step 0 moved {cross[0]:.3f}")

Code walkthrough

  • change_last swaps the last token. Each line changes one input and measures the largest move at the named positions: Lesson 1's causal check, on each part.
CHANGE ONE LAST TOKEN, MEASURE THE MOVE (MEASURED HERE) encoder position 0 moved 0.075 bidirectional decoder steps 0 to 14 moved 0.000 causal cross-attention output step 0 moved 25.323 reads all source
Measured on the trained transformer: the encoder is bidirectional (0.075), the decoder is causal (exactly 0.000 before the change), and cross-attention reads the whole source (25.323).

Measured here: in the encoder, changing the last token moved position 0 by 0.075, so it is bidirectional. In the decoder, steps 0 to 14 moved by exactly 0.000 and step 15 by 3.627, so it is causal. Through cross-attention, the first output step moved by 25.323: it reads the whole source.

In production

THE FRAMINGS IN PUBLISHED MODELS (CITED) BERT, 2018 encoder-only left and right context GPT-2 decoder-only course reference model Whisper, 2022 encoder-decoder audio in, text out
All three framings are in use: BERT reads text both ways, GPT-2 writes left to right, and Whisper reads audio with an encoder and writes text with a decoder.

All three framings ship. BERT (Devlin and colleagues, 2018) is encoder-only, conditioning "on both left and right context in all layers". GPT-2, the course's reference model, is decoder-only. OpenAI's Whisper speech recogniser (Radford and colleagues, 2022) chose "an encoder-decoder Transformer (Vaswani et al., 2017) as this architecture has been well validated to scale reliably". Module 7 compares the families in detail.

What you added to tinyformer-lab

  • The reversal task in encoder-decoder form, at lengths 8, 16 and 32.
  • Four matched models: fixed, Bahdanau, Luong and a one-layer transformer, with accuracy and time.
  • An alignment reader and an intervention test that checks what a weight actually does.
  • A framing check for bidirectional, causal and cross-attention parts. The module lab's experiments/m01_rnn_vs_attention.py adds this lesson's multiplicative attention to the GRU as its second racer; the other models and checks stay in this lesson.

Check yourself

  1. Why did a wider fixed model still score only 0.344 at length 32?
  2. What is the one structural difference that let Luong's attention run as a single matrix product?
  3. Why can a model with argmax attention not learn where to look?
  4. The alignment peak sits one place right of the answer. Why is "the model looks in the wrong place" the wrong conclusion?
  5. Which framing would you choose to transcribe speech, and why?
  1. Its decoder still sees only the encoder's last state: the limit is the single vector.
  2. It attends after the recurrent step, so the recurrence never waits on attention.
  3. Argmax has no gradient, so the scoring weights never get a learning signal.
  4. Forcing the weight onto the peak gave 1.000: the model finds the token just read and takes its neighbour from that state.
  5. Encoder-decoder: audio is a separate input read in full, text is written token by token, as in Whisper.

Exercise

Luong's paper reports that the plain "dot" score "works well for the global attention". Write LuongDot, a Luong subclass scoring with S @ H.transpose(1, 2) and no weight matrix, then run the check. Predict first: better or worse than 0.780?

python
class LuongDot(Luong):
    """Luong's 'dot' score: s_t . h_j, no weight matrix."""
    def forward(self, src, dec_in, force=None):
        H, h = self.encoder(src)
        S, _ = self.gru(self.embed(dec_in), h)
        a = (S @ H.transpose(1, 2)).softmax(-1)          # (B, T, S) with no Wa at all
        self.align = a
        return self.head(torch.tanh(self.combine(torch.cat([S, a @ H], -1))))

Code walkthrough

  • Only the score line changes. Wa is still created by Luong.__init__ but never used.

Check:

python
set_seed(); dot_model = LuongDot()
acc, secs = train(dot_model, REV[32], 32)
print(f"dot score, length 32: {acc:.3f}   general: {RESULTS['luong', 32]:.3f}   fixed: {RESULTS['fixed', 32]:.3f}")
assert dot_model.Wa.weight.grad is None, "the score must not use Wa"
assert acc > RESULTS["fixed", 32] + 0.3, "attention should beat the bottleneck clearly"
print("check passed")

Code walkthrough

  • Trains your model with the lesson's recipe and checks that the score uses no weight matrix and clearly beats the fixed model. The author's run printed 0.911 (one seed).

Key terms

TermMeaning
Encoder-decoderOne network reads the input, another writes the output
BottleneckEverything about the input must pass through that one vector
Teacher forcingThe decoder reads the true previous token in training
Additive attentionScore = v · tanh(W s + U h): Bahdanau, 2014
Multiplicative attentionScore = s · W h or s · h: Luong, 2015; one matrix product
Soft lookupA softmax-weighted blend of values; differentiable
Alignment matrixAttention weights, output steps by source positions
InterventionSet the weights by hand and measure the output
Cross-attentionDecoder positions attending to the encoder's outputs

Sources

  • Dzmitry Bahdanau, Kyunghyun Cho, Yoshua Bengio, "Neural Machine Translation by Jointly Learning to Align and Translate", 2014 (ICLR 2015). https://arxiv.org/abs/1409.0473
  • Minh-Thang Luong, Hieu Pham, Christopher D. Manning, "Effective Approaches to Attention-based Neural Machine Translation", 2015. https://arxiv.org/abs/1508.04025
  • Sarthak Jain, Byron C. Wallace, "Attention is not Explanation", NAACL 2019. https://arxiv.org/abs/1902.10186
  • Ashish Vaswani et al., "Attention Is All You Need", 2017. https://arxiv.org/abs/1706.03762
  • Jacob Devlin et al., "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding", 2018. https://arxiv.org/abs/1810.04805
  • Alec Radford et al. (OpenAI), "Robust Speech Recognition via Large-Scale Weak Supervision" (Whisper), 2022. https://arxiv.org/abs/2212.04356