DL S035 · Why Attention — Bahdanau's Scribble That Ate NLP

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 35 / 80 · Module M06 · Enriched pass 2026-07-27 · Module M06 finale

The story

0 · Montréal, late September 2014 — the scribble

Yoshua Bengio's lab at the Université de Montréal, autumn 2014. Dzmitry Bahdanau, a first-year PhD student, is working with Kyunghyun Cho on machine translation. They've read Sutskever's seq2seq paper the moment it hit arXiv three weeks earlier. They've been trying to reproduce it on WMT English-French, and — like Sutskever himself — they're noticing the length problem. Translations start well and then, somewhere around word 20, deteriorate into fluent-sounding but semantically wrong French.

Bahdanau tells the story (in a 2019 interview) that he was staring at a whiteboard trying to figure out where in the pipeline the information got lost. His scribble was something like: "the encoder produces one hidden state per input word. Why does the decoder throw away all of them except the last one? What if at each decoder step, the decoder could look back at the whole sequence of encoder states, weighted by how relevant each one is right now?"

That was the germ. He, Cho, and Bengio worked out the details over the next month or two, and by December 2014 the paper Neural Machine Translation by Jointly Learning to Align and Translate was on arXiv. It was accepted to ICLR 2015. It buys 3+ BLEU points on WMT En-Fr — a huge margin at the time — and, more importantly, it flattens the length-degradation curve. Translations of 60-word sentences are almost as accurate as translations of 10-word sentences. The bottleneck is gone.

The mechanism they introduced was a small MLP that, at each decoder step, scored every encoder hidden state for relevance, softmaxed those scores into a probability distribution ("attention weights"), and used them to compute a weighted sum of encoder states as the context vector for that step. A different context vector for each decoded word. That was it. Six additional weight matrices and a softmax.

Three years later, in June 2017, Vaswani et al. at Google Brain would publish Attention Is All You Need, stripping away the LSTM entirely and building a network out of nothing but stacked self-attention layers. That paper defines the architecture of essentially every large language model deployed in 2025 — GPT, Claude, Llama, DeepSeek, Gemini, Qwen, Mixtral. The transformer is Bahdanau's attention, generalised: instead of the decoder attending to the encoder, every token in the sequence attends to every other token.

Today's session is the pivot on which the entire modern LLM edifice balances. The rest of Module M07 (self-attention, multi-head, positional encodings, transformer blocks, nanoGPT, RoPE, KV cache, FlashAttention) is what happens when you take today's idea and push it to its logical extreme.

1 · What you'll build today

One sentence: You will add Bahdanau (additive) attention to the seq2seq model from S034, verify that the length-vs-accuracy curve flattens dramatically, then implement Luong (multiplicative) attention as a lightweight alternative — arriving at the exact abstraction ("query, key, value → weighted sum") that transformers will generalise.

By the end you will:

• Explain the seq2seq bottleneck in one crisp sentence and give the geometric fix (weighted sum over encoder states).
• Write additive attention $\alpha_{t,i} \propto \exp(v^\top \tanh(W_h h_i^\text{enc} + W_s s_{t-1}^\text{dec}))$ from memory.
• Write dot-product / scaled dot-product attention and know when to prefer it over additive.
• Visualise attention weights as a heatmap and see the model learn word alignment for free.
• Recognise the query/key/value abstraction and see it prefigure the transformer.
• Trace the full arc: seq2seq (2014) → Bahdanau attention (2015) → transformer (2017) → GPT-3 (2020) → ChatGPT (2022) → Llama-4 / DeepSeek-V3 / Claude 4 (2024–2025).

2 · The pain, restated in one paragraph

