Tell about multi-head attention in detail: what are Q/K/V, how are attention weights calculated, why are there multiple heads?

Multi-head attention is a key mechanism in transformers that allows the model to focus on different parts of the input sequence simultaneously.

• Q (Query), K (Key), V (Value) are three matrices obtained from the input data through trainable linear transformations. Each position in the sequence is transformed into three vectors:

  • Query — the request for relevant information.
  • Key — the key used to compare with the query.
  • Value — the value returned in response to the query.

• Calculating attention weights:

  1. Compute the dot product of Query and Key for each pair of positions.
  2. Scale the result (usually by dividing by the square root of the key dimension) to stabilize gradients.
  3. Apply softmax to obtain a probability distribution of weights.
  4. Multiply the weights by the corresponding Values, and sum the results — this is the attention output.

• Why multiple heads? Multiple heads allow the model to look at information from different "angles," meaning each head learns to extract different dependencies and patterns in the data. This increases the expressive power and quality of the model.

Simplified example:

# For each head:
Q_i = X * W_Q_i
K_i = X * W_K_i
V_i = X * W_V_i

attention_weights = softmax((Q_i @ K_i.T) / sqrt(d_k))
output_i = attention_weights @ V_i

# Concatenate outputs of all heads
output = concat(output_1, ..., output_h) * W_O

Where W_Q_i, W_K_i, W_V_i, W_O are trainable matrices, d_k is the dimension of the keys, and h is the number of heads.