Is there any Reference About Backpropagation of the Transformer’s Multi‑Head Layer?

The Transformer architecture has become the backbone of modern natural‑language processing, computer‑vision, and multimodal models. While the forward pass of the multi‑head attention (MHA) layer is well documented, many practitioners wonder where to find a clear, step‑by‑step description of its backpropagation (gradient) computation. This post gathers the most useful references, explains the mathematics behind the gradients, and offers practical tips for implementing or debugging the backward pass.

1. Quick Recap: What Is Multi‑Head Attention?

Given an input matrix X ∈ ℝ^{N×d_model} (where N is the sequence length and d_model the hidden dimension), the MHA layer performs:

1. Linear projections to obtain queries, keys, and values for each head:
   Q_h = XW_Q^{(h)}, K_h = XW_K^{(h)}, V_h = XW_V^{(h)}
2. Scaled dot‑product attention for each head:
   Attention_h = softmax(Q_h K_h^T / √d_k) V_h
3. Concatenation of all heads and a final linear projection:
   Output = concat(Attention_1,…,Attention_H) W_O

All matrices W_Q^{(h)}, W_K^{(h)}, W_V^{(h)}, W_O are learnable parameters. The backward pass must propagate gradients through the softmax, the matrix multiplications, and the concatenation.

2. Where to Find Formal Derivations?

Below is a curated list of papers, tutorials, and code repositories that explicitly discuss the gradient flow through MHA.

• Original Transformer paper – “Attention Is All You Need” (Vaswani et al., 2017). Section 3.2 describes the forward equations; the supplementary material provides a brief note on backpropagation through the softmax.
• “The Illustrated Transformer” – Jay Alammar’s blog. The post includes an interactive diagram of the backward pass for the attention scores.
• “A Structured Self‑Attentive Sentence Embedding” – Lin et al., 2017. Appendix A derives gradients for the attention weighting, which directly applies to each head.
• “Understanding the Gradient Flow in Transformers” – Zhang et al., 2021. This paper analyses how gradients propagate through the softmax and the scaling factor √d_k, offering closed‑form expressions.
• PyTorch source code – torch.nn.MultiheadAttention. The autograd implementation is a reliable reference; reading the backward function in the C++/CUDA kernels shows the exact order of operations.
• TensorFlow/Keras implementation – tf.keras.layers.MultiHeadAttention. The _compute_attention and _apply_attention methods are annotated with gradient comments.

3. Deriving the Gradients – Step by Step

3.1. Gradient of the Softmax

For a single head, let S = QK^T / √d_k and A = softmax(S). The Jacobian of the softmax for element i,j is:

∂A_{ij} / ∂S_{ik} = A_{ij}(δ_{jk} - A_{ik})

During backpropagation, the upstream gradient ∂L/∂A is multiplied by this Jacobian to obtain ∂L/∂S. In practice, deep‑learning libraries fuse this operation into a single efficient kernel.

3.2. Gradient w.r.t. Queries, Keys, Values

Given V and the upstream gradient G = ∂L/∂Output (after the final linear projection), the chain rule yields:

∂L/∂V = A^T G
∂L/∂A = G V^T
∂L/∂S = ∂L/∂A ⊙ (A ⊙ (1 - A))   // element‑wise product with softmax Jacobian
∂L/∂Q = (∂L/∂S) K / √d_k
∂L/∂K = (∂L/∂S)^T Q / √d_k

Here “⊙” denotes element‑wise multiplication. The scaling factor 1/√d_k appears in both ∂L/∂Q and ∂L/∂K.

3.3. Gradient w.r.t. Projection Matrices

Each head has its own projection matrices. For the query matrix W_Q^{(h)}:

∂L/∂W_Q^{(h)} = X^T (∂L/∂Q_h)

Analogous formulas hold for W_K^{(h)} and W_V^{(h)}. The final output projection W_O receives the concatenated attention outputs Concat(Attention_h):

∂L/∂W_O = Concat(Attention_1,…,Attention_H)^T G

4. Practical Tips for Implementers

• Use built‑in layers. Both PyTorch’s nn.MultiheadAttention and TensorFlow’s MultiHeadAttention already implement the correct backward pass. Re‑implementing from scratch is only advisable for research experiments.
• Watch out for masking. When applying causal or padding masks, the mask must be applied *before* the softmax. During backpropagation the mask simply zeroes the corresponding gradients.
• Numerical stability. The softmax is usually computed as softmax(S - max(S, axis=-1, keepdims=True)). The same subtraction must be mirrored in the gradient computation (handled automatically by autograd).
• Gradient clipping. Transformers often require clipping (e.g., torch.nn.utils.clip_grad_norm_) because the attention gradients can become large, especially with long sequences.
• Mixed‑precision training. When using FP16, ensure that the scaling factor √d_k is cast to the same precision to avoid overflow/underflow in the QK^T product.

5. Summary

The backpropagation through a Transformer's multi‑head attention layer follows a well‑defined chain of matrix‑multiplication, softmax, and linear‑projection gradients. The key references listed above—especially the original “Attention Is All You Need” paper, the gradient‑focused analysis by Zhang et al. (2021), and the source code of major deep‑learning libraries—provide both the theoretical derivation and practical implementation details.

By consulting these resources and keeping the practical tips in mind, you can confidently debug, extend, or re‑implement the MHA backward pass in any AI project.