Transformer Self-Attention Computed by Hand

Work a single self-attention head through query-key scores, softmax weights, and value mixing on three short tokens with tiny vectors.

Setup

Self-attention lets each token gather information from every other token, weighted by learned relevance. We take three tokens with 2D embeddings and identity projection matrices so the arithmetic stays transparent: queries, keys, and values equal the inputs.

Inputs

Let x1=(1,0), x2=(0,1), x3=(1,1). With Q=K=V=x, the raw attention score between token i and token j is the dot product q_i . k_j, scaled by 1/sqrt(d) with d=2.

import numpy as np
X=np.array([[1.,0.],[0.,1.],[1.,1.]])
d=X.shape[1]
S=X@X.T/np.sqrt(d)                 # scaled scores
S=S-S.max(1,keepdims=True)
W=np.exp(S); W/=W.sum(1,keepdims=True)  # softmax rows
out=W@X
print(np.round(W,3))
print(np.round(out,3))

Reading the weights

Token 3, whose embedding (1,1) aligns with both others, spreads its attention broadly, while token 1 attends most to itself and to token 3 because (1,0).(1,1)=1 beats (1,0).(0,1)=0. Each output row is a convex combination of the value vectors, so no single token can dominate unless its score is much larger.

• The 1/sqrt(d) scaling keeps scores from growing with dimension, which would otherwise push softmax into a near one-hot regime and kill gradients.
• Softmax rows sum to one, so attention redistributes information rather than amplifying it.
• Real heads use distinct learned Q, K, V matrices; multiple heads run in parallel and their outputs are concatenated.