Skip to Content

Run several independent self-attention heads in parallel; concatenate and project. Each head can specialise.

Data Flow

Q, K, V from input
           |
     +-----+-----+-----+-----+
     |h_1  |h_2  |h_3  |h_4  |    split into h smaller Q, K, V
     +--+--+--+--+--+--+--+--+
        |     |     |     |
        v     v     v     v
     self_attn(self_attn(self_attn(self_attn(
        each with its own W_Q, W_K, W_V
        |
        v
     concat heads   ->   W_O   ->   output

Why multiple heads?

Code

def multi_head(Q, K, V, n_heads):
    B, L, d = Q.shape
    d_head = d // n_heads
    Qh = Q.view(B, L, n_heads, d_head).transpose(1, 2)
    Kh = K.view(B, L, n_heads, d_head).transpose(1, 2)
    Vh = V.view(B, L, n_heads, d_head).transpose(1, 2)
    out = self_attention(Qh, Kh, Vh)
    return out.transpose(1, 2).reshape(B, L, d)

Pitfalls

Analogy

A panel of detectives, each tracking one type of clue (motive, alibi, fingerprint), pooling their notes to solve the case.

Interview tip: “Why not just one big attention matrix?” — representational diversity at the same compute cost.

Advertisement