Lesson 103: Seq2Seq Encoder-Decoder with Bahdanau Attention

1. Seq2Seq Encoder-Decoder + Bahdanau Attention

Title

USAAIO -- Lesson 103 -- Phase 3

The architecture behind neural machine translation, summarization, and code generation: a GRU encoder compresses the source into hidden states, a decoder attends over them to generate each target token. Today: attention math, teacher forcing, exposure bias, scheduled sampling, and copy mechanisms.

2. By the end of this lesson you can

Objectives

1. Trace the seq2seq forward pass (encoder hidden states -> context -> decoder output) with concrete tensor shapes
2. Compute Bahdanau attention energy scores, softmax weights, and the context vector by hand
3. Explain teacher forcing: why it speeds convergence and what the training objective actually is
4. Define exposure bias and explain why it causes train-inference mismatch in seq2seq models
5. Describe scheduled sampling as a remedy and state the linear decay schedule
6. Explain the copy mechanism as an attention-based pointer distribution over source tokens

3. What survived from RLHF & DPO — Instruction Fine-Tuning?

Warm-up

4. Seq2Seq architecture -- encoder to decoder

Section

Part 1 of 4

5. The encoder-decoder idea

Concept

A seq2seq model maps a variable-length input sequence to a variable-length output sequence. Unlike classification (fixed output), both lengths are unknown at design time: a 7-word English sentence can become a 5-word French sentence or a 20-word German one.

• Encoder (Lesson 95 GRU): reads all source tokens left-to-right and produces a hidden state per token. Final hidden state h_T = compressed source summary.
• Decoder: conditioned on the encoder's hidden states, generates target tokens one at a time using its own GRU -- stopping at a special <EOS> token.
• With attention (Bahdanau): the decoder does not rely on a single fixed h_T; it attends over ALL encoder hidden states {h_1, ..., h_T} at each generation step.

6. Fill in: input for The encoder-decoder idea

Comparison

7. Bahdanau (additive) attention

Concept

At decoder step t, Bahdanau attention assigns a soft weight alpha_i to each encoder hidden state h_i. The weighted sum is the context vector fed into the decoder GRU alongside the input embedding.

\[ e_{t,i} = v^\top \tanh(W_h\, h_i + W_s\, s_{t-1}) \]

\[ \alpha_{t,i} = \frac{\exp(e_{t,i})}{\sum_{j=1}^{T}\exp(e_{t,j})}, \quad c_t = \sum_{i=1}^{T} \alpha_{t,i}\, h_i \]

Three learned parameters: W_h in R^{H x H}, W_s in R^{H x H}, v in R^H. No query-key dot product -- the additive tanh(W_h h + W_s s) is the distinguishing feature of Bahdanau vs Luong (dot-product) attention.

8. Guess the shape of the answer: Bahdanau attention -- numeric trace

Estimation

9. Bahdanau attention -- numeric trace

Worked example

Toy example: 3-token source, H=4. All weights drawn from torch.randn with torch.manual_seed(0). Predict which source position gets the highest attention weight before reading on.

import torch
torch.manual_seed(0)
SRC_LEN, H = 3, 4
h  = torch.randn(SRC_LEN, H)  # encoder states
s  = torch.randn(H)           # decoder hidden state
Wh = torch.randn(H, H)
Ws = torch.randn(H, H)
v  = torch.randn(H)
Wh_h   = h @ Wh.T                     # (3, H)
Ws_s   = (Ws @ s).unsqueeze(0)        # (1, H) broadcasts
energy = torch.tanh(Wh_h + Ws_s) @ v  # (3,)
alpha  = torch.softmax(energy, dim=0)
ctx    = (alpha.unsqueeze(1) * h).sum(0)
print([round(x,4) for x in energy.tolist()])
print([round(x,4) for x in alpha.tolist()])
print([round(x,4) for x in ctx.tolist()])

10. What each one costs: Bahdanau attention -- numeric trace

Trade off

11. Something is wrong here: Bahdanau vs Luong attention

Anomaly

12. Trap: Bahdanau vs Luong attention

Trap

13. Break it on purpose: Bahdanau vs Luong attention

Break the constraint

14. Teacher forcing & exposure bias

Section

Part 2 of 4

15. Teacher forcing -- the training trick

Concept

At training time, the decoder knows the ground-truth target sequence. Teacher forcing feeds the true previous token as the decoder's input at each step, regardless of what the model actually predicted.

