I red the paper about transformers, and fully understood every single piece of that, so now I'm implementing it from scratch in tensroflow, without using any shipped layer from the library.

The only missing part is how they intend to take a single tensor (batch size, time-steps, embeddings) and give it to the multihead module.

For what I can tell, it seems that they take the embeddings of size $d_{model}$, then split the into $n$ heads, so each head get a piece of the embedding of size $d_{model}/n$, however I don't quite see what's the intuition why this should work

Am I missing something? are they just duplicating the input for each head instead?


1 Answer 1


Yes, the input is duplicated for each head. This allows the model to jointly attend to different information about the input at the same time. Passing only a subset of the embedding vector would likely result in a worse representation, since each head has less information about the input.

The confusion probabily arises from the fact that the output dimension of each Attention head in [1] is $d_{\text{model}}/h$, a fraction of the input embedding dimension $d_{\text{model}}$, where $h$ is the number of attention heads. However, the reduction of dimensionality is not obtained by selecting a subset of the input dimension $d_\text{model}$, but by performing the linear projections (i.e matrix multiplications) $QW_i^Q$, $K W_i^K, VW_i^V$ since $Q \in \mathbb{R}^{n \times d_{\text{model}}}$ and $W_i^Q, W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$.

This reduction is useful to reduce the computational cost of each head, which is given by $$\mathcal{O}(n^2 \cdot d_i)$$ This way, the total computational cost of the Multi-Head Attention is similar to that of a single-head attention with full dimensionality, but is shown in [1] that multi-head attention works better.

The output dimension is mantained, since values coming from each attention head are concatenated and then projected with a matrix $W^O \in \mathbb{R}^{h d_v \times d_\text{model}}$, obtaining a representation of the same shape as the input.

$$\begin{equation} \text{MultiHead}(Q,K,V) = \text{Concat}(h_1, \dots, h_i) W^O \\ h_i = \text{Attention}(Q W_i^Q, KW_i^K,VW_i^V) \end{equation}$$ Please note that with the choice $d_k = d_v = d_\text{model}/h$, $W^O$ is a square matrix, but with different choices of $d_k$ and $d_v$ it would project the attention output in the same space of the input.


[1] Vaswani, Ashish, et al. "Attention is all you need." Advances in neural information processing systems 30 (2017).


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.