What I Finally Understood About Transformers

A few things finally clicked for me after re-reading the original paper and building a minimal decoder-only Transformer from scratch.

Dynamic Lookup Tables

The core realization is that attention mechanisms are not magical; they are dynamic lookup tables where keys and queries determine how much weight every token pays to every other token. Once you visualize the Q, K, V matrix multiplications as soft routing, the elegance of multi-head self-attention becomes strikingly simple.

import torch
import torch.nn as nn
import math

class SelfAttention(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.d_model = d_model
        
    def forward(self, Q, K, V):
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_model)
        weights = torch.softmax(scores, dim=-1)
        return torch.matmul(weights, V)

Key Insights

Here are the key takeaways from implementing multi-head self-attention:

• Soft Routing: Every token dynamically chooses which other tokens in the context window to attend to.
• Scaling Factor: Dividing by sqrt(d_model) prevents gradients from vanishing in large dimension spaces.
• Parallel Computation: Unlike RNNs, sequence processing happens simultaneously across all tokens during training.

Attention is not about paying attention to everything; it is about learning what to ignore.

Common Language Implementations

Below is a quick snippet in TypeScript for calculating standard dot-product attention scores:

function computeAttention(query: number[], key: number[]): number {
  let score = 0;
  for (let i = 0; i < query.length; i++) {
    score += query[i] * key[i];
  }
  return score / Math.sqrt(query.length);
}