Backpropagating Through Scaled Dot-Product Attention to Queries, Keys and Values

… more less With \( S = QK^T/\sqrt{d} \), \( P = \mathrm{softmax}(S) \), \( O = PV \): \( dV = P^T dO \), \( dP = dO\,V^T \), each row of \( dS \) is \( (\mathrm{diag}(p)-pp^T)\,dp \), then \( dQ = dS\,K/\sqrt{d} \).