I have a difficult time understanding the "multi-head" notion in the original transformer paper. What makes the learning in each head unique? Why doesn't the neural network learn the same set of parameters for each attention head? Is it because we break query, key and value vectors into smaller dimensions and feed each portion to a different head?

The reason each head is different is because they each learn a different set of weight matrices $$\{ W_i^Q, W_i^K, W_i^V \}$$ where $$i$$ is the index of the head. To clarify, the input to each attention head is the same. For attention head $$i$$:

\begin{align} Q_i(x) &= x W_i^Q \\ K_i(x) &= x W_i^K \\ V_i(x) &= x W_i^V \\ \text{attention}_i(x) &= \text{softmax} \left(\frac{Q_i(x) K_i(x)^T}{\sqrt{d_k}} \right) V_i(x). \end{align}

Notice that the input to each head is $$x$$ (either the semantic + positional embedding of the decoder input for the first decoder layer, or the output of the previous decoder layer). More info

The question as to why gradient descent learns each set of weight matrices $$\{ W_i^Q, W_i^K, W_i^V \}$$ to be different across each attention head is very similar to "Is there anything that ensures that convolutional filters end up the same?", so maybe you might find the answer there helpful for you:

No, nothing really prevents the weights from being different. In practice though they end up almost always different because it makes the model more expressive (i.e. more powerful), so gradient descent learns to do that. If a model has n features, but 2 of them are the same, then the model effectively has n−1 features, which is a less expressive model than that of n features, and therefore usually has a larger loss function.

• I think this problem would be a bit different from convolution networks. In convolution layers, the input to each kernel is the same. I thin in attention, the inputs are also different. Is it the case that we break the inputs into smaller chunks in attention? Dec 12 '20 at 23:26
• The original transformer paper words it a bit confusing, but the input to each attention head is the same. For attention head $i$, $Q_i(x) = x W_i^Q$, $K_i(x) = x W_i^K$, and $V_i(x) = x W_i^V$. The output of attention head $i$ is $attention_i(x) = softmax \left(\frac{Q_i(x) K_i(x)^T}{\sqrt{d_k}} \right) V_i(x)$. The input to each head is $x$ (either the semantic + positional embedding of the decoder input for the first decoder layer, or the output of the previous decoder layer). For more info, read the section "Self-Attention in Detail": jalammar.github.io/illustrated-transformer Dec 12 '20 at 23:58
• I did more research into this and it seems that both ways exist in attention literature. We have "narrow self-attention" in which the original input is split into smaller chunks and each head get it's own small input. We also have "wide self-attention" in which the whole input gets fed into each head separately. Wide self-attention gives better results at the cost of memory and computation time. See peterbloem.nl/blog/transformers Dec 13 '20 at 0:55
• Interesting, thanks for the share! Are there more publications on wide vs narrow self-attention? I wanted to research further, but couldn't find any literature or academic papers on it. I know that the original transformer code in Tensor2Tensor, and also BERT's code, uses attention the way I described (which is considered "wide self-attention") Dec 13 '20 at 1:17
• No unfortunately that's all I could find. Dec 13 '20 at 2:10

Multiple attention heads in a single layer in a transformer is analogous to multiple kernels in a single layer in a CNN: they have the same architecture, and operate on the same feature-space, but since they are separate 'copies' with different sets of weights, they are hence 'free' to learn different functions.

In a CNN this may correspond to different definitions of visual features, and in a Transformer this may correspond to different definitions of relevance:1

For example:

Architecture Input (Layer 1)
(Layer 1)