Is there any reference about backpropagation of the Transformer's multi-head layer?

Understanding Backpropagation in the Transformer’s Multi‑Head Attention Layer

The Transformer architecture has become the backbone of modern natural language processing, computer vision, and many other AI domains. While the forward pass of the multi‑head attention (MHA) mechanism is well documented, developers and researchers often wonder where to find a clear, step‑by‑step description of its backpropagation (gradient) computation. This post gathers the most relevant references and explains the gradient flow through the MHA layer.

Why Backpropagation Matters for Multi‑Head Attention

• Parameter Updates: The query, key, value projection matrices (W_Q, W_K, W_V) and the final output projection (W_O) are learned via gradient descent.
• Stability: Understanding gradients helps diagnose exploding/vanishing issues that are common in deep Transformer stacks.
• Custom Extensions: Researchers adding sparsity, low‑rank factorization, or alternative attention kernels need the exact gradient formulas.

Core Equations of the Forward Pass

For a single head, the forward computation is:

Q = XW_Q,   K = XW_K,   V = XW_V
Attention(Q, K, V) = softmax(QKᵀ / √d_k) V

With h heads, the inputs are split into h sub‑spaces, each head computes its own attention, and the results are concatenated and projected:

ConcatHeads = [head₁; …; head_h]
Output = ConcatHeads W_O

Gradient Flow Through a Single Head

Let L be the loss. The key gradients are:

1. Gradient w.r.t. the output projection:

   dL/dW_O = ConcatHeadsᵀ · dL/dOutput

2. Gradient w.r.t. the concatenated heads:

   dL/dConcatHeads = dL/dOutput · W_Oᵀ

3. Split the gradient for each head: dL/dHead_i is the slice of dL/dConcatHeads belonging to head i.
4. Gradient through the attention matrix: For head i,

   Let A = softmax(QKᵀ / √d_k)
   dL/dV = Aᵀ · dL/dHead_i
   dL/dA = dL/dHead_i · Vᵀ
   

   The derivative of the softmax is:

   dL/dS = (A ⊙ dL/dA) - A ⊙ (sum_rows(A ⊙ dL/dA))
   where S = QKᵀ / √d_k
   

5. Gradients w.r.t. Q and K:

   dL/dQ = dL/dS · K / √d_k
   dL/dK = (dL/dS)ᵀ · Q / √d_k
   

6. Gradients w.r.t. projection matrices:

   dL/dW_Q = Xᵀ · dL/dQ
   dL/dW_K = Xᵀ · dL/dK
   dL/dW_V = Xᵀ · dL/dV
   

All heads share the same pattern; the only difference is that each head operates on a different slice of the hidden dimension.

Key References

• Vaswani et al., “Attention Is All You Need” (2017) – Original Transformer paper; includes the forward equations and a brief description of back‑propagation in the appendix.
• Alammar, “The Illustrated Transformer” (2018) – Visual guide that walks through the gradient flow for a single attention head.
• Wang et al., “Understanding the Difficulty of Training Transformers” (2020) – Provides a detailed derivation of gradients for multi‑head attention and discusses scaling factors.
• Karpathy’s minGPT implementation – Open‑source code where the backward pass is handled automatically by PyTorch, but the comments outline the same gradient steps described above.
• 🤗 Transformers library – The MultiHeadAttention module in modeling_utils.py contains the exact tensor operations that PyTorch autograd differentiates.

Practical Tips for Implementing Custom Backpropagation

1. Use Automatic Differentiation: In PyTorch, TensorFlow, or JAX you rarely need to code the gradients manually; just define the forward pass as shown and let the framework handle the rest.
2. Check Gradient Shapes: Print .shape of dL/dW_Q, dL/dW_K, dL/dW_V, and dL/dW_O to ensure they match the expected (d_model, d_k) or (h·d_v, d_model) dimensions.
3. Stabilize Softmax Gradients: Apply torch.nn.functional.softmax with dim=-1 and consider torch.nn.functional.log_softmax + nn.NLLLoss for numerical stability.
4. Gradient Clipping: Transformers often benefit from clipping (e.g., torch.nn.utils.clip_grad_norm_) to prevent exploding gradients across many layers.
5. Inspect Attention Maps: Visualizing A (the softmax matrix) during training can reveal whether gradients are flowing through all tokens or getting stuck on a few.

Summary

The backpropagation through a Transformer's multi‑head attention layer follows a clear, linear algebraic path: loss → output projection → concatenated heads → each head’s attention matrix → softmax → queries, keys, values → projection matrices. The derivations above, together with the cited papers and open‑source implementations, give you a complete reference for both theoretical understanding and practical coding.

Whether you are debugging a custom attention variant or simply want to deepen your grasp of how Transformers learn, these resources and formulas should serve as a solid foundation.