Lesson 10

Attention scores, masking, and context

Why attention was a major change

An RNN carries one changing state through a sequence. Information from an early token must travel through every later step. Attention gives each position a more direct route to other positions. A token can compare itself with all allowed context positions and collect a weighted mixture of their information.

Consider the final character in h e l l. To guess the next character, the position asks which earlier representations are useful. Attention does not select one token with a hard rule. It creates non-negative weights that sum to one and uses them to average information.

Query, key, and value

Learned linear layers produce query, key, and value vectors from each token representation. For one head, query and key have feature width d. Their dot product measures alignment. Scores are divided by √d before softmax. This scaling keeps large feature widths from producing extremely sharp probabilities too early.

Worked example: one attention head

Predict the weighted sum: suppose one query gives weights [0.75, 0.25] over values [4, 0] and [0, 8]. The output is [3, 2]. Attention mixes values; it does not return the weights themselves.

Input has shape (B,T,C). Here there is one sequence, four positions, and six input features. Query, key, and value projections each use three features. Matrix multiplication between queries and transposed keys creates one score for every pair of positions: shape (B,T,T).

import math
import torch
from torch import nn

torch.manual_seed(10)
x = torch.randn(1, 4, 6)
query_layer = nn.Linear(6, 3, bias=False)
key_layer = nn.Linear(6, 3, bias=False)
value_layer = nn.Linear(6, 3, bias=False)

query = query_layer(x)
key = key_layer(x)
value = value_layer(x)

scores = query @ key.transpose(-2, -1) / math.sqrt(3)
weights = torch.softmax(scores, dim=-1)
context = weights @ value

print(tuple(scores.shape))
print(tuple(context.shape))
print(torch.allclose(weights.sum(dim=-1), torch.ones(1, 4)))

# Expected:
# (1, 4, 4)
# (1, 4, 3)
# True

Every row of weights belongs to one querying position. Every column refers to one source position. Softmax runs along the source-position axis, so each row sums to one. Multiplying weights by values produces one mixed value vector for every query position.

The causal mask

Check one row of the mask before checking the whole matrix. At position 1, only positions 0 and 1 may receive nonzero weight. If position 2 has any weight, the model can read a future answer during training.

GPT predicts the next token and must not see future tokens during training. A causal mask blocks entries above the score matrix diagonal. We replace forbidden scores with negative infinity before softmax. Their resulting probability becomes zero.

future = torch.triu(torch.ones(4, 4, dtype=torch.bool), diagonal=1)
causal_scores = scores.masked_fill(future, float("-inf"))
causal_weights = torch.softmax(causal_scores, dim=-1)

print(causal_weights[0, 0].tolist())
print(causal_weights[0, 2, 3].item())

# Expected:
# [1.0, 0.0, 0.0, 0.0]
# 0.0

Position zero can use only itself. Position two can use positions zero, one, and two, but not position three. BERT-style masked-token training is different: it needs context on both sides and does not apply a GPT causal mask. Padding masks are also different; they hide added padding tokens rather than future real tokens.

Multi-head attention

Several heads let the model learn several matching patterns. One might focus on nearby syntax while another tracks a repeated name. The heads are joined and projected back to the model width. The model can still learn redundant heads; the interpretation is a useful intuition, not a guarantee.

attention = nn.MultiheadAttention(
    embed_dim=6, num_heads=2, dropout=0.0, batch_first=True
)
output, all_weights = attention(x, x, x, need_weights=True)
print(tuple(output.shape))       # Expected: (1, 4, 6)
print(tuple(all_weights.shape))  # Expected: (1, 4, 4)

Try it

For two sequences, five positions, and key width four, what is the score tensor shape? Under a causal mask, how many source positions may query position three use when positions start at zero?

Recap

Attention compares queries with keys, scales the scores, normalizes them with softmax, and mixes values. Score and weight matrices have shape (B,T,T). Multiple heads learn separate projections. Causal masking prevents next-token models from reading future answers; bidirectional encoder tasks use both directions instead.

Reference: PyTorch MultiheadAttention documentation.