Loading your learning…
Loading your learning…
A token needs context to predict what comes next. Learn how queries and keys select information from earlier tokens, then mix their values without looking into the future.
import torch def causal_attention(q, k, v): # Each input: [tokens, features], one sequence and one head. scores = q @ k.T / q.shape[-1] ** 0.5 future = torch.ones(q.shape[0], k.shape[0], device=q.device, dtype=torch.bool).triu(1) weights = scores.masked_fill(future, float("-inf")).softmax(dim=-1) return weights @ v, weights if __name__ == "__main__": q = torch.tensor([[1., 0.], [0., 1.], [1., 1.]]) k = q.clone() v = torch.tensor([[2., 0.], [0., 4.], [2., 2.]]) output, weights = causal_attention(q, k, v) torch.testing.assert_close(weights.sum(dim=-1), torch.ones(3)) assert torch.count_nonzero(weights.triu(1)) == 0 print(output) print(weights)A token is a piece of text. Give each token a query (what information it seeks), a key (what it matches), and a value (what it contributes). Here each input is a floating-point matrix [T, D]: T tokens with D features each. All three have the same shape and device. The example uses three tokens and two features. In a trained model, learned linear projections of token representations create these vectors; these hand-written vectors isolate the attention operation.
import torch def causal_attention(q, k, v): # Each input: [tokens, features], one sequence and one head. scores = q @ k.T / q.shape[-1] ** 0.5 future = torch.ones(q.shape[0], k.shape[0], device=q.device, dtype=torch.bool).triu(1)q @ k.T computes one dot product for every query–key pair, producing [T, T] scores. Row i describes the matches for token i. For query [1, 1] and key [0, 1], the dot product is 1. A dot product adds D products. When features are roughly unit-sized and unrelated, the spread of that sum grows like sqrt(D), so dividing by sqrt(D) keeps scores in a similar range at any D; without it, large scores would make softmax put almost all weight on a single key and learning would slow down.
def causal_attention(q, k, v): # Each input: [tokens, features], one sequence and one head. scores = q @ k.T / q.shape[-1] ** 0.5 future = torch.ones(q.shape[0], k.shape[0], device=q.device, dtype=torch.bool).triu(1)