Are single-head and multi-head attention equivalent in terms of input and output?

Question

So I'm trying to write PyTorch code that performs the single-head computation, so you get some input vectors, attention mechanism/linear transformations, some output vectors.

Then I tried to re-use the same matrices by splitting them up into submatrices and get the same output vectors by using a 2-head attention approach. If I'm not mistaken, you should be able to just concatenate the outputs together.

The dimensions were correct, so output of single head and output of multi-head had the same dimensions, making me believe that multi-head is just better for performance reasons but not really different if we talk about input and output.

However, the output values were not the same despite both using the same matrix.

Is multi-head now inherently different from single head or how can I have an example compute the exact same input and output using single head and multi-head? I thought it should be possible if you re-use the same matrices but I was wrong...

Answer 1

Your approach of splitting the same matrices into submatrices for different heads and then concatenating together as output doesn't work as intended because each head in multi-head attention should have distinct learnable weight matrices which enable each head to focus on different aspects of the input sequence. Simply splitting the weight matrices into two submatrices for the two heads defeats the purpose of having multiple heads in the first place, essentially forcing the two heads to learn the same original aspect of self-attention and might causing the output values were not the same despite both using the same matrix.

Answer 2

Splitting weight matrices into weights for heads should be fine, but of course your output will not be the same as if you were using single head attention, since you are calculating attention scores and in the end values independently for each head. This means your backprop will take a different path for the different splits of the weight matrices.

As far as I remember the huggingface implementation of GPT-2 (and maybe others) actually uses splitting: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py#L123