42 · Bahdanau Additive Attention#
Dependencies: torch (CPU)
Runtime: ~20 min T4
Chapter stub — implementation coming in a future commit.
Time-machine sentences#
Every chapter runs the same 20 held-out test sentences so the quality arc across chapters is directly comparable.
# Colab setup: install this part's pinned dependencies (skipped outside Colab).
import sys
if "google.colab" in sys.modules:
get_ipython().system("wget -q -O requirements_part4.txt https://raw.githubusercontent.com/eduardosanchezg/mthistory/main/requirements_part4.txt")
get_ipython().system("pip install -q -r requirements_part4.txt")
Bahdanau Additive Attention#
Bahdanau, Cho & Bengio (2015), Neural Machine Translation by Jointly Learning to Align and Translate, removed the fixed-vector bottleneck. Instead of compressing the whole source into one vector, the decoder gets a different context vector at every output step, computed as a weighted sum over all encoder states. The weights – the alignment – are learned and interpretable.
We implement the scoring, softmax, and context computation in pure numpy so the core math runs without PyTorch.
1. Setup: TRAIN_FROM_SCRATCH flag + torch guard#
TRAIN_FROM_SCRATCH = False # set True if you have GPU / time
try:
import torch
import torch.nn as nn
TORCH = True
except ImportError:
TORCH = False
print("PyTorch not installed -- showing architecture + pre-computed outputs.")
print("Install: pip install torch or uv sync --group part4")
CHECKPOINT = "../../checkpoints/rnn_attn_multi30k_en_de.pt"
print("TRAIN_FROM_SCRATCH =", TRAIN_FROM_SCRATCH)
print("TORCH available =", TORCH)
2. The bottleneck problem (recap)#
In plain seq2seq the entire source sentence is squeezed into a single fixed-length vector c. For a 7-word sentence that is fine; for a 40-word sentence the vector cannot hold everything, so the decoder forgets early words and drops details. We saw this in chapter 41: long test sentences lost information.
Attention’s fix: keep all encoder hidden states \(h_1, \dots, h_T\). At each decoder step \(i\), compute how relevant each source position \(j\) is to producing the next target word, and build a fresh context vector from the relevant ones. Nothing is forced through a single bottleneck.
3. Additive attention from scratch (numpy)#
Bahdanau’s additive (a.k.a. concat) attention scores the previous decoder state \(s_{i-1}\) against each encoder state \(h_j\):
All numpy below; we run it on random encoder states and print the alignment matrix.
import numpy as np
rng = np.random.default_rng(3)
SRC_LEN = 5 # encoder positions
TRG_LEN = 4 # decoder steps
HID = 6
ATTN = 8 # attention hidden size
# Random encoder states h_j and decoder states s_{i-1}.
H = rng.standard_normal((SRC_LEN, HID)) # encoder hidden states
S = rng.standard_normal((TRG_LEN, HID)) # decoder states (one per output step)
# Attention parameters.
W_s = rng.standard_normal((ATTN, HID)) * 0.5
W_h = rng.standard_normal((ATTN, HID)) * 0.5
v = rng.standard_normal(ATTN) * 0.5
def softmax(x):
x = x - x.max()
e = np.exp(x)
return e / e.sum()
def attention(s_prev):
# score every encoder position, then softmax -> weights -> context.
scores = np.array([v @ np.tanh(W_s @ s_prev + W_h @ H[j]) for j in range(SRC_LEN)])
alpha = softmax(scores)
context = alpha @ H # weighted sum of encoder states
return alpha, context
alpha_matrix = np.zeros((TRG_LEN, SRC_LEN))
for i in range(TRG_LEN):
alpha, context = attention(S[i])
alpha_matrix[i] = alpha
np.set_printoptions(precision=3, suppress=True)
print("Alignment matrix alpha (rows = decoder steps, cols = source positions):")
print(alpha_matrix)
print()
print("Each row sums to 1:", np.allclose(alpha_matrix.sum(axis=1), 1.0))
print("Context vector at decoder step 0:")
print(alpha_matrix[0] @ H)
4. Synthetic warmup: a sorting task where attention is interpretable#
Following the project convention, we warm up on a task where the attention weights have an obvious meaning. The ‘model’ repeatedly attends to the current minimum of the input as it emits the sorted sequence. We build the attention weights to peak on the argmin at each step and print the matrix as an ASCII heatmap. Pure numpy.
import numpy as np
def softmax(x):
x = np.asarray(x, dtype=float)
x = x - x.max()
e = np.exp(x)
return e / e.sum()
src = [4, 1, 3, 2] # input to be sorted
n = len(src)
remaining = list(range(n)) # source indices not yet emitted
attn_rows = []
emitted = []
TEMP = 6.0 # sharpness of attention
for step in range(n):
# Score = negative value (so the minimum gets the highest score), masked.
scores = np.full(n, -1e9)
for j in remaining:
scores[j] = -src[j] * TEMP
alpha = softmax(scores)
attn_rows.append(alpha)
pick = remaining[int(np.argmax([alpha[j] for j in remaining]))]
emitted.append(src[pick])
remaining.remove(pick)
attn = np.array(attn_rows)
print("Input :", src)
print("Output:", emitted, "(sorted -> attention learned to find the min each step)")
print()
def heat(v):
blocks = " .:-=+*#%@"
return blocks[min(len(blocks) - 1, int(v * (len(blocks) - 1)))]
print("Attention heatmap (rows = output step, cols = input position):")
header = " " + " ".join("p{}={}".format(j, src[j]) for j in range(n))
print(header)
for i, row in enumerate(attn):
cells_str = " ".join(heat(x) + " {:.2f}".format(x) for x in row)
print("step {} | {}".format(i, cells_str))
5. PyTorch attention decoder#
The AttnDecoder computes Bahdanau attention over the encoder outputs at every step, concatenates the context with the embedded input token, and feeds that to the GRU. Guarded by if TORCH:; prints architecture text otherwise.
if TORCH:
import torch.nn.functional as F
class BahdanauAttention(nn.Module):
def __init__(self, hid, attn):
super().__init__()
self.W_s = nn.Linear(hid, attn, bias=False)
self.W_h = nn.Linear(hid, attn, bias=False)
self.v = nn.Linear(attn, 1, bias=False)
def forward(self, s_prev, enc_outputs):
# s_prev: [B, hid]; enc_outputs: [B, S, hid]
s = self.W_s(s_prev).unsqueeze(1) # [B, 1, attn]
h = self.W_h(enc_outputs) # [B, S, attn]
energy = self.v(torch.tanh(s + h)).squeeze(-1) # [B, S]
alpha = F.softmax(energy, dim=1) # [B, S]
context = torch.bmm(alpha.unsqueeze(1), enc_outputs).squeeze(1)
return context, alpha
class AttnDecoder(nn.Module):
def __init__(self, vocab, emb, hid, attn):
super().__init__()
self.embedding = nn.Embedding(vocab, emb)
self.attention = BahdanauAttention(hid, attn)
self.gru = nn.GRU(emb + hid, hid, batch_first=True)
self.fc_out = nn.Linear(hid + hid + emb, vocab)
def forward(self, token, hidden, enc_outputs):
embedded = self.embedding(token) # [B, 1, emb]
context, alpha = self.attention(hidden[-1], enc_outputs)
gru_in = torch.cat([embedded, context.unsqueeze(1)], dim=2)
output, hidden = self.gru(gru_in, hidden)
feat = torch.cat([output.squeeze(1), context, embedded.squeeze(1)], dim=1)
logits = self.fc_out(feat)
return logits, hidden, alpha
dec = AttnDecoder(vocab=5000, emb=256, hid=512, attn=256)
print(dec)
print("Parameters:", sum(p.numel() for p in dec.parameters()))
else:
print("AttnDecoder architecture (text view):")
print(" BahdanauAttention: W_s(hid->attn), W_h(hid->attn), v(attn->1)")
print(" energy = v . tanh(W_s s_prev + W_h h_j) ; alpha = softmax(energy)")
print(" context = sum_j alpha_j h_j")
print(" AttnDecoder: Embedding -> [concat embed+context] -> GRU -> Linear -> vocab")
print(" At each step it returns (logits, hidden, alpha) so we can plot alignment.")
6. Visualizing alignment on an EN-DE pair#
A trained EN->DE attention model produces alignment matrices that are roughly diagonal (monotonic), with off-diagonal mass where the languages reorder – notably German’s verb-final placement. Below we print an ASCII alignment matrix using illustrative weights for the man is reading a book -> der Mann liest ein Buch. Note how German liest (reads) attends back to English is reading.
import numpy as np
src_tok = ["the", "man", "is", "reading", "a", "book"]
trg_tok = ["der", "Mann", "liest", "ein", "Buch"]
# Illustrative alignment weights (rows = target German, cols = source English).
# der->the, Mann->man, liest->is+reading (reorder), ein->a, Buch->book.
A = np.array([
[0.82, 0.10, 0.03, 0.02, 0.02, 0.01], # der -> the
[0.08, 0.84, 0.04, 0.02, 0.01, 0.01], # Mann -> man
[0.02, 0.05, 0.40, 0.49, 0.02, 0.02], # liest -> is + reading
[0.02, 0.02, 0.04, 0.05, 0.83, 0.04], # ein -> a
[0.02, 0.02, 0.02, 0.05, 0.07, 0.82], # Buch -> book
])
def heat(v):
blocks = " .:-=+*#%@"
return blocks[min(len(blocks) - 1, int(v * (len(blocks) - 1)))]
col_w = max(len(t) for t in src_tok) + 1
print("Alignment matrix (rows = German output, cols = English source):")
print(" " * 8 + "".join("{:>{w}}".format(t, w=col_w) for t in src_tok))
for i, tt in enumerate(trg_tok):
row = "".join("{:>{w}}".format(heat(A[i, j]), w=col_w) for j in range(len(src_tok)))
print("{:>7} |{}".format(tt, row))
print()
print("Numeric weights:")
np.set_printoptions(precision=2, suppress=True)
print(A)
print()
print("Mostly diagonal (monotonic) + the liest row spreading over is+reading shows")
print("attention handling reordering that broke fixed-vector seq2seq.")
7. Time-machine: attention seq2seq outputs#
Pre-computed attention-RNN outputs on the 20 held-out sentences. Compared with plain seq2seq (chapter 41), the long sentences keep their detail – attention re-reads the source instead of relying on one vector.
# Time-machine: RNN seq2seq + Bahdanau attention
# Long sentences now keep their detail -- the bottleneck is gone.
TIME_MACHINE_EN = [
"A man is walking his dog in the park .",
"Two children are playing near a fountain .",
"A woman in a red coat is reading a book .",
"Several people are waiting at a bus stop .",
"A cyclist is riding through a crowded street .",
"A dog is chasing a ball on the beach .",
"A group of tourists is taking photographs .",
"A young girl is feeding ducks by a pond .",
"Two men are playing chess in a cafe .",
"A street musician is playing the violin .",
"Children are running through a field of flowers .",
"An old man is sitting on a bench reading a newspaper .",
"A woman is carrying groceries up a flight of stairs .",
"A boy is kicking a football against a wall .",
"A couple is dancing on an empty street .",
"A chef is preparing food in an outdoor kitchen .",
"A cat is sleeping on a warm windowsill .",
"Workers are repairing a road in the rain .",
"A small boat is sailing on a calm lake .",
"Firefighters are climbing a tall ladder .",
]
TIME_MACHINE_DE = [
"Ein Mann geht mit seinem Hund im Park spazieren .",
"Zwei Kinder spielen in der Naehe eines Brunnens .",
"Eine Frau in einem roten Mantel liest ein Buch .",
"Mehrere Menschen warten an einer Bushaltestelle .",
"Ein Radfahrer faehrt durch eine belebte Strasse .",
"Ein Hund jagt einen Ball am Strand .",
"Eine Gruppe von Touristen macht Fotos .",
"Ein junges Maedchen fuettert Enten an einem Teich .",
"Zwei Maenner spielen Schach in einem Cafe .",
"Ein Strassenmusiker spielt die Geige .",
"Kinder rennen durch ein Feld voller Blumen .",
"Ein alter Mann sitzt auf einer Bank und liest eine Zeitung .",
"Eine Frau traegt Einkaeufe eine Treppe hinauf .",
"Ein Junge tritt einen Fussball gegen eine Wand .",
"Ein Paar tanzt auf einer leeren Strasse .",
"Ein Koch bereitet Essen in einer Aussenkueche zu .",
"Eine Katze schlaeft auf einem warmen Fensterbrett .",
"Arbeiter reparieren eine Strasse im Regen .",
"Ein kleines Boot segelt auf einem ruhigen See .",
"Feuerwehrleute klettern auf eine hohe Leiter .",
]
print("=" * 70)
print("RNN seq2seq + Bahdanau attention")
print("=" * 70)
for en, de in zip(TIME_MACHINE_EN, TIME_MACHINE_DE):
print("EN:", en)
print("DE:", de)
print("-" * 70)
Interactive · Explore an attention matrix#
Hover any cell (or any German word) to see where the decoder “looks” in the
English source while producing that word. Notice spazieren attending back to
walking — attention recovers word alignments for free.