NLP with RNNs + Attention

Before any model sees a word, the word has to become an integer. Tokenization is the choice of how to split text into pieces and assign each piece an ID. The pieces are called tokens.

Three approaches you should be able to name and compare:

Whitespace / word-level. Split on spaces and punctuation. One token per word. Simple, but the vocabulary explodes (English has hundreds of thousands of distinct words once you count inflections and proper nouns), out-of-vocabulary words are common, and you cannot generalise to misspellings.

Character-level. One token per Unicode code point. Vocabulary of ~100 (ASCII) to ~150,000 (all Unicode), no out-of-vocab problem, but sequence length explodes (a 100-character sentence is 100 tokens) and the model spends most of its compute on redundancy.

Subword (byte-pair encoding, WordPiece, SentencePiece). The dominant choice since GPT-2. Start with single characters or bytes. Repeatedly find the most frequent adjacent pair in the training corpus, merge them into a new token. After ~50k merges, common words become single tokens (the, and, of), rare words become subword sequences (tokenization → token, ization). Pre-trained models all use some variant: GPT uses byte-level BPE, BERT uses WordPiece, T5 and Llama use SentencePiece.

# A tiny byte-pair encoding pass
from collections import Counter

def get_pair_counts(seq: list[int]) -> Counter:
    return Counter(zip(seq, seq[1:]))

def merge(seq: list[int], pair: tuple, new_id: int) -> list[int]:
    out, i = [], 0
    while i < len(seq):
        if i < len(seq) - 1 and (seq[i], seq[i+1]) == pair:
            out.append(new_id); i += 2
        else:
            out.append(seq[i]); i += 1
    return out

def train_bpe(text: str, target_vocab: int = 1000) -> tuple[list[int], dict]:
    seq = list(text.encode("utf-8"))
    merges, next_id = {}, 256
    while next_id < target_vocab:
        pairs = get_pair_counts(seq)
        if not pairs: break
        best = pairs.most_common(1)[0][0]
        seq = merge(seq, best, next_id)
        merges[best] = next_id
        next_id += 1
    return seq, merges

This is the warm-up. another chapter has the full BPE deep-dive (with regex pre-tokenization, special tokens, the GPT-2 pattern). For now, internalise the idea: tokenization is a learned compression of bytes into integers, with a vocabulary size you choose, and the choice has downstream consequences.

Once you have token IDs, you turn each one into a vector. The lookup is exactly nn.Embedding(vocab_size, d_embed): a learned table of shape (vocab_size, d_embed) where each row is one token's vector.

What makes embeddings interesting is not the lookup; it is what the rows learn to encode. After training, similar words end up at similar positions in the embedding space. The canonical demonstration is the analogy structure: $\text{king} - \text{man} + \text{woman} \approx \text{queen}$ in word2vec's vector space. This is not a property anyone designed in. It emerges because the training objective (predict surrounding context from a target word, or vice versa) forces the geometry to encode contextual co-occurrence patterns, and those patterns turn out to capture meaningful semantic axes.

Three pre-transformer ways to learn embeddings:

word2vec (Mikolov et al. 2013). Two variants: skip-gram (given a target word, predict the surrounding context words) and CBOW (given context, predict the target). Both train a shallow network. The middle-layer weights become the embeddings.

GloVe (Pennington, Socher, Manning 2014). Factorise the global word-word co-occurrence matrix (after log-transform) into low-rank embedding matrices. A matrix-factorisation reframing of word2vec.

fastText (Bojanowski et al. 2017). Like word2vec, but each word is represented as the sum of its character n-gram embeddings — an n-gram being a run of $n$ consecutive characters, so where at $n=3$ contributes <wh, whe, her, ere, re>. Handles out-of-vocabulary words (you can embed any string by summing its n-gram embeddings) and learns better representations for morphologically rich languages.

model = Word2Vec(sentences, vector_size=100, window=5,
                 sg=1, negative=5, min_count=1)
vec = model.wv["cat"]                 # learned input-side embedding
model.wv.most_similar(positive=["king", "woman"], negative=["man"])

class SkipGramNeg(nn.Module):
    def __init__(self, vocab_size, d_embed):
        super().__init__()
        self.in_emb = nn.Embedding(vocab_size, d_embed)
        self.out_emb = nn.Embedding(vocab_size, d_embed)

    def forward(self, target, context, neg):
        v_t = self.in_emb(target)          # (B, D)
        v_c = self.out_emb(context)        # (B, D)
        v_n = self.out_emb(neg)            # (B, k, D)
        pos_score = (v_t * v_c).sum(dim=-1)
        neg_score = torch.bmm(v_n, v_t.unsqueeze(-1)).squeeze(-1)
        pos_loss = -torch.log(torch.sigmoid(pos_score) + 1e-9)
        neg_loss = -torch.log(torch.sigmoid(-neg_score) + 1e-9).sum(dim=-1)
        return (pos_loss + neg_loss).mean()

The two-embedding-tables structure is structural: in_emb is the "what word am I right now" representation, out_emb is the "what word should I predict" representation. You typically only keep in_emb as the final embedding.

For modern usage, you almost never train word embeddings from scratch. You either use pre-trained ones (GloVe, fastText vectors are downloadable in seconds) or you let your end-to-end model (BERT, GPT) learn them as part of the architecture. The pre-transformer word-embedding chapter ends here. The structural intuitions remain useful.

another chapter already introduced the encoder-decoder. Quick recap because the rest of this chapter modifies it.

The encoder is an RNN (typically GRU or LSTM, often bidirectional) that reads the source sequence and produces a hidden state at every position. The "context" passed to the decoder is, in the simplest seq2seq, just the encoder's final hidden state.

The decoder is another RNN that starts from the encoder's context as its initial hidden state and generates the target sequence one token at a time. At training time, it is teacher-forced: it sees the ground-truth previous token at each step. At inference time, it sees its own previous prediction.

class Seq2Seq(nn.Module):
    def __init__(self, src_vocab: int, tgt_vocab: int, d: int = 256):
        super().__init__()
        self.src_emb = nn.Embedding(src_vocab, d)
        self.tgt_emb = nn.Embedding(tgt_vocab, d)
        self.encoder = nn.GRU(d, d, batch_first=True)
        self.decoder = nn.GRU(d, d, batch_first=True)
        self.head = nn.Linear(d, tgt_vocab)

    def forward(self, src_ids, tgt_ids):
        _, h_enc = self.encoder(self.src_emb(src_ids))
        out, _ = self.decoder(self.tgt_emb(tgt_ids), h_enc)
        return self.head(out)

This is the architecture Bahdanau et al. 2014 started from. They identified a specific failure mode: when the source sentence is long (>20 tokens or so), the single final hidden state cannot carry all the information the decoder needs. Translation quality degrades sharply with source length. The next section is the fix.

The problem with vanilla seq2seq: a single fixed-size context vector is asked to encode an arbitrary-length source sentence. For short sources, fine. For long sources, the context vector becomes a bottleneck.

Bahdanau, Cho, Bengio 2014 ("Neural Machine Translation by Jointly Learning to Align and Translate") changed the framing. Instead of compressing the source into one vector and giving that to the decoder, keep all the encoder's hidden states and let the decoder, at every step, choose which encoder states to look at. The choice is a learned soft selection: at each decoder step, compute a weight for each source position, softmax them, and take a weighted average of the encoder hidden states.

Jay Alammar's 01-explorables/jalammar-visualizing-neural-machine-translation-mechanics-of-seq2seq-models-with-attention §lets-pay-attention has the cleanest visualisation of this: the encoder produces $T_{\text{src}}$ hidden states, the decoder produces its own hidden states, and at each decoder step a distribution over source positions is computed and used to mix the encoder states into a position-specific context vector.

That mixing operation is attention. Everything after this in the chapter (and most of another chapter) is variations on it.

The first attention mechanism (Bahdanau et al. 2014) computes alignment scores between the decoder's current hidden state $s_{t-1}$ and each encoder hidden state $h_i$ using an additive function:

where $W_a, U_a$ are learned matrices and $v_a$ is a learned vector. The score is a function of the concatenation of the two states, projected through a tanh nonlinearity, then projected down to a scalar.

The scores are softmaxed over source positions to get attention weights:

The context vector at decoder step $t$ is the weighted average of encoder hidden states:

The decoder uses $c_t$ along with its own hidden state to produce the next output. Concretely, the decoder input becomes $[\text{prev\_token\_embedding}, c_t]$, and the output projection sees $[s_t, c_t]$.

# There is no library call equal to the scratch block below.
# nn.MultiheadAttention is DOT-PRODUCT attention (QK^T/sqrt d), not additive:
mha = nn.MultiheadAttention(embed_dim=d, num_heads=1, batch_first=True)
ctx, w = mha(query=s.unsqueeze(1), key=h, value=h)  # different math: no tanh, no v

class BahdanauAttention(nn.Module):
    def __init__(self, encoder_hidden, decoder_hidden):
        super().__init__()
        self.W = nn.Linear(decoder_hidden, decoder_hidden, bias=False)
        self.U = nn.Linear(encoder_hidden, decoder_hidden, bias=False)
        self.v = nn.Linear(decoder_hidden, 1, bias=False)

    def forward(self, s, h):                      # s: (B, dec_h); h: (B, T_src, enc_h)
        s_exp = s.unsqueeze(1)                     # (B, 1, dec_h)
        scores = self.v(torch.tanh(self.W(s_exp) + self.U(h))).squeeze(-1)  # (B, T_src)
        weights = F.softmax(scores, dim=-1)
        context = torch.bmm(weights.unsqueeze(1), h).squeeze(1)             # (B, enc_h)
        return context, weights

The "additive" name comes from the fact that the scoring function adds the projected query and key inside the tanh. This is in contrast to "multiplicative" attention, which takes a dot product. Additive attention was first because Bahdanau's group came from a recurrent-network background where the natural compositional operation was concatenation-then-linear. The multiplicative form turned out to be cheaper and equally accurate, and it is what won.

Luong, Pham, Manning 2015 ("Effective Approaches to Attention-based Neural Machine Translation") proposed two simplifications.

First, replace the additive scoring function with a multiplicative one:

The dot product is cheaper than the additive form (one matmul instead of two MLPs) and trains comparably well. This is the operation that became scaled dot-product attention in Vaswani 2017.

Second, change when the attention is computed. Bahdanau used the previous decoder state $s_{t-1}$ to attend over encoder states, then combined the context with the previous-token embedding to compute the new decoder state $s_t$. Luong's "global" formulation computed $s_t$ first (from the previous decoder state and the new input token) and then used $s_t$ to attend over encoder states. The difference is small in practice but the Luong order is cleaner.

class LuongAttention(nn.Module):
    def __init__(self, d: int, score_type: str = "general"):
        super().__init__()
        self.score_type = score_type
        if score_type == "general":
            self.W = nn.Linear(d, d, bias=False)

    def forward(self, s: torch.Tensor, h: torch.Tensor) -> tuple:
        # s: (B, d). h: (B, T_src, d).
        if self.score_type == "dot":
            scores = torch.bmm(s.unsqueeze(1), h.transpose(1, 2)).squeeze(1)
        else:   # general
            scores = torch.bmm(self.W(s).unsqueeze(1), h.transpose(1, 2)).squeeze(1)
        weights = torch.softmax(scores, dim=-1)
        context = torch.bmm(weights.unsqueeze(1), h).squeeze(1)
        return context, weights

Bahdanau and Luong used different vocabularies, but the operations were structurally the same: compute compatibility scores between a "thing you're querying with" and "things you're querying over", softmax the scores, take a weighted sum of "values associated with each queried thing".

The unifying vocabulary that won (Vaswani et al. 2017, retroactively applied to all of attention) is:

• Query (Q): what the current step is asking
• Key (K): what each source position offers as a matching feature
• Value (V): what each source position contributes if matched

The attention operation:

In Bahdanau / Luong attention, Q is the decoder hidden state, K and V are both the encoder hidden states (often identical). In self-attention (next section), Q, K, V are all linear projections of the same input. In transformer cross-attention, Q comes from the decoder and K, V come from the encoder. Same equation, different sources.

def scaled_dot_product_attention(Q, K, V, mask=None):
    """Q: (B, T_q, d_k). K, V: (B, T_k, d_k). mask: optional (T_q, T_k) bool."""
    d_k = Q.size(-1)
    scores = Q @ K.transpose(-2, -1) / (d_k ** 0.5)
    if mask is not None:
        scores = scores.masked_fill(mask, float("-inf"))
    weights = torch.softmax(scores, dim=-1)
    return weights @ V, weights

This is the operation. Everything in another chapter is variations on it: stacking multiple heads, applying it to itself, masking it for causality. Internalise it now.

This is the moment the field pivots. Self-attention is the generalisation of Bahdanau-style attention where Q, K, and V all come from the same sequence. Every position attends to every other position in its own sequence, in parallel, with no recurrence.

The three properties that mattered, named only briefly here because they get earned in another chapter:

• No recurrence. One matmul per layer, parallelisable across all positions.
• Constant path length. Position $i$ can attend to position $j$ directly regardless of $|i-j|$. No vanishing-gradient-through-sequence-depth.
• Permutation equivariant. Self-attention has no notion of order on its own; you bolt positional embeddings on as a separate step (another chapter covers the full menu: absolute, sinusoidal, RoPE).

The first self-attention layers in NLP were used as a refinement on top of RNN outputs (Cheng, Dong, Lapata 2016, "Long Short-Term Memory-Networks for Machine Reading"). Vaswani et al. 2017 ("Attention Is All You Need") replaced the RNN entirely. The latter is what stuck.

The full math (scaled dot-product, multi-head, causal mask, batched implementation, shape-by-shape debug walkthrough) is derived from scratch in Ch 15 §3. The PyTorch class with proper multi-head reshaping lives there too. Here, the only thing to internalise is the equation in its barest form:

This is the operation another chapter builds the entire transformer architecture out of. The next chapter is essentially "now multiply this self-attention block by 12, stack with FFNs and LayerNorms, and train on text".

A subtlety that bites everyone who trains seq2seq for the first time.

At training time, the decoder is fed the ground-truth previous token at each step. This is called teacher forcing. It makes training fast and stable: the decoder always sees a correct context, so gradient signal is clean, and the loss is computed per-position independently.

At inference time, the decoder is fed its own previous prediction. If the model makes a mistake at position 3, that mistake feeds into the input at position 4, which compounds. By position 30, the model is operating in a regime it was never trained on (its own error-corrupted history). This is exposure bias.

Fixes are heuristic:

• Scheduled sampling (Bengio et al. 2015). During training, with probability $p$ feed the model its own prediction instead of the ground truth. Anneal $p$ from 0 toward 1 over training. Mostly works, mildly debated.
• Beam search (next section). Reduces the impact of a single bad prediction by maintaining multiple candidates.
• Train-with-noise tricks (Sample $p\%$ of decoder inputs from the model's own distribution). Various flavours.

In practice, large modern models (transformers) ignore exposure bias mostly because they are large enough that the discrepancy between teacher-forced training and free-run inference is small. The problem is most acute for small RNN-based seq2seq systems.

Greedy decoding takes the argmax at each step. It is fast and often wrong: a locally-optimal first token can lock you into a globally bad sequence. Beam search keeps the top-$k$ candidate sequences at every step and only commits at the end.

The algorithm:

1. Start with $k$ copies of the start token, each with score 0.
2. At each step, for each of the $k$ candidates, expand it by every possible next token. You now have $k \cdot V$ candidates.
3. Score each: existing-score + log-prob of the new token.
4. Keep the top $k$ by score.
5. Stop when all $k$ candidates have emitted end-of-sequence.

import math

@torch.no_grad()
def beam_search(model, start_id: int, end_id: int, max_len: int = 50, k: int = 4):
    beams = [([start_id], 0.0)]
    finished = []
    for _ in range(max_len):
        new_beams = []
        for seq, score in beams:
            if seq[-1] == end_id:
                finished.append((seq, score))
                continue
            idx = torch.tensor([seq])
            logits = model(idx)[:, -1, :]
            log_probs = torch.log_softmax(logits, dim=-1).squeeze(0)
            top_logp, top_idx = log_probs.topk(k)
            for lp, tid in zip(top_logp.tolist(), top_idx.tolist()):
                new_beams.append((seq + [tid], score + lp))
        new_beams.sort(key=lambda x: x[1], reverse=True)
        beams = new_beams[:k]
        if not beams: break
    finished.extend(beams)
    # Length-normalised re-ranking (otherwise beam search prefers short sequences)
    finished.sort(key=lambda x: x[1] / len(x[0]), reverse=True)
    return finished[0]

Beam search is rarely used for open-ended language generation today (it produces too deterministic and repetitive outputs). It remains standard for machine translation, summarisation, and other tasks where there is a "correct" answer the model should converge toward.

From Bahdanau (2014) to Vaswani (2017), the field built attention on top of RNN backbones. The encoder was an RNN, the decoder was an RNN, and attention was a bolt-on between them. The RNN did the sequence modeling; attention did the cross-sequence alignment.

Vaswani et al. 2017's insight ("Attention Is All You Need") was that you could remove the RNN entirely. If self-attention can mix information across all positions in one matmul, you do not need recurrence to model sequential dependencies. You can build the entire encoder out of stacked self-attention + feedforward layers, the entire decoder out of stacked self-attention + cross-attention + feedforward layers, and train the whole thing in parallel across positions.

The payoffs:

• Training parallelises across sequence positions, so per-step wall-clock drops by orders of magnitude on long sequences.
• No vanishing gradient through depth-of-sequence.
• The model can attend to any position directly. This turned out to encode dependencies BiLSTM models could only approximate.

The cost: $O(L^2)$ memory and compute in the sequence length, versus $O(L)$ for an RNN. For sequences up to a few thousand tokens (the regime where most NLP lives), this is a win. For very long sequences (DNA, code, long documents), specialised attention variants (Longformer, BigBird, FlashAttention) and state-space models (Mamba) reclaim some of the linear-time territory.

This is the chapter's exit door. another chapter picks up here: same attention math, no RNN, full transformer.

Bahdanau attention's weight matrix $\alpha_{t,i}$ has a useful interpretation: it is a soft alignment between target positions and source positions. If you display it as a heatmap with source positions on the x-axis and target positions on the y-axis, a translation model trained well shows a near-diagonal pattern with some scrambling.

Jay Alammar's 01-explorables/jalammar-visualizing-neural-machine-translation-mechanics-of-seq2seq-models-with-attention §alignment-visualization shows this for English-to-French. When the model emits "économique" it attends heavily to "economic" in the source. When it emits "européenne" it attends to "European". The alignment is interpretable: you can audit a translation model by inspecting the attention weights.

This interpretability property partially transferred to transformers and partially did not. Transformer attention heads have many roles (positional, syntactic, semantic, copy-mechanism). Some are interpretable in the Bahdanau-alignment sense; many are not. The mech-interp literature (another chapter) is largely about figuring out what transformer attention heads actually do.

# Visualise attention weights for a trained Bahdanau model
import matplotlib.pyplot as plt

def plot_attention(weights: torch.Tensor, src_tokens: list[str], tgt_tokens: list[str]):
    """weights: (T_tgt, T_src)."""
    fig, ax = plt.subplots(figsize=(len(src_tokens), len(tgt_tokens)))
    ax.imshow(weights.cpu().numpy(), aspect="auto", cmap="viridis")
    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)
    ax.set_xlabel("source"); ax.set_ylabel("target")
    return fig