Attention Walkthrough

In this post we will walk through a very simple example of multiheaded attention is computed. The goal is to keep the dimensions small so that it’s easier to understand what’s going on.

Step 0: Define some constants

We will use the three tokens ['major', 'league', 'baseball'] with model dimension 4 (d_model=4) and two attention heads (num_heads=2).

num_heads = 2
d_model = 4
d_k = d_v = d_model // num_heads

Here we also define the dimensions of the key and value here so that after after concatenation in the multihead attention algorithm, the final result will have dimension (seq_len, d_model).

Step 1: Get token embeddings

Usually we would tokenize the input string “major league baseball” and then look up each corresponding embedding in the model’s token embedding table. Here, we will just use random numbers to illustrate the math.

major = array([0.4, 0.2, 0.1, 0.3])
league = array([0.5, 0.3, 0. , 0.9])
baseball = array([0.8, 0.4, 0.2, 0.6])

x = np.stack([major, league, baseball])

Step 2: Define weight matrices

Here we define four weight matrices:

• The query, key, and value weight matrices have dimensions (n_heads, d_model, d_k). It is pretty standard for d_k == d_v, although it is technically not mathematically required to be the case.
• The output projection matrix is typically a square matrix with dimension (d_model, d_model).

wq = array([
    [
        [0.5, 0.6],
        [0.8, 0.3],
        [0.7, 0.2],
        [0.1, 0.8]
    ],
    [
        [0.3, 0.6],
        [0.5, 0.8],
        [0.9, 0.1],
        [0.5, 0.8]
    ]
])

wk = array([
    [
        [0. , 0.7],
        [0.4, 0.3],
        [0.8, 0.3],
        [0.7, 0.4]
    ],
    [
        [0.4, 0.9],
        [0.1, 0.7],
        [0.6, 0.6],
        [0.1, 0.6]
    ]
])

wv = array([
    [
        [0.4, 0.5],
        [0.7, 0.6],
        [0.9, 0.8],
        [0. , 0.3]
    ],
    [
        [0.9, 0.2],
        [0.2, 0.8],
        [0.7, 0. ],
        [0.1, 0.4]
    ]
])

wo = array([
    [0.2, 0.2, 0.6, 0.1],
    [0.3, 0.2, 0.2, 0.8],
    [0.2, 0.5, 0.1, 0.3],
    [0.6, 0.9, 0.5, 0.7]
])

Step 3: The multi-head attention algorithm

Now that we have our pieces ready, we can compute multihead attention on the input. The following is an example to illustrate the mechanism of multihead attention. Real implementations are vectorized to avoid slow loops.

# the input x has shape (seq_len, d_model)
def mha(x):
    heads = []
    for h in range(num_heads):
        q = x@wq[h]
        k = x@wk[h]
        v = x@wv[h]
        head = sdpa(q, k, v) # head has shape (seq_len, d_v)
        heads.append(head)
    concat = np.concatenate(heads, axis=1)
    return concat @ wo

As you can see, for each head we:

• compute the query matrix by matmul of input x with the query weights for that head. Notice that in this operation the input x has shape (seq_len, d_model) and the query weights for that head will have shape (d_model, d_k), so the resulting matrix will have shape (seq_len, d_k).
• compute the key and value matrices in the same way
• compute the final value of the head after scaled dot product attention between that head’s query, key, and value matrices. We will talk about SDPA in the next step
• The final value of the head will have shape (seq_len, d_v). If we think of these dimensions as rows and columns, the “head” is really operating on a group of columns to make a “skinny” matrix.

Finally, we concatenate all heads over the last dimension. This combines all “skinny” head matrices into a final matrix with shape (seq_len, d_model). We then matmul that combined matrix by the out projection matrix.

Step 4: Scaled dot product attention

Scaled dot product attention here is the same as in single-head attention (i.e. self-attention). The difference here is that we are sending “skinnier” matrices into the SPDA function.

def sdpa(q, k, v):
    raw_attn_scores = q@k.T
    norm_attn_scores = raw_attn_scores / d_k ** 0.5
    weights = np.stack([
        softmax(row)
        for row in norm_attn_scores
    ])
    return weights @ v

First, we matmul the query by the transpose of the key. Both the query and key matrices have dimension (seq_len, d_v), so raw_attn_scores will have dimension (seq_len, seq_len).

Next, we divide raw_attn_scores by the square root of d_k. This helps to stabilize training since the variance here scales with d_k, specifically in the softmax in the next step.

After that, we compute the softmax values for each row. Since each corresponds to a token in the input, this can be thought of as emphasizing which tokens in the sequence are most relevant to each query token. This operation doesn’t change the input shape.

Finally, we matmul the weights with the value matrix, (seq_len, seq_len) x (seq_len, d_v) resulting in a final matrix of (seq_len, d_v).