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.

Transformer model architecture

  • 2
    $\begingroup$ The Annotated Transformer explains exactly how to implement a transformer from scratch in pytorch. nlp.seas.harvard.edu/annotated-transformer $\endgroup$
    – Sycorax
    Commented Mar 27, 2023 at 18:24
  • $\begingroup$ @Sycorax thank you, but I don't want to look at any code unless absolutely necessary. I'd like it if someone could explain it to me without code, so I can try my hand at implementing it. $\endgroup$ Commented Mar 27, 2023 at 19:24

1 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$.


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.