Learn LLM

The problem with fixed windows

Recurrent neural networks process a sentence one token at a time. They keep a hidden state that theoretically remembers everything, but in practice gradients vanish and the model forgets distant words. Convolutional networks see only a fixed local window.

Attention solves this by letting every token look directly at every other token. There is no fixed window and no sequential bottleneck. The model learns how much each token should care about every other token.

Query, Key, Value

For each token, we create three vectors from its embedding: a query, a key, and a value. The query represents what the token is looking for. The key represents what the token offers. The value is the actual information to be passed on.

Q = X W_Q, \quad K = X W_K, \quad V = X W_V

X is the input sequence matrix where each row is a token embedding. The weight matrices W_Q, W_K, W_V are learned. Their dimensions are chosen so that every token ends up with a query, key, and value vector.

Scaled dot-product attention

To decide how much token i attends to token j, we compare the query of i with the key of j. The dot product measures similarity. We scale by the square root of the key dimension to keep gradients stable, then apply softmax.

\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V

The softmax produces a probability distribution over positions for each query. We use these weights to take a weighted sum of the value vectors. The result is a new representation for token i that blends information from the whole sequence.

Self-attention heatmap

Type a sentence and watch the attention scores. Brighter cells mean the query token pays more attention to that key token.

Visualization attention is coming soon.

Multi-head attention

One attention function can only capture one kind of relationship. Multi-head attention runs h attention functions in parallel, each with its own learned projections. The outputs are concatenated and projected again.

\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W_O

Different heads can specialize: one may track pronoun references, another may track verb-subject agreement, another may track nearby punctuation. The model learns these specializations from data.

Causal masking for generation

When generating text, a token should not look at future tokens. A causal or autoregressive mask sets attention scores for future positions to negative infinity before the softmax, forcing the model to use only past tokens.

M_{ij} = \begin{cases} 0 & i \ge j \ -\infty & i < j \end{cases}

import numpy as np

def softmax(x):
    e = np.exp(x - np.max(x, axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def causal_attention(Q, K, V):
    scores = Q @ K.T / np.sqrt(Q.shape[-1])
    # mask future positions
    mask = np.triu(np.ones_like(scores), k=1) * -1e9
    scores = scores + mask
    weights = softmax(scores)
    return weights @ V

Why scaled?

When the dimension of the query and key vectors is large, their dot products have high variance. Large values push softmax into a very sharp distribution, which gives tiny gradients. Dividing by sqrt(d_k) keeps the distribution soft and gradients healthy.

Key takeaway

Self-attention compares every token to every token using queries, keys, and values. It captures long-range dependencies in parallel. Causal masking makes it suitable for left-to-right text generation.