• Without TF: if step 1 is wrong, step 2 receives a bad input -- error compounds, gradients are noisy, convergence is slow.
• With TF: each step always receives a clean, correct conditioning signal -- the gradient at each step is independent of earlier mistakes.
• Effect: much faster convergence, especially in early training when the model's predictions are mostly random.

16. Exposure bias -- the train-inference gap

Concept

At inference time, there is no ground-truth available. The decoder must use its own previous prediction as the next input -- the free-running regime.

This creates exposure bias: the decoder was never exposed to its own errors during training, so it has never learned to recover from them. One early mistake cascades through the rest of the sequence.

17. Break it if you can: Exposure bias -- the train-inference gap

Counterexample

18. What has to be given first: Teacher forcing vs free-run -- divergence…

Missing information

19. Teacher forcing vs free-run -- divergence demo

Worked example

Tiny seq2seq: VOCAB=10, E=6, H=8, src=[2,5,1,3], tgt=[0,4,6,3] (BOS then real tokens). With torch.manual_seed(99): compare TF inputs to free-run inputs over 3 steps.

import torch, torch.nn as nn
torch.manual_seed(99)
VOCAB, E, H = 10, 6, 8
enc = Encoder(VOCAB, E, H)   # defined in milestone
dec = Decoder(VOCAB, E, H)
src = torch.tensor([[2, 5, 1, 3]])
tgt_seq = torch.tensor([[0, 4, 6, 3]])
enc_out, enc_hid = enc(src)
for mode in ['teacher_forcing', 'free_run']:
    hid = enc_hid
    inp_list, pred_list = [], []
    tok = tgt_seq[:, 0:1]  # BOS
    for t in range(3):
        logit, hid, _ = dec(tok, hid, enc_out)
        pred = logit.argmax(-1).item()
        inp_list.append(tok.item())
        pred_list.append(pred)
        tok = (tgt_seq[:, t+1:t+2] if mode == 'teacher_forcing'
               else torch.tensor([[pred]]))
    print(f'{mode}: inputs={inp_list}, preds={pred_list}')

20. Watch it run: Teacher forcing vs free-run -- divergence demo

Pattern

21. Scheduled sampling & copy mechanism

Section

Part 3 of 4

22. Scheduled sampling -- bridging the gap

Concept

Scheduled sampling (Bengio 2015) mixes teacher forcing and free-running during training. At each step, with probability epsilon the decoder receives the ground-truth token; with probability 1 - epsilon it receives its own prediction.

\[ \text{input}_t = \begin{cases} y_{t-1} & \text{with prob } \epsilon_k \\ \hat{y}_{t-1} & \text{with prob } 1 - \epsilon_k \end{cases} \]

epsilon_k is annealed from 1.0 (pure teacher forcing) toward 0.0 (pure free-run) as training epoch k increases. The most common schedule is linear: epsilon_k = max(0, 1 - k/K).

23. By analogy: Scheduled sampling -- bridging the gap

Analogy

24. Guess the shape of the answer: Scheduled sampling -- linear decay trace

Estimation

25. Scheduled sampling -- linear decay trace

Worked example

Linear schedule: epsilon_k = max(0, 1 - k/K) with K=100. At each epoch, a Bernoulli draw decides which input token the decoder sees. Trace epsilon across epochs.

import numpy as np
K = 100
steps = [0, 10, 25, 50, 75, 100]
print(f"{'epoch':>8} {'epsilon':>10} {'mode':>12}")
for k in steps:
    eps = max(0.0, 1.0 - k / K)
    mode = ('teacher' if eps > 0.75
            else 'mixed' if eps > 0.25
            else 'free-run')
    print(f"{k:>8} {eps:>10.2f} {mode:>12}")

26. Watch it run: Scheduled sampling -- linear decay trace

Pattern

27. Copy mechanism -- pointing to the source

Concept

In summarization and question answering, the decoder often needs to copy a token verbatim from the input (named entities, numbers, rare words not in the vocab). The copy mechanism (Gu 2016, Vinyals pointer networks 2015) reuses the attention distribution as a pointer to source positions.

• Generate probability p_gen: a gating scalar in [0,1] learned from context. Controls how much output comes from the vocabulary vs. copying.
• Final distribution: P(w) = p_gen * P_vocab(w) + (1 - p_gen) * sum_{i: src_i=w} alpha_i
• Effect: rare proper nouns that never appeared in training vocab can still be generated by copying from the source.

