How do I go from embeddings to queries, keys and values in the Transformer model?

Question

I am trying to implement Attention Is All You Need paper from scratch in PyTorch. So far, I implemented the Scaled Dot-Product Attention layer and the Multi-Head Attention layer. As I began to write the code for the Encoder, I am facing a question I have not yet found an answer to: How do I go from embeddings to queries, keys and values in the Transformer?

As you can see from Figure 1 in the paper below, embeddings enter the Encoder and then somehow they turn into queries, keys and values which enter the Multi-Head Attention layer. I do not know how to get the queries, keys and values from the embeddings. The way I implemented it, my Multi-Head Attention layer expects queries, keys and values in its forward method (the method for the forward pass in PyTorch). Maybe it should expect something else? I am a bit confused and I'm hoping someone can clarify what happens here. I looked at various resources online (various Cross Validated answers and Medium articles (such as Illustrated Attention)), but couldn't find a clear answer to this question.

Answer 1

The queries, keys and values are computed from the embeddings $x$ using matrix multiplication: \begin{align} q &= xW_Q \\ k &= xW_K \\ v &= xW_V, \end{align} where $W_Q$, $W_K$ and $W_V$ are trainable parameters of the attention layer. Basically the attention layer learns how to map the embeddings to queries, keys and values.

Probably the only trainable parameter currently in your attention layer is the matrix $W_O$, that is used for combining the outputs from the multiple heads (described in 3.2.2 of the paper). You could modify your attention layer to contain three additional trainable parameters - the matrices $W_Q, W_K, W_V$.