Attention and language models

Prerequisites: Vectors, dot products, probabilities, and the neural-network chapter.

• Identify the shapes of Q, K, V, and the score matrix.
• Calculate a normalized weighted sum of values.
• Explain the causal training shift and its limits.

Tokens enter as vectors

A tokenizer maps text into a sequence of token identifiers. Tokens may be word pieces, bytes, or other units; one token is not necessarily one word. An embedding lookup maps each identifier to a learned vector. Positional information gives the model access to order. Two sequences with the same token counts can mean different things because their order differs.

A causal language model estimates a distribution over the next token from a prefix. For a sequence t1…tT, its joint probability factors into the product of P(t_i | t1…t_(i−1)), with any start-of-sequence convention included. Training commonly minimizes the negative logarithm of the probability assigned to observed next tokens.

Attention mixes value vectors

For n positions, let Q and K each have shape n×d_k and V have shape n×d_v. A scaled dot-product attention head forms an n×n score matrix QKᵀ/√d_k, applies a mask when needed, takes softmax separately across each row, and multiplies by V. The output has shape n×d_v. Multiple heads learn different projections. The original Transformer also includes feed-forward layers, residual connections, normalization, and positional encoding. See sections 3.1 through 3.5 of Attention Is All You Need.

Work a deliberately small example. Let q=[0,1], keys k1=[1,0] and k2=[0,1], and values v1=[6,0], v2=[0,3]. The unscaled dot products are [0,1]. Dividing by √2 gives [0,0.7071]. Exponentiating and normalizing produces weights approximately [0.3302,0.6698]. The output is 0.3302*[6,0]+0.6698*[0,3]≈[1.9814,2.0093].

The weights add to one for this head’s row before any attention dropout. Values, rather than keys, are the vectors combined. A dot product determines compatibility within the learned representation; it is not a universal measure of semantic truth.

from math import exp, sqrt

def softmax(scores):
    largest = max(scores)
    weights = [exp(s-largest) for s in scores]
    total = sum(weights)
    return [w/total for w in weights]

q = [0, 1]
keys = [[1, 0], [0, 1]]
values = [[6, 0], [0, 3]]
scores = [sum(a*b for a, b in zip(q, k))/sqrt(2) for k in keys]
weights = softmax(scores)
output = [sum(weight*v[j] for weight, v in zip(weights, values))
          for j in range(2)]
assert abs(sum(weights)-1) < 1e-12
assert abs(output[0]-1.981430704) < 1e-8
print([round(x, 4) for x in output])  # [1.9814, 2.0093]

Subtracting the largest score before exponentiation leaves the normalized distribution unchanged and avoids unnecessarily large exponentials. The code assumes at least one permitted key; an entirely masked row needs explicit handling in an implementation.

A causal mask prevents looking ahead

At input position i, causal self-attention allows keys at positions j≤i, including the current input token. Training targets are shifted: the output at position i predicts token i+1. Future keys receive a score of negative infinity before softmax, giving them zero weight. An encoder used for other tasks can instead attend bidirectionally. “Transformer” does not imply one universal mask.

For input [START, red, fox] and targets [red, fox, END], the position containing “red” can use START and red while predicting fox. It cannot read the fox input from the next position. Allowing that future token would leak the target and make training performance misleading.

Separate likelihood from verification

High probability under a language model means a continuation fits the learned conditional distribution. It does not prove that a statement is true, that an action is authorized, or that a cited document exists. That limitation follows from the objective: next-token loss evaluates probability assigned to observed text, not a general truth predicate.

Temperature rescales logits before sampling. Lower positive temperature concentrates probability around larger logits; it does not add new evidence. Greedy decoding chooses the largest logit and can repeat the same error consistently. Changing decoding is therefore different from checking a claim against a reliable source.

For a model evaluation, specify the task and inspect concrete failures. Keep factual correctness, source support, and task completion distinct. The attention calculation explains a mechanism inside a model; it does not establish the reliability of a complete application. Use the exercise to calculate a second attention row and then explain the causal boundary without relying on the formula alone.

Practice exercise

Take one query q=[1,0], keys k1=[1,0] and k2=[0,1], values v1=[2,0] and v2=[0,4], and key dimension d_k=2. Calculate the scaled scores, softmax weights, and output. Then mask the second key as a future position and repeat. Explain why changing the values changes the output without changing these attention weights.

Compare with the worked solution

Scaled scores are [1/√2,0]≈[0.7071,0]. Softmax gives [0.6698,0.3302]. The output is approximately [1.3395,1.3210]. With key 2 masked, its score is −infinity, its exponential is 0, and the weights become [1,0]; output=[2,0]. Weights in this calculation depend on Q and K. V supplies the vectors mixed using those weights. In a full model, all three are learned projections of hidden representations.

Check your understanding

What is the output when only the first key is allowed?

• [2, 0]
• [0, 4]
• [1, 2]
• [0, 0]

Check the reasoning

Softmax over one allowed score gives weight 1 to that value vector.

When a causal language model predicts token t+1 from position t, which input positions may that position attend to?

• Only t+1.
• Every position including future tokens.
• Positions up to and including t.
• No input positions.

Check the reasoning

The diagonal is allowed. Training labels are shifted so the target next token is unavailable to the position making that prediction.

Learn AI curriculum · Next: A reproducible experiment