Self-Attention in Python

In the last lesson every position took an equal share from itself and the positions before it. Self-attention lets each position choose how much to take from each one, based on the characters. Each position makes three vectors of 16 numbers with three linear layers, already set up for you: a query (what it is looking for), a key (what it contains) and a value (what it hands over).

The score of a position for an earlier one is the dot product of its query with that position's key: big when they match. t.transpose(i, j) gives back t with dimensions i and j swapped, so a @ b.transpose(0, 1) takes the dot product of every row of a with every row of b. Your tensors have the 32 windows first, so there you swap dimensions 1 and 2 instead.

a = torch.tensor([[1.0, 2.0], [0.0, 1.0]])   # 2 queries
b = torch.tensor([[3.0, 1.0], [1.0, 1.0]])   # 2 keys
print((a @ b.transpose(0, 1)).tolist())        # [[5.0, 3.0], [1.0, 1.0]]

Divide the scores by the square root of head_size (√16 = 4): large scores make softmax give almost all the weight to one position. Then mask them and softmax them as in the last lesson; dim=-1 means the last dimension. Finally weights @ v gives each position a weighted sum of the values it may see. The three layers are untrained, so the weights are random for now: training will make the matches useful.

When it's done, your script prints the shapes, the weights of the first window and four checks:

h: [32, 8, 32] q: [32, 8, 16] k: [32, 8, 16] v: [32, 8, 16]
scores: [32, 8, 8] weights: [32, 8, 8] out: [32, 8, 16]
weights in window 0, 'kelolo\nr'
position 0 [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
position 1 [0.155, 0.845, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
position 2 [0.442, 0.21, 0.348, 0.0, 0.0, 0.0, 0.0, 0.0]
position 3 [0.217, 0.158, 0.187, 0.438, 0.0, 0.0, 0.0, 0.0]
position 4 [0.24, 0.207, 0.151, 0.16, 0.242, 0.0, 0.0, 0.0]
position 5 [0.211, 0.134, 0.148, 0.152, 0.18, 0.175, 0.0, 0.0]
position 6 [0.1, 0.097, 0.264, 0.263, 0.068, 0.19, 0.018, 0.0]
position 7 [0.157, 0.202, 0.079, 0.094, 0.182, 0.141, 0.041, 0.103]
every row adds up to 1: True
nothing taken from later positions: True
position 0 only takes its own value: True
out is one attention head: True

Your task: below tril, store the queries query(h) in q and the keys key(h) in k; store every query times every key, divided by the square root of head_size, in scores; mask scores with tril; softmax them over the last dimension into weights; store the values value(h) in v; and multiply weights by v into out.

Assignment
Follow the Coding Tutorial to write one head of self-attention, from the queries and keys to the weighted values.

Hint
Check the shapes: q, k and v are [32, 8, 16], scores and weights are [32, 8, 8], out is [32, 8, 16]. The mask is the same masked_fill line as in the last lesson.