Seq2seq compresses the entire source sentence into one fixed-size vector (the encoder's final hidden state) and hands it to the decoder. Every decoded word then depends on the source only through this one vector. If the source is 5 words, the vector holds enough. If the source is 50 words, the vector cannot possibly encode all of them with enough fidelity to translate word 47 correctly. You saw this in S034 §9 — accuracy fell off a cliff around length 10 on the toy task; on real WMT it fell off around length 30–40.

Attention's fix: don't compress. Keep all $T_\text{src}$ encoder hidden states around. At each decoder step, compute a fresh context vector as a weighted sum of the encoder states, with the weights chosen dynamically per step based on what the decoder is trying to produce next.

3 · Additive (Bahdanau) attention — the equations

Setup. Encoder produces states $h_1, h_2, \ldots, h_{T_\text{src}} \in \mathbb{R}^{H}$ (each is the output of the encoder LSTM at that timestep — you use .forward returning all outputs now, not just the final state). Decoder is producing token $t$, with previous hidden state $s_{t-1} \in \mathbb{R}^{H}$.

Step 1 — score each encoder state for relevance.

for $i = 1, \ldots, T_\text{src}$. Here $W_h, W_s \in \mathbb{R}^{H' \times H}$ are learned matrices and $v \in \mathbb{R}^{H'}$ is a learned vector. This is a one-hidden-layer MLP that takes the concatenation of (encoder state $h_i$, decoder previous state $s_{t-1}$) and outputs a single scalar score.

Step 2 — normalise into a probability distribution.

So $\alpha_{t, i}$ is the "attention weight" that the decoder at step $t$ puts on encoder position $i$. All the $\alpha_{t, \cdot}$ sum to 1.

Step 3 — compute the context vector as a weighted sum.

Step 4 — use $c_t$ in the decoder step. Concatenate $c_t$ with the previous decoder input embedding and feed to the decoder LSTM:

Then the output logits are $\text{Linear}(\text{output}_t)$ as before, or the more expressive $\text{Linear}([\text{output}_t;\, c_t])$.

That is Bahdanau attention. Four equations. Two new matrices ($W_h, W_s$) and one new vector ($v$). Everything else is unchanged from seq2seq.

3.1 A worked example — 3 encoder states, 1 decoder step

Set $H = H' = 2$ so we can inspect every number. Encoder states:

Decoder previous state $s_{t-1} = [0.5, 0.5]$. Weights:

Scores:

• $e_1 = v^\top \tanh(W_h h_1 + W_s s_{t-1}) = [1,1]^\top \tanh([1.5, 0.5]) = \tanh(1.5) + \tanh(0.5) = 0.905 + 0.462 = 1.367$
• $e_2 = [1,1]^\top \tanh([0.5, 1.5]) = 0.462 + 0.905 = 1.367$
• $e_3 = [1,1]^\top \tanh([1.5, 1.5]) = 0.905 + 0.905 = 1.810$

Softmax:

• $\exp(e_1) = 3.923, \exp(e_2) = 3.923, \exp(e_3) = 6.111$. Sum = 13.957.
• $\alpha_1 = 0.281, \alpha_2 = 0.281, \alpha_3 = 0.438$.

Context vector:

The decoder gets a context vector that leans toward $h_3$ (highest score) but still incorporates $h_1$ and $h_2$. If the training objective ends up favouring $h_2$ for producing the next target word, gradient will push $W_s$ and $W_h$ in directions that raise $e_2$ relative to the others, and next epoch's attention weight $\alpha_2$ will grow. Attention is differentiable soft alignment.

4 · Implement it — extending the S034 model

We need to modify the encoder to return all outputs (not just the final state), and rewrite the decoder to use attention at each step.

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class EncoderAll(nn.Module):
    def __init__(self, src_vocab, embed_dim=128, hidden_dim=256):
        super().__init__()
        self.embed = nn.Embedding(src_vocab, embed_dim, padding_idx=0)
        self.rnn = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
 
    def forward(self, src):
        e = self.embed(src)
        outputs, (h, c) = self.rnn(e)          # outputs: (B, T_src, H)
        return outputs, (h, c)
 
class BahdanauAttention(nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()
        self.Wh = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.Ws = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.v  = nn.Linear(hidden_dim, 1, bias=False)
 
    def forward(self, encoder_outputs, decoder_state, src_mask=None):
        # encoder_outputs: (B, T_src, H)
        # decoder_state:   (B, H)
        # src_mask:        (B, T_src) with 1 for real tokens, 0 for pad
        s = decoder_state.unsqueeze(1)          # (B, 1, H)
        scores = self.v(torch.tanh(self.Wh(encoder_outputs) + self.Ws(s)))
        scores = scores.squeeze(-1)             # (B, T_src)
        if src_mask is not None:
            scores = scores.masked_fill(~src_mask, float('-inf'))
        alpha = F.softmax(scores, dim=-1)       # (B, T_src)
        context = torch.bmm(alpha.unsqueeze(1), encoder_outputs).squeeze(1)  # (B, H)
        return context, alpha
 
class AttnDecoder(nn.Module):
    def __init__(self, tgt_vocab, embed_dim=128, hidden_dim=256):
        super().__init__()
        self.embed = nn.Embedding(tgt_vocab, embed_dim, padding_idx=0)
        self.attn = BahdanauAttention(hidden_dim)
        self.rnn = nn.LSTMCell(embed_dim + hidden_dim, hidden_dim)
        self.head = nn.Linear(hidden_dim * 2, tgt_vocab)   # concat output + context
 
    def forward_step(self, y_prev, state, encoder_outputs, src_mask):
        h_prev, c_prev = state                  # (B, H) each
        e = self.embed(y_prev).squeeze(1)       # (B, E)
        context, alpha = self.attn(encoder_outputs, h_prev, src_mask)
        rnn_input = torch.cat([e, context], dim=-1)
        h, c = self.rnn(rnn_input, (h_prev, c_prev))
        logits = self.head(torch.cat([h, context], dim=-1))
        return logits, (h, c), alpha

Two implementation subtleties:

1. Masking pad positions with -inf. If the source has padding, the encoder still emits (garbage) hidden states at pad positions. We must not let attention weight lie on them. Setting the pre-softmax score to $-\infty$ makes the softmax return exactly 0 for those positions. torch.masked_fill(~src_mask, float('-inf')) is the idiom.
2. torch.bmm(alpha.unsqueeze(1), encoder_outputs). This is the vectorised weighted-sum: (B, 1, T_src) @ (B, T_src, H) = (B, 1, H). Do not write a Python loop.

4.1 The training loop is almost unchanged

The Seq2Seq.forward from S034 becomes:

class Seq2SeqAttn(nn.Module):
    def __init__(self, src_vocab, tgt_vocab, embed_dim=128, hidden_dim=256):
        super().__init__()
        self.encoder = EncoderAll(src_vocab, embed_dim, hidden_dim)
        self.decoder = AttnDecoder(tgt_vocab, embed_dim, hidden_dim)
        self.tgt_vocab = tgt_vocab
 
    def forward(self, src, src_mask, tgt, teacher_forcing_ratio=1.0):
        B, T_tgt = tgt.shape
        enc_outputs, (h, c) = self.encoder(src)
        state = (h.squeeze(0), c.squeeze(0))
        y_prev = tgt[:, :1]
        all_logits = torch.zeros(B, T_tgt-1, self.tgt_vocab, device=src.device)
        all_alpha = torch.zeros(B, T_tgt-1, enc_outputs.size(1), device=src.device)
        for t in range(1, T_tgt):
            logits, state, alpha = self.decoder.forward_step(y_prev, state, enc_outputs, src_mask)
            all_logits[:, t-1] = logits
            all_alpha[:, t-1] = alpha
            use_teacher = torch.rand(1).item() < teacher_forcing_ratio
            y_prev = tgt[:, t:t+1] if use_teacher else logits.argmax(-1, keepdim=True)
        return all_logits, all_alpha

Loss and optimiser are identical to S034. Training time per step goes up by maybe 30% (extra MLP + softmax per step) but on the same number of steps quality is dramatically better.

4.2 The length curve, revisited

Re-run the S034 §9 length-vs-accuracy experiment with the attention model. Typical output:

src_len= 1  acc=1.000
src_len= 5  acc=0.997
src_len=10  acc=0.985
src_len=15  acc=0.972
src_len=20  acc=0.958
src_len=30  acc=0.941
src_len=40  acc=0.912

Compare to the S034 vanilla seq2seq curve, which fell to 0.31 at length 11+. Attention took the flat-line-then-cliff shape and turned it into a gentle slope. This is why Bahdanau et al.'s paper was so influential: the improvement is not marginal — it fixes the class of failure that motivated the paper in the first place.

5 · Attention weights as alignment — a free superpower

all_alpha is a (B, T_tgt, T_src) tensor of attention weights. For any input-output pair, alpha[b, t, i] is "how much did the model look at source position $i$ when producing target position $t$?"

Plot as a heatmap:

import matplotlib.pyplot as plt
import numpy as np
 
# after inference on one sentence pair
alpha = model_alpha[0].cpu().numpy()      # (T_tgt, T_src)
fig, ax = plt.subplots(figsize=(6, 6))
ax.imshow(alpha, cmap='Blues')
ax.set_xticks(range(len(src_tokens))); ax.set_xticklabels(src_tokens, rotation=45)
ax.set_yticks(range(len(tgt_tokens))); ax.set_yticklabels(tgt_tokens)
plt.tight_layout(); plt.show()

For a well-trained English → French model, you'll see a bright roughly-diagonal band (English word 3 attends most to French word 3-ish), with off-diagonal spikes exactly where you'd expect grammatical reordering (French subject-verb-object may require the model to attend to English word 5 when producing French word 2). Alignment falls out of the attention weights without any supervision. In pre-neural statistical MT, learning word alignment was an entire separate model (IBM Model 4, then HMM alignment, then discriminative aligners). Attention gave it away for free as a byproduct.

Bahdanau's Figure 3 in the 2015 paper is one of the most reproduced images in ML — a heatmap showing beautifully sensible English-French alignment learned by pure attention on parallel text.

6 · Luong (multiplicative / dot-product) attention

Minh-Thang Luong et al. (Stanford, EMNLP 2015) simplified Bahdanau by observing that the extra MLP is unnecessary — a plain dot product between encoder states and the decoder state gives essentially the same performance and is much cheaper. Luong compared three "score functions":

• dot — $e_{t,i} = h_i^\top s_t$
• general — $e_{t,i} = h_i^\top W_a s_t$
• concat / additive — the Bahdanau formulation

The dot and general variants have no per-position MLP; they're a single matmul across the whole sequence. On GPUs this is dramatically faster. Modern transformers use a lightly modified version — scaled dot-product attention — which is:

The $\sqrt{d}$ scaling comes from Vaswani et al. 2017: without it, large hidden dimensions produce very large dot products that push softmax into saturated (near-one-hot) regimes, killing gradient. Divide by $\sqrt{d}$ and the pre-softmax scores stay in a reasonable range regardless of $d$.

class LuongDotAttention(nn.Module):
    def forward(self, encoder_outputs, decoder_state, src_mask=None):
        # encoder_outputs: (B, T_src, H), decoder_state: (B, H)
        d = encoder_outputs.size(-1)
        scores = torch.bmm(encoder_outputs, decoder_state.unsqueeze(-1)).squeeze(-1)
        scores = scores / (d ** 0.5)                # scaled dot-product
        if src_mask is not None:
            scores = scores.masked_fill(~src_mask, float('-inf'))
        alpha = F.softmax(scores, dim=-1)
        context = torch.bmm(alpha.unsqueeze(1), encoder_outputs).squeeze(1)
        return context, alpha

Zero learned parameters in the attention module itself. Drop-in replacement for Bahdanau. On our toy task the two perform indistinguishably. On real NMT the difference is small (~0.1 BLEU) but Luong wins on training and inference speed by 20–40%.

7 · The Q/K/V abstraction — the doorway to the transformer

Look at what attention actually does, purely mechanically:

• You have a query $q$ (the decoder state).
• You have a set of keys $k_i$ (the encoder states, in Luong; a transformed version, in transformers).
• You have a set of values $v_i$ (also the encoder states, or a transformed version).
• You compute similarity between $q$ and each $k_i$, softmax to get weights, and return $\sum_i \alpha_i v_i$.

This is soft associative-memory lookup. Think of it as dict.get(query, ...) where the dictionary matches fuzzily and returns a blended average of its values weighted by how well each key matches the query.

Vaswani et al. 2017 recognised that the decoder state and encoder state don't have to play distinct roles. In self-attention, every token in a sequence produces its own $(q, k, v)$ triple (via three learned linear projections), and every token attends to every other token in the same sequence:

That is the entire mechanism of the transformer. Every LLM you've heard of runs this operation dozens of times per layer, dozens of layers deep, with millions of tokens per training example, using specialised kernels (FlashAttention, FlashAttention-2, FlashAttention-3 by Tri Dao's group in 2022-2024) to make it fast.

Bahdanau attention is the ancestor. Luong attention is the simplification. Self-attention is the generalisation. All three are the same operation with different inputs.

8 · Modern 2024–2025 twist — attention variants, alternatives, revival of recurrence

Attention has not stood still since 2017. The recent frontier:

8.1 Faster and longer attention

• FlashAttention (Dao et al. 2022, arXiv:2205.14135) — IO-aware GPU implementation that computes exact attention with 2–4× wall-clock speedup by keeping intermediates in SRAM instead of HBM. Now baked into PyTorch's scaled_dot_product_attention.
• FlashAttention-2 (Dao 2023) and FlashAttention-3 (Shah, Bikshandi, Zhang, Thakkar, Ramani, Dao 2024, arXiv:2407.08608) — squeeze more efficiency on Hopper (H100) hardware, hitting 75% of theoretical peak on FP16 matmul.
• Ring Attention (Liu, Zaharia, Abbeel 2023, arXiv:2310.01889) — distributes attention across devices to reach million-token contexts.
• Blockwise Parallel Transformers — used in Gemini and Claude for very long context (>1M tokens on 2024-2025 flagship models).

8.2 Long-context tricks on top of attention

• RoPE — Rotary Positional Embedding (Su et al. 2021, arXiv:2104.09864) — the positional encoding used in Llama, Mistral, DeepSeek, Qwen. We'll cover it in S042.
• YaRN (Peng et al. 2023, arXiv:2309.00071) and LongRoPE (Ding et al. 2024, arXiv:2402.13753) — extend RoPE's effective context length past pretraining without retraining.
• Multi-head Latent Attention (MLA) (DeepSeek-V2 & V3, 2024) — compresses the KV cache to a small latent representation, saving ~10× memory during inference. Central to DeepSeek's cost efficiency.

8.3 Beyond attention — the SSM comeback

The most exciting sequence-modelling story of 2024–2025 is that recurrence is back as a serious challenger to attention. Structured State Space Models (SSMs) provide a fully parallelisable, linear-cost alternative:

• Mamba (Gu & Dao 2023, arXiv:2312.00752) — selective SSM matching transformer quality up to 3B params with linear compute in sequence length.
• Mamba-2 (Dao & Gu 2024, arXiv:2405.21060, ICML 2024) — 2–8× faster than Mamba via "state space duality", scaling to 8B params.
• Jamba-1.5 (AI21 Labs 2024, arXiv:2408.12570) — hybrid Mamba + attention + MoE architecture; Jamba-1.5-Large (94B active / 398B total) hits Llama-3-70B quality with 256k context.
• RWKV-7 "Goose" (BlinkDL et al. 2025) — pure-recurrent architecture matching Llama-3-8B on many benchmarks, constant memory per token during inference.
• xLSTM-7B (Beck, Pöppel, Hochreiter et al. 2024, arXiv:2405.04517) — Sepp Hochreiter's own team's LSTM revival, released as a 7B open-weight model in September 2024.

The pendulum is swinging back. From 2017 to 2023, "attention is all you need" was gospel. Since late 2023 the frontier has been "attention is expensive at long context; can we get 90% of the quality with a parallelisable recurrence?". The answer looks like yes. Expect the 2026 large model landscape to be a mix of pure transformers (GPT, Claude), pure/mostly SSM (Mamba-3, RWKV-8), and hybrid (Jamba, Griffin) architectures.

9 · Attention as explanation — a healthy debate

Bahdanau's alignment heatmaps looked so interpretable that the ML community briefly treated attention weights as explanations for model decisions. Jain & Wallace 2019 (Attention is not Explanation) and Wiegreffe & Pinter 2019 (Attention is not not Explanation) had a public back-and-forth, with the modern consensus: attention weights describe what the model attended to, but are neither necessary nor sufficient for explaining why it made a particular prediction — many different attention distributions can produce identical outputs, and adversarially manipulated attention weights don't necessarily change predictions.

The 2024–2025 mechanistic interpretability line (Anthropic's Scaling Monosemanticity; Ferrando et al. 2024 A Primer on the Inner Workings of Transformer-based Language Models) has largely moved past raw attention to sparse autoencoders and attribution methods — but the aesthetic pull of Bahdanau's heatmaps remains. They will always be the first interpretable window we had into what a neural sequence model was doing.

10 · War stories

11 · Try it yourself

12 · Recap in five sentences

Bahdanau attention (Bahdanau, Cho, Bengio 2015) fixes the seq2seq fixed-context-vector bottleneck by computing a fresh context vector at every decoding step, as a softmax-weighted sum of every encoder hidden state, with the weights coming from a small MLP that scores each encoder state against the current decoder state. Luong 2015 simplified the scoring to a scaled dot product; that variant is what Vaswani et al. 2017 used inside the transformer, generalised to self-attention (query, key, value all from the same sequence). Attention weights give free interpretable soft alignments — the reason Bahdanau's Figure 3 became one of the most reproduced images in ML — though "attention as explanation" has been rightly complicated by later work. The 2024–2025 frontier is a mix of ever-faster attention kernels (FlashAttention-3, MLA, Ring Attention, long-context RoPE extensions) and a genuine recurrence comeback (Mamba, Mamba-2, Jamba, xLSTM, RWKV-7) that competes with attention on long-context tasks. You now understand the mechanism that every modern language model rests on, and Module M07 will build it up from scratch into a full GPT.

🧠 Retention scaffold

One-line summary (write it in your own words): ________________________________________________

Spaced review: re-read §3 (equations) and §7 (Q/K/V abstraction) in 24 hours. Revisit the full session on day 7. In week 2, revisit §8 (modern variants + SSM comeback) as the perfect bridge into M07.

Next session (S036): we cross into Module M07 — self-attention from scratch. Every idea from today gets applied within a single sequence, and we start building the transformer block that will power everything from S036 through the capstone.

Sticky note (keep on your desk): Attention = weighted sum where weights come from learned similarity. Query asks, keys match, values are pooled. Every LLM operation traces back to this.

Further reading

• Founding papers. Bahdanau, Cho, Bengio 2015 arXiv:1409.0473; Luong, Pham, Manning 2015 arXiv:1508.04025; Vaswani et al. 2017 Attention Is All You Need.
• Fast attention kernels. Dao et al. 2022 FlashAttention; Shah et al. 2024 FlashAttention-3.
• Long-context extensions. Su et al. 2021 RoPE; Peng et al. 2023 YaRN; Ding et al. 2024 LongRoPE; Liu et al. 2023 Ring Attention.
• Structured recurrence revival. Gu & Dao 2023 Mamba; Dao & Gu 2024 Mamba-2; AI21 2024 Jamba-1.5; Beck et al. 2024 xLSTM; Peng et al. 2023 RWKV; DeepMind 2024 Griffin / RecurrentGemma.
• Attention interpretability debate. Jain & Wallace 2019 Attention is not Explanation; Wiegreffe & Pinter 2019 Attention is not not Explanation; Ferrando et al. 2024 Primer on Transformer Internals.
• Modern efficient inference. DeepSeek-V3 tech report (2024) — Multi-head Latent Attention explanation. Kwon et al. 2023 vLLM for KV cache management.

Module M06 complete — you now stand at the door of the transformer

You've come a long way. From nn.Embedding in S031 to a full attention-augmented seq2seq in S035. You understand the vanishing gradient, the LSTM's fix, the seq2seq bottleneck, and the attention mechanism that fixes it. You know the specific reasons transformers replaced RNNs in 2017 and the specific reasons recurrence is fighting back in 2024–2025.

Module M07 (S036–S043) rebuilds everything you learned today into a decoder-only transformer — nanoGPT-style — trained on TinyShakespeare. You'll write self-attention, multi-head, positional encodings, transformer blocks, KV cache, RoPE, and FlashAttention integration from scratch. By the end of M07 you will have implemented (a small version of) the architecture of GPT-4.

Previous: ← DL S034 · Next: DL S036 →