Q, K, V Projections: Where the Asymmetry Comes From

Each token's input vector is the same when it serves as a query, a key, or a value. The asymmetry comes from three separate learned projection matrices that produce Q, K, V from the same input.

Q = X · W_Q, K = X · W_K, V = X · W_V

The three matrices are independently learned. They can — and during training do — diverge in what features they emphasize. W_Q learns to project tokens into a "what would I like to find?" subspace; W_K projects them into a "what am I when I'm being searched for?" subspace; W_V projects them into a "what do I contribute to other tokens' representations?" subspace.

Common shapes

Standard MHA: W_Q, W_K, W_V each have shape (d_model, d_model). After projection and reshape into heads, each head sees (seq_len, d_head) tensors where d_head = d_model / n_heads.
GQA / MQA: W_Q is still (d_model, d_model), but W_K and W_V are smaller — only enough rows for n_kv_heads instead of n_heads. (Track 4, lesson on GQA.)
RoPE: applied to Q and K after these projections, before the dot product. V is not rotated.

Code

Multi-head Q/K/V projection in PyTorch·python

class MultiHeadProjection(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
    def forward(self, x):
        B, L, _ = x.shape
        Q = self.W_q(x).view(B, L, self.n_heads, self.d_head)
        K = self.W_k(x).view(B, L, self.n_heads, self.d_head)
        V = self.W_v(x).view(B, L, self.n_heads, self.d_head)
        # Transpose to (B, n_heads, L, d_head) for attention
        return Q.transpose(1, 2), K.transpose(1, 2), V.transpose(1, 2)

Q, K, V Projections: Where the Asymmetry Comes From

Common shapes

Code

External links

Exercise

Progress

Comments 0