Multi-Headed Self-Attention: Understanding the Core of Transformers

Multi-headed self-attention is the core modeling component of Transformers. In this question, we'll get some practice working with the self-attention equations, and motivate why multi-headed self-attention can be preferable to single-headed self-attention.

Recall that attention can be viewed as an operation on a query q ∈ R^d, a set of value vectors {v₁, . . . , v_n}, v_i ∈ R^d, and a set of key vectors {k₁, . . . , k_n}, k_i ∈ R^d, specified as follows:

c = ∑_i=1ⁿ v_iα_i (1)

α_i = exp(k_i^⊤q) / ∑_j=1ⁿ exp(k_j^⊤q) (2)

with α_i termed the 'attention weights'. Observe that the output c ∈ R^d is an average over the value vectors weighted with respect to α_i.

(a) (4 points) Copying in attention. One advantage of attention is that it's particularly easy to 'copy' a value vector to the output c. In this problem, we'll motivate why this is the case.

i. (1 point) Explain why α can be interpreted as a categorical probability distribution.

ii. (2 points) The distribution α is typically relatively 'diffuse'; the probability mass is spread out between many different α_i. However, this is not always the case. Describe (in one sentence) under what conditions the categorical distribution α puts almost all of its weight on some α_j, where j ∈ {1, . . . , n} (i.e. α_j ≫ ∑_i≠j α_i). What must be true about the query q and/or the keys {k₁, . . . , k_n}?

iii. (1 point) Under the conditions you gave in (ii), describe what properties the output c might have.

iv. (1 point) Explain (in two sentences or fewer) what your answer to (ii) and (iii) means intuitively.

i. α can be interpreted as a categorical probability distribution because it is a set of non-negative values that sum to 1, which is the definition of a probability distribution.

ii. The categorical distribution α puts almost all of its weight on some α_j when the query q is very similar to one of the keys k_j. This causes the dot product k_i^⊤q to be much larger for j than for any other i, leading to a much larger value of α_j.

iii. Under these conditions, the output c will be very similar to the value vector v_j corresponding to α_j, because it is weighted much more heavily than the other value vectors.

iv. Intuitively, this means that attention can easily 'copy' a value vector to the output if the query is very similar to one of the keys. This is useful because it allows the model to directly use relevant information from the input sequence without having to rely on sequential processing. However, it also means that attention can be misled if the query is not actually similar to any of the keys, which is why multi-headed attention can be useful.