28. Where does each piece belong: Lesson 103: Seq2Seq Encoder-Decoder with…

Sorting

29. Guess the shape of the answer: Copy mechanism -- pointer distribution

Estimation

30. Copy mechanism -- pointer distribution

Worked example

Source: ['the', 'customer', 'ordered', 'tea']. Simulate the copy distribution: torch.manual_seed(1), 4-dim softmax over source positions. Which token does the model most likely copy?

import torch
torch.manual_seed(1)
src_tokens = ['the', 'customer', 'ordered', 'tea']
attn_logits = torch.randn(4)          # from attention module
copy_probs  = torch.softmax(attn_logits, dim=0)
for tok, p in zip(src_tokens, copy_probs.tolist()):
    print(f"  '{tok}': {round(p, 4)}")
best = src_tokens[copy_probs.argmax().item()]
print(f"  -> most-likely copied: '{best}'")

31. Fill in: rank for Copy mechanism -- pointer distribution

Comparison

32. Something is wrong here: teacher forcing causes overfitting

Anomaly

33. Trap: teacher forcing causes overfitting

Trap

34. Which of these survive contact with Lesson 103: Seq2Seq Encoder-Decoder with…?

Two truths and a lie

35. Full Seq2Seq in PyTorch from scratch

Section

Part 4 of 4

36. Architecture summary: shapes at every step

Concept

Concrete shapes for B=1, src_len=4, H=8, E=6, VOCAB=10 (our toy from the worked example).

Total parameters (VOCAB=10, E=6, H=8): encoder=444, decoder=942, total=1,386 (verified with sum(p.numel()...)). Scale to VOCAB=30000, H=512 for real NMT: ~30M params.

37. By analogy: Architecture summary: shapes at every step

Analogy

38. Without one step: The seq2seq + attention recipe

Constraint

39. The seq2seq + attention recipe

Pattern

1. Encoder: nn.Embedding + nn.GRU(batch_first=True). Keep ALL hidden states enc_out (B,T,H), not just the final one.
2. Bahdanau attention: at each decoder step, energy = v^T tanh(W_h enc_out + W_s s), alpha = softmax(energy), ctx = sum(alpha * enc_out).
3. Decoder GRU input: cat([embed(prev_token), ctx], dim=-1) -- (E+H)-dim input at each step.
4. Output head: Linear(H*2, V) applied to cat([dec_out, ctx]) -- concatenating context again improves performance.
5. Teacher forcing (training): feed y_{t-1} (GT) as decoder input; cross-entropy loss on predicted vocab distribution.
6. Scheduled sampling: use Bernoulli(epsilon_k) to decide GT vs model prediction; anneal epsilon_k linearly from 1.0 to 0.0.
7. Copy mechanism (if needed): gate p_gen * P_vocab + (1-p_gen) * sum_i alpha_i [src_i=w] over the output distribution.

40. Where does it stop working: The seq2seq + attention recipe

Edge cases

41. Rule out three: Check yourself -- Bahdanau energy

Elimination

42. Check yourself -- Bahdanau energy

Check

Compute on paper before clicking.

43. Answer it before you see the options: Check yourself -- teacher forcing…

Prediction

44. Check yourself -- teacher forcing objective

Check

Think about what loss is actually being optimized.

45. Rule out three: Check yourself -- scheduled sampling

Elimination

46. Check yourself -- scheduled sampling

Check

Apply the linear schedule formula.

47. Your turn: Seq2Seq from scratch

Section

Project

48. Project: seq2seq with Bahdanau attention

Concept

Build a complete Encoder + BahdanauAttn + Decoder from scratch in PyTorch. Three milestones: attention math -> full forward pass -> teacher-forcing training loop.

Build rules: print every tensor shape after each operation; manually verify alpha.sum(-1) == 1; use torch.manual_seed(99) throughout to match the deck's verified numbers.

49. Break it if you can: Project: seq2seq with Bahdanau attention

Counterexample

50. Milestone 1 -- BahdanauAttn module

Worked example

Your turn: implement BahdanauAttn(H) with three nn.Linear layers (Wh, Ws, v). Given enc_out (B, T, H) and decoder hidden dec_h (B, H), return (alpha, context). Predict: does alpha.sum(dim=1) equal [1.0]?

