S107 · Attention Intuition — Why RNNs Failed, Why Attention Won

# mini_attention.py — a from-scratch self-attention layer with visualisation.
# Run: uv run mini_attention.py
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
 
 
class SingleHeadAttention(nn.Module):
    """One head of self-attention. Q, K, V from same input; attention weights returned."""
 
    def __init__(self, d_model: int, causal: bool = False):
        super().__init__()
        self.d_model = d_model
        self.causal = causal
        # Three linear projections — no bias, following the Transformer paper.
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
 
    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        # x: [B, T, d]. Project to Q, K, V.
        B, T, D = x.shape
        Q = self.W_q(x)  # [B, T, D]
        K = self.W_k(x)  # [B, T, D]
        V = self.W_v(x)  # [B, T, D]
 
        # Compute scaled dot-product attention scores.
        # scores[b, i, j] = how much token i attends to token j.
        scores = Q @ K.transpose(-2, -1)               # [B, T, T]
        scores = scores / math.sqrt(D)                 # scale — see Session S108 for why
 
        # Causal mask: token i cannot see token j > i.
        if self.causal:
            mask = torch.triu(torch.ones(T, T, device=x.device), diagonal=1).bool()
            scores = scores.masked_fill(mask, float("-inf"))
 
        # Softmax across the KEYS (last dim).
        weights = F.softmax(scores, dim=-1)            # [B, T, T]
 
        # Weighted sum of values.
        output = weights @ V                            # [B, T, D]
        return output, weights
 
 
def demo() -> None:
    torch.manual_seed(42)
 
    # A tiny vocabulary + embedding for demonstration.
    sentence = ["The", "cat", "sat", "on", "the", "mat"]
    vocab = {word: i for i, word in enumerate(sorted(set(sentence)))}
    ids = torch.tensor([[vocab[w] for w in sentence]])  # [1, 6]
 
    d_model = 16
    embed = nn.Embedding(len(vocab), d_model)
    attn = SingleHeadAttention(d_model, causal=False)
 
    x = embed(ids)                                        # [1, 6, 16]
    out, weights = attn(x)
    print(f"input shape:  {x.shape}")
    print(f"output shape: {out.shape}")
    print(f"weights shape: {weights.shape} (should be [1, 6, 6])")
 
    # Pretty-print the attention matrix.
    w = weights[0].detach().numpy()
    print("\nAttention pattern (rows = query token, cols = key token):")
    print(f"          {'  '.join(w_.rjust(4) for w_ in sentence)}")
    for i, word in enumerate(sentence):
        row = "  ".join(f"{v:.2f}" for v in w[i])
        print(f"  {word:<7}  {row}")
 
    # Compare causal vs non-causal.
    print("\n--- Causal (GPT-style) ---")
    attn_causal = SingleHeadAttention(d_model, causal=True)
    _, weights_c = attn_causal(x)
    w = weights_c[0].detach().numpy()
    for i, word in enumerate(sentence):
        row = "  ".join(f"{v:.2f}" for v in w[i])
        print(f"  {word:<7}  {row}")
    print("(Notice the upper-right triangle is 0 — the future is masked.)")
 
 
if __name__ == "__main__":
    demo()