(40 points) Backpropagation of matrix multiplication. The key operation in transformer is to calculate
the query-key multiplication matrix.
Forward input: Q E Rnxd, K E Rnxd Forward output: Z = QKT, where Z E Rnxn
Backward input:
aJ
a.J
OZin
:
oz
aJ
aJ
aJ
OZn1
0Zn2
aZnn
aJ
aJ
and
for i,j =
Find the backward outputs
n.
,
aQij
SKij

Question

Accepted Answer

⚹ To solve the backpropagation problem for matrix multiplication in the context of transformers, we need to compute the gradients ∂J/∂Q and ∂J/∂K based on the forward output Z = QK^T. ⚹ ⚹ The backward outputs can be derived using the chain rule of calculus. ⚹ ⚹ The gradients can be computed as follows: ⚹ 1. ⚹ The forward output is given by $ Z = QK^T $. ⚹ 2. ⚹ The Jacobian matrix for the backward input is $ \frac{\partial J}{\partial Z} $. ⚹ 3. ⚹ To find $ \frac{\partial J}{\partial Q} $, we use the relation: $$ \frac{\partial J}{\partial Q} = \frac{\partial J}{\partial Z} \cdot K $$ where $ K $ is the matrix of keys. ⚹ 4. ⚹ To find $ \frac{\partial J}{\partial K} $, we use the relation: $$ \frac{\partial J}{\partial K} = Q^T \cdot \frac{\partial J}{\partial Z} $$ where $ Q^T $ is the transpose of the query matrix. ⚹ ⚹ Therefore, the backward outputs are given by the above equations for $ i, j $ ranging from 1 to $ n $. ⚹

Sia