Hint: self.Ws(dec_h).unsqueeze(1) broadcasts over the time dimension -- shape (B, 1, H) adds to self.Wh(enc_out) of shape (B, T, H) correctly.

import torch, torch.nn as nn

class BahdanauAttn(nn.Module):
    def __init__(self, H):
        super().__init__()
        self.Wh = nn.Linear(H, H, bias=False)
        self.Ws = nn.Linear(H, H, bias=False)
        self.v  = nn.Linear(H, 1, bias=False)
    def forward(self, enc_out, dec_h):
        # enc_out: (B,T,H)   dec_h: (B,H)
        energy = self.v(
            torch.tanh(
                self.Wh(enc_out)
                + self.Ws(dec_h).unsqueeze(1)
            )
        )                         # (B, T, 1)
        alpha = torch.softmax(energy, dim=1)  # (B, T, 1)
        ctx   = (alpha * enc_out).sum(1)       # (B, H)
        return alpha.squeeze(-1), ctx

torch.manual_seed(99)
attn = BahdanauAttn(8)
enc_out = torch.randn(1, 4, 8)
dec_h   = torch.randn(1, 8)
alpha, ctx = attn(enc_out, dec_h)
print(alpha.shape, alpha.sum(dim=1))   # (1,4)  tensor([[1.0]])
print(ctx.shape)                       # (1, 8)

51. Which is which, by shape

Discrimination

52. Milestone 2 -- Encoder + Decoder, single step

Worked example

Your turn: implement Encoder and Decoder using the BahdanauAttn you just built. Run one decoder step from BOS and print the predicted token. Predict: with random weights, will the prediction be stable across multiple runs without manual_seed?

Hint: Decoder GRU input is cat([emb(prev_tok), ctx], dim=-1) -- shape (B, 1, E+H). Output head is Linear(H*2, V) applied to cat([dec_gru_out, ctx]).

class Encoder(nn.Module):
    def __init__(self, V, E, H):
        super().__init__()
        self.emb = nn.Embedding(V, E)
        self.gru = nn.GRU(E, H, batch_first=True)
    def forward(self, x):
        return self.gru(self.emb(x))  # (B,T,H), (1,B,H)

class Decoder(nn.Module):
    def __init__(self, V, E, H):
        super().__init__()
        self.emb  = nn.Embedding(V, E)
        self.attn = BahdanauAttn(H)
        self.gru  = nn.GRU(E + H, H, batch_first=True)
        self.fc   = nn.Linear(H * 2, V)
    def forward(self, tok, hid, enc_out):
        e = self.emb(tok)                             # (B,1,E)
        h = hid.squeeze(0)                            # (B,H)
        alpha, ctx = self.attn(enc_out, h)
        inp = torch.cat([e, ctx.unsqueeze(1)], -1)   # (B,1,E+H)
        out, new_hid = self.gru(inp, hid)
        logit = self.fc(
            torch.cat([out.squeeze(1), ctx], -1))    # (B,V)
        return logit, new_hid, alpha

torch.manual_seed(99)
VOCAB, E, H = 10, 6, 8
enc = Encoder(VOCAB, E, H)
dec = Decoder(VOCAB, E, H)
src = torch.tensor([[2, 5, 1, 3]])
enc_out, enc_hid = enc(src)
logit, _, alpha = dec(torch.tensor([[0]]), enc_hid, enc_out)
print('enc_out:', list(enc_out.shape))
print('pred:', logit.argmax(-1).item(), '  alpha:', [round(x,3) for x in alpha[0].tolist()])

53. Milestone 3 -- teacher-forcing training loop

Worked example

Your turn: write the teacher-forcing training loop. At each decoder step, feed the GROUND-TRUTH previous token. Accumulate cross-entropy loss across all steps and backprop once.

Hint: slice tgt_seq[:, t:t+1] as the input at step t (the (t)th ground-truth token, offset by 1 from the target). The loss target is tgt_seq[:, t+1] (the next token to predict).

import torch, torch.nn as nn
torch.manual_seed(99)
VOCAB, E, H, TGT_LEN = 10, 6, 8, 3
enc = Encoder(VOCAB, E, H)
dec = Decoder(VOCAB, E, H)
params = list(enc.parameters()) + list(dec.parameters())
opt = torch.optim.Adam(params, lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
src = torch.tensor([[2, 5, 1, 3]])
tgt = torch.tensor([[0, 4, 6, 3]])  # BOS, tok4, tok6, EOS
for epoch in range(3):   # show 3 epochs
    opt.zero_grad()
    enc_out, hid = enc(src)
    loss = torch.tensor(0.0)
    for t in range(TGT_LEN):
        logit, hid, _ = dec(tgt[:, t:t+1], hid, enc_out)
        loss = loss + loss_fn(logit, tgt[:, t + 1])
    loss.backward(); opt.step()
    print(f'epoch {epoch}: loss = {loss.item():.4f}')

54. What each one costs: Milestone 3 -- teacher-forcing training loop

Trade off

55. The full program

Concept

import torch, torch.nn as nn
torch.manual_seed(99)

class BahdanauAttn(nn.Module):
    def __init__(self, H):
        super().__init__()
        self.Wh = nn.Linear(H, H, bias=False)
        self.Ws = nn.Linear(H, H, bias=False)
        self.v  = nn.Linear(H, 1, bias=False)
    def forward(self, enc_out, dec_h):
        e = self.v(torch.tanh(self.Wh(enc_out) + self.Ws(dec_h).unsqueeze(1)))
        alpha = torch.softmax(e, dim=1)
        return alpha.squeeze(-1), (alpha * enc_out).sum(1)

class Encoder(nn.Module):
    def __init__(self, V, E, H):
        super().__init__()
        self.emb = nn.Embedding(V, E)
        self.gru = nn.GRU(E, H, batch_first=True)
    def forward(self, x): return self.gru(self.emb(x))

class Decoder(nn.Module):
    def __init__(self, V, E, H):
        super().__init__()
        self.emb  = nn.Embedding(V, E)
        self.attn = BahdanauAttn(H)
        self.gru  = nn.GRU(E + H, H, batch_first=True)
        self.fc   = nn.Linear(H * 2, V)
    def forward(self, tok, hid, enc_out):
        e = self.emb(tok); h = hid.squeeze(0)
        alpha, ctx = self.attn(enc_out, h)
        out, new_hid = self.gru(torch.cat([e, ctx.unsqueeze(1)], -1), hid)
        return self.fc(torch.cat([out.squeeze(1), ctx], -1)), new_hid, alpha

VOCAB, E, H = 10, 6, 8
enc = Encoder(VOCAB, E, H); dec = Decoder(VOCAB, E, H)
opt = torch.optim.Adam(list(enc.parameters())+list(dec.parameters()), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
src = torch.tensor([[2,5,1,3]]); tgt = torch.tensor([[0,4,6,3]])
for ep in range(200):
    opt.zero_grad(); enc_out, hid = enc(src); loss = torch.tensor(0.)
    for t in range(3):
        logit, hid, _ = dec(tgt[:, t:t+1], hid, enc_out)
        loss = loss + loss_fn(logit, tgt[:, t+1])
    loss.backward(); opt.step()
print(f'final loss: {loss.item():.4f}')

56. Fill in: rationale for The full program

Comparison

57. Show it off

Concept

Out loud, slides closed: (1) explain the Bahdanau attention computation from energy scores to context vector, naming every matrix and its shape; (2) define exposure bias and state exactly which distribution mismatch causes it; (3) describe scheduled sampling and give the linear decay formula.

Stretch (homework from the lesson plan): implement scheduled sampling -- add a Bernoulli draw at each decoder step and anneal epsilon linearly over 200 epochs; compare final BLEU on a held-out sequence vs pure teacher forcing. Next: Lesson 104 -- advanced NLP and pretraining strategies.

58. Connect it up: Lesson 103: Seq2Seq Encoder-Decoder with Bahdanau Attention

Connect it up

59. What you can do now

Recap

• Trace the seq2seq forward pass: encoder GRU -> enc_out (B,T,H) -> Bahdanau attention -> context -> decoder GRU -> logits
• Compute Bahdanau energy e_i = v^T tanh(W_h h_i + W_s s), softmax alpha, and weighted-sum context vector
• Explain teacher forcing: feeds GT tokens at every training step; faster convergence but creates exposure bias
• Define exposure bias: model never sees its own errors during training; at inference errors cascade
• State scheduled sampling: epsilon_k (linear from 1 to 0) controls GT vs model-prediction probability per step
• Describe the copy mechanism: p_gen * P_vocab + (1-p_gen) * sum_i alpha_i [src_i=w] lets the decoder point to source tokens