🎓 A1Learner
🔍
0 of 53 lessons completed

PyTorch Tokens That Talk

Self-attention is the step where each position reads the earlier positions and mixes what they offer into its own vector. The mix is weighted, so a position can take a lot from one neighbour and almost nothing from another.

Three tokens, three vectors

A window of three characters becomes three vectors in a table. Each row is one position, and each column is one feature of that character.

import torch

tokens = torch.tensor([[1.0, 0.0, 2.0, 0.0],
                       [0.0, 1.0, 0.0, 2.0],
                       [1.0, 1.0, 1.0, 1.0]])
print(tokens.shape)
print(tokens[0].tolist())
print(tokens[2].tolist())

One row of weights per position

A weight is a plain number that says how much of one vector enters the mix. The rows below give the first position only itself, split the second evenly, and let the third lean on itself.

import torch

weights = torch.tensor([[1.0, 0.0, 0.0],
                        [0.5, 0.5, 0.0],
                        [0.25, 0.25, 0.5]])
print(weights.tolist())
print(weights.sum(dim=-1).tolist())

The mix is one matrix product

Multiplying the weight table by the token table mixes every position in a single step. Row two of the result is the average of the first two token vectors, exactly as its weights demand.

import torch

tokens = torch.tensor([[1.0, 0.0, 2.0, 0.0],
                       [0.0, 1.0, 0.0, 2.0],
                       [1.0, 1.0, 1.0, 1.0]])
weights = torch.tensor([[1.0, 0.0, 0.0],
                        [0.5, 0.5, 0.0],
                        [0.25, 0.25, 0.5]])
mixed = weights @ tokens
print(mixed.tolist())

Example

import torch

# Three token vectors, one per position in a tiny window.
tokens = torch.tensor([[1.0, 0.0, 2.0, 0.0],
                       [0.0, 1.0, 0.0, 2.0],
                       [1.0, 1.0, 1.0, 1.0]])

# Hand-set weights. Row i mixes the tokens that position i reads.
# Every entry above the diagonal is zero, so no position reads ahead.
weights = torch.tensor([[1.0, 0.0, 0.0],
                        [0.5, 0.5, 0.0],
                        [0.25, 0.25, 0.5]])

mixed = weights @ tokens

print("tokens:")
for row in tokens.tolist():
    print(row)
print("weights:")
for row in weights.tolist():
    print(row)
print("mixed:")
for row in mixed.tolist():
    print(row)
print("row sums:", weights.sum(dim=-1).tolist())

Key points

  • Token vector: the list of numbers that stands for one character at one position
  • Weighted mix: a sum where each vector counts in proportion to its weight
  • A weight of zero shuts a position out of the mix completely
  • Each position gets its own row of weights, so each position mixes differently
  • Weights above the diagonal stay zero, so no position reads the future

In this example, three hand-written token vectors stand for a three-character window. The weight table is lower-triangular, which means every entry above the diagonal is zero and no position can read ahead. The matrix product mixes all three positions at once. The first mixed row equals the first token, because its row gives it full weight. The second mixed row averages the first two tokens. The third blends all three, with half of its weight on itself. The row sums print as 1.0, so each mix is a true weighted average rather than a growing sum. Later lessons replace these hand-set weights with weights the model computes and learns.

Note

The weights here are set by hand, and the next lessons learn them from the text instead

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

In the closing example, why does the first mixed row equal the first token vector exactly?