The paper Attention Is All You Need describes the Transformer architecture, which describes attention as a function of the queries $Q = x W^Q$, keys $K = x W^K$, and values $V = x W^V$:

$\text{Attention(Q, K, V)} = \text{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right) V \\ = \text{softmax}\left( \frac{x W^Q (W^K)^T x}{\sqrt{d_k}} \right) x W^V$

In the Transformer, there are 3 different flavors of attention:

  1. Self-attention in the Encoder, where the queries, keys, and values all come from the input to the Encoder.
  2. Encoder-Decoder attention in the Decoder, where the queries come from the input to the Decoder, and the keys and values come from the output of the Encoder
  3. Masked self-attention in the Decoder, where the queries, keys and values all come from the input to the Decoder, and, for each token, the $\text{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right)$ operation is masked out (zero'd out) for all tokens to the right of that token (to prevent look-ahead, which is cheating during training).

What is the gradient (i.e. the partial derivatives of the loss function w.r.t. $x$, $W^Q$, $W^K$, $W^V$, and any bias term(s)) of each of these attention units? I am having a difficult time wrapping my head around derivating a gradient equation because I'm not sure how the softmax function interacts with the partial derivatives, and also, for the Encoder-Decoder attention in the Decoder, I'm not clear how to incorporate the encoder output into the equation.

  • $\begingroup$ Let me try to understand what you're really looking for. You're not looking for all the details of the partial derivatives with respect to the parameters, i.e. you're not looking for someone to describe to you how back-prop works, right? If I am right, can you please explain what type of answer you're looking for? Anyway, I suspect that the original authors of the transformer did not have to compute the gradients by hand, but let the library perform automatic differentiation (i.e. back-prop) automatically, as it's usually the case. You just need to make sure your operations are differentiable. $\endgroup$
    – nbro
    Dec 8, 2020 at 0:57
  • 3
    $\begingroup$ I understand how backprop works in general for deep networks, networks with convolutions, and networks with residuals. I guess I am asking for what the mathematical equation of a gradient is for a specific layer of the neural network, where the layer is an attention unit. Specifically, deriving the partial derivatives of the loss w.r.t. the input $x$, weights $W^Q$, $W^K$, $W^V$, and the bias terms (if the attention unit has any -- I wasn't clear on that from the paper), assuming we have the partial derivatives of the later layers already. This may be solved by software, but I'm still curious! $\endgroup$ Dec 8, 2020 at 1:16
  • $\begingroup$ Have you got any solution for this. If so please let me know. I'm also trying the same $\endgroup$ Apr 17, 2022 at 14:18

1 Answer 1


I have written a blog to answer this question, please see https://say-hello2y.github.io/2022-09-07/attention-gradient

I use Matrix calculus to solve this question, here I put the final result of the gradient of an attention unit. $$ \frac{\partial f(X)}{\partial W_h^K }=\gamma K^T\mathbb{P}_h\frac{\partial f(X)}{\partial A }^TQW_h^Q $$ $$ \frac{\partial f(X)}{\partial W_h^Q }=\gamma Q^T\frac{\partial f(X)}{\partial A }\mathbb{P}_{h}^{T}KW_h^K $$ $$ \frac{\partial f(X)}{\partial W^v }=A^T\frac{\partial f(X)}{\partial X }(W^O)^T $$ $$ \frac{\partial f(X)}{\partial W^O }=\mathrm{Concat}(\mathrm{head_1}, ..., \mathrm{head_H})^T\frac{\partial f(X)}{\partial X } $$ For more detail, please see https://say-hello2y.github.io/2022-09-07/attention-gradient.If you have any questions, feel free to contact me ,my email is [email protected].


You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .