Typorithm · hard

Self-Attention

self-attentiontransformersattentionqkvsoftmax

Saved on this device

Full lesson · 8 learning steps Open any exercise for focused practice

Step 1 of 8 · intro card

Learn Self-Attention

Attention & Transformers

Self-attention lets every token gather information from all the others. It projects each token into a query, key, and value, scores queries against keys with a scaled dot product, softmaxes those scores into weights, and returns a weighted sum of the values.

Self-attention is the heart of the transformer. Each input token's embedding is linearly projected into three vectors: a query (Q), a key (K), and a value (V).

To decide how much one token should attend to another, the query of the first is dotted with the key of the second. Doing this for all pairs gives a score matrix; dividing by sqrt(dK) keeps the scores well-scaled, and softmax turns each row into attention weights that sum to one.

The output for each token is the weighted sum of every token's value vector, using those attention weights. This is how each position pulls in context from the rest of the sequence in a single step.

Q, K, V are learned linear projections of the input.

scores = (Q · Kᵀ) / sqrt(dK), then softmax to weights.

output is a weighted sum of the value vectors.

// Q, K, V are linear projections of x
const scores = Q.matmul(K.transpose()).div(Math.sqrt(dK));
const weights = softmax(scores);
const output = weights.matmul(V);

Read the implementation

Connect each part of the code to the concept before you type.

1. Q, K, V projections

Each token's embedding is projected by learned weight matrices into a query, a key, and a value.

// Q, K, V are linear projections of x

2. Dot-product scores

Multiplying queries by transposed keys gives every query a similarity score against every key.

Q.matmul(K.transpose())

3. Scale by sqrt(dK)

Dividing by the square root of the key dimension keeps the scores from growing too large and saturating softmax.

.div(Math.sqrt(dK))

4. Softmax to weights

Softmax turns the raw scores into a probability distribution that sums to one across the keys.

const weights = softmax(scores);

5. Weighted sum of values

The output is a weighted sum of the value vectors, so each token gathers information from the others.

const output = weights.matmul(V);

A complete learning loop

Understand

Read the concept map and follow the code breakdown.

Reconstruct

Type the implementation, repair mistakes and work through every recall exercise.

Verify

Use the final check without hints; revisit weak concepts in short practice sessions.