AI Concepts in Python

Attention and Self-Attention in Python

Attention lets each token look at the other tokens and decide which ones matter. Every token has a query (what it looks for), a key (what it offers) and a value (the information it passes on).

The concept in plain words

Attention lets each token look at the other tokens and decide which ones matter. Every token has a query (what it looks for), a key (what it offers) and a value (the information it passes on).

Compute scaled dot-product attention by hand with queries, keys and values, then add a causal mask like GPT.

Key ideas

Attention weights come from a softmax, so each row adds up to 1.
Dividing by √d stops the scores from growing too large for long vectors.
A causal mask stops a token from looking at future tokens during generation.
Multi-head attention runs several attentions in parallel, and positional encoding adds word order.

When should you use it?

Understand what happens inside neural networks and AI libraries.
Learn the idea before you implement it in the matching challenges.
Prepare for machine learning courses, interviews and projects.

See it in action

Click a row of the heatmap to see where that token looks. Turn on the causal mask to hide future tokens, and change the sharpness to focus the attention.

Example code

Run code →

main.py

import math

def softmax(row):
    exps = [math.exp(v - max(row)) for v in row]
    return [e / sum(exps) for e in exps]

# Three tokens, each with a query, key and value vector of size 2.
Q = [[1, 0], [0, 1], [1, 1]]
K = [[1, 0], [0, 1], [1, 1]]
V = [[10, 0], [0, 10], [5, 5]]
d = len(K[0])

def attention(causal=False):
    out = []
    for i, q in enumerate(Q):
        # Score: how well this query matches every key, scaled by sqrt(d).
        scores = [sum(a * b for a, b in zip(q, k)) / math.sqrt(d) for k in K]
        if causal:  # a token may only look at itself and earlier tokens
            scores = [s if j <= i else -1e9 for j, s in enumerate(scores)]
        weights = softmax(scores)
        out.append([round(sum(w * v[c] for w, v in zip(weights, V)), 2) for c in range(2)])
    return out

print("attention:", attention())
print("causal attention:", attention(causal=True))

Expected output

attention: [[6.02, 3.98], [3.98, 6.02], [5.0, 5.0]]
causal attention: [[10.0, 0.0], [3.3, 6.7], [5.0, 5.0]]

How it works

For each query, the scores are its dot products with every key, divided by √d. Softmax turns the scores into weights, and the output is the weighted average of the values. With the causal mask the first token can see only itself, so its output is exactly its own value [10, 0].

Change the values and run the program in the CodeUtility online Python compiler without installing Python locally.

Practice exercises

Modify the runnable example with the exercises below to build understanding beyond copying the result.

Change the inputs or weights and predict the output before running.
Print the intermediate values to follow each step of the calculation.
Solve the linked challenges below to implement the idea yourself.

Run in Python IDE →

Practice with a challenge

Use what this lesson covers to solve a real problem. Each challenge runs in your browser with instant tests and an AI review.

All challenges