# How can the Jacobian of the softmax function be used for calculating the gradient of a linear layer?

I am creating a visual transformer from scratch as an exercise to have a better understanding of machine learning. Within the network it multiplies 2 matrices, then the softmax function is applied, and it is not the last layer, so I am assuming cross entropy would not be used, and multiplied by a constant I am going to refer to as $$C$$ and then more layers are applied afterward. Still, I am only planning on focusing on what I covered.

So take the matrices $$Q$$ and $$K$$ where $$Q$$ is a $$3\times2$$ matrix and $$K$$ is a $$2\times3$$ matrix, and let $$A$$ be the result of $$Q$$ and $$K$$ being multiplied which results in a $$3\times3$$ matrix.

Where, $$Q = \begin{pmatrix}q_{1,1} & q_{1,2}\\ q_{2,1} & q_{2,2}\\q_{3,1} & q_{3,2}\end{pmatrix}$$

and

$$K = \begin{pmatrix}k_{1,1} & k_{1,2} & k_{1,3}\\ k_{2,1} & k_{2,2} & k_{2,3}\end{pmatrix}$$

making $$A = \begin{pmatrix}q_{1,1}*k_{1,1} + q_{1,2}*k_{2,1} & q_{1,1}*k_{1,2} + q_{1,2}*k_{2,2} & q_{1,1}*k_{1,3} + q_{1,2}*k_{2,3}\\ q_{2,1}*k_{1,1} + q_{2,2}*k_{2,1} & q_{2,1}*k_{1,2} + q_{2,2}*k_{2,2} & q_{2,1}*k_{1,3} + q_{2,2}*k_{2,3}\\ q_{3,1}*k_{1,1} + q_{3,2}*k_{2,1} & q_{3,1}*k_{1,2} + q_{3,2}*k_{2,2} & q_{3,1}*k_{1,3} + q_{3,2}*k_{2,3}\end{pmatrix} = \begin{pmatrix} a_{1,1} & a_{1,2} & a_{1,3} \\ a_{2,1} & a_{2,2} & a_{2,3} \\ a_{3,1} & a_{3,2} & a_{3,3}\end{pmatrix}$$

Then the softmax function is applied on A which is defined as $$S(x_{i}) = \frac{e^{x_{i}}}{\sum_{j=0}^n e^{x_{j}}}$$

Afterwards $$B$$ = $$S(A)*C$$

So then on the backward pass, with respect to the loss, $$L$$, Let $$Z$$ = $$B*C$$, $$\frac{\partial L}{\partial A} = S'(Z)$$

where the derivative of the softmax function is given by

$$z_{i}(1-z_{i})$$ when $$i = j$$ and $$-z_{i}z_{j}$$ when $$i != j$$ from the $$Z$$ matrix shown above making $$\frac{\partial L}{\partial A} = \begin{pmatrix} z_{1,1}(1-z_{1,1}) & -z_{1,1}z_{1,2} & \cdots & -z_{1,1}z_{3,3} \\ -z_{1,2}z_{1,1} & z_{1,2}(1-z_{1,2}) & \cdots & -z_{1,2}z_{3,3} \\ \vdots & \vdots & \ddots & \vdots \\ -z_{3,3}z_{1,1} & -z_{3,3}z_{1,2} & \cdots & z_{3,3}(1-z_{3,3})\end{pmatrix}$$

Which is a $$9x9$$ matrix

What I don't understand is how I can calculate $$\frac{\partial A}{\partial Q}$$ and $$\frac{\partial A}{\partial K}$$ to find $$\frac{\partial L}{\partial K}$$ and $$\frac{\partial L}{\partial Q}$$ in order to allow for continuous back propagation. From what I understand I have to compute the jacobian of $$A$$ with respect to $$Q$$ to find $$\frac{\partial A}{\partial Q}$$ But I am not entirely sure how to do that to then revert back to the same dimensions as $$Q$$

You can just calculate the derivatives directly from your formula for $$A$$. I can write your formula as $$a_{ij} = \sum_{\alpha} q_{i\alpha} k_{\alpha j}$$ so then we have $$\frac{da_{ij}}{dq_{rs}} = \delta_{ir} k_{sj}$$ where $$\delta$$ is Kronecker delta.

In general, if you have a function from $$\mathbb{R}^n \to \mathbb{R}^m$$, then the derivative will be a $$(n+m)$$-dimensional tensor. In case you aren't used to the word "tensor", just think a vector is a 1D block of numbers, a matrix is a 2D block of numbers, and then this derivative will be a $$(n+m)$$-D block of numbers. That should make sense because as we see in the formula above, to fully understand $$\frac{dA}{dQ}$$ we need to understand the derivative of any entry in $$A$$ with respect to any entry in $$Q$$.

The next logical question would be - how do we use that for back propagation? In this case your real goal is to compute $$\frac{dL}{dQ}$$ and you know $$L$$ as a function $$L(A(Q,K))$$. You've already computed $$\frac{dL}{dA}$$, which was is a 2D tensor, which matches my formula above because the function that computes $$L$$ from $$A$$ has 2D input and 0D output.

The formula you want will be: \begin{align} \frac{dL}{dQ_{rs}} &= \sum_{ij} \frac{dL}{dA_{ij}} \frac{dA_{ij}}{dQ_{rs}} \\ &= \sum_{ij} \frac{dL}{dA_{ij}} \delta_{ir} k_{sj} \\ &= \sum_{j} \frac{dL}{dA_{ij}} k_{sj} \end{align} where $$\frac{dL}{dA_{ij}}$$ is the derivative you already computed in your question statement.

Of course, you can repeat the same steps for $$\frac{dL}{dK}$$.

## Example that might help the intuition?

Say we have a function $$f:\mathbb{R} \to \mathbb{R}$$ given by $$f(x) = g(h(x))$$ where $$h(x) = (x^2, x^3)$$ and $$g(a, b) = a \ln(b)$$. Then what's $$\frac{df}{dx}$$?

You probably learned to handle this kind of situation in a multivariable calculus class at some point. We can compute \begin{align} \frac{df}{dx} &= \frac{df}{dh_1} \frac{dh_1}{dx} + \frac{df}{dh_2}{dh_2}{dx} \\ &= \ln(h_2(x)) (2x) + \frac{h1(x)}{h2(x)} (3x^2) \\ &= 2x \ln(x^3) + 3x^2 \frac{x^2}{x^3} \\ &= 6x \ln(x) + 3x \end{align} which you could confirm if you want by manually differentiating $$f(x) = x^2 \ln(x^3)$$.

Can you see how this computation is just a really simple version of the same method I was using for your real problem above? In this case the mappings just go $$\mathbb{R} \to \mathbb{R}^2 \to \mathbb{R}$$ but you still apply the same idea where $$\frac{df}{dx} = \sum_{j} \frac{df}{dh_j} \frac{dh_j}{dx}$$.

I can get a bunch of good reading material on this topic by Googling chain rule for matrices using tensors. Here are some example hits that look reasonable to me at first glance: 1 and 2 and 3. But you don't need to trust these particular sources; really I think you can learn more yourself now that I've told you the right keywords to look up. Good luck :)

## EDIT: Computing $$\mathbf{\frac{dS(A)}{dQ}}$$

Per discussion in the comments, OP would like me to also write out how to compute $$\frac{dS(A)}{dQ}$$.

We want the derivative of the map $$Q \mapsto S(A)$$. We'll compute it via chain rule, using the composition $$Q \mapsto A \mapsto S(A)$$. We already computed the derivative of the first step, $$\frac{dQ}{dA}$$. For the second step $$\frac{dS(A)}{dA}$$ we are differentiating a $$3 \times 3$$ matrix with respect to a $$3 \times 3$$ matrix, so we expect a 4 dimensional output tensor which will be $$3 \times 3 \times 3 \times 3$$.

Using the derivative formula you already computed, we have \begin{align} \frac{dS(A)_{ij}}{dA_{rs}} &= \delta_{ir} \delta_{js} a_{ij} - a_{ij} a_{rs}. \end{align} Now we use Chain Rule to compute the final derivative we're looking for: \begin{align} \frac{dS(A)_{ij}}{dQ_{cd}} &= \sum_{r} \sum_{s} \frac{dS(A)_{ij}}{dA_{rs}} \frac{dA_{rs}}{dQ_{cd}} \\ &= \sum_{r,s} \left(\delta_{ir} \delta_{js} a_{ij} - a_{ij} a_{rs} \right) \left( \delta_{rc} k_{ds} \right) \\ &= \delta_{ic} k_{dj} - \sum_s a_{ij} a_{cs} k_{ds} \end{align}

Note this is a 4D tensor with dimensions $$3 \times 3 \times 3 \times 2$$, since the indices $$i,j,c$$ run from 1 to 3 and $$d$$ runs from 1 to 2.

• There is still one thing I do not entirely understand. If I were to apply my above example to this solution. I am assuming the final summation of j would go from 1 to 9 because dL/dA is a 9x9 matrix. However, K is a 2x3 matrix. This would make it impossible to get anything above 6 elements from K. How would I handle this? Commented Dec 21, 2022 at 22:01
• @SamMoldenha Oh yeah I didn't look at your $\frac{dL}{dA}$ part much because it wasn't what you were asking about, but it's not set up correctly. $A$ is $3 \times 3$ and $L$ is $1 \times 1$, so $\frac{dL}{dA}$ will be only $3 \times 3$. The entries of $\frac{dL}{A}$ will be like $\frac{dL}{da_{ij}}$ with $1 \le i,j \le 3$. If you write it out in coordinates, I'm not sure how you're even filling up a $9 \times 9$ matrix. Commented Dec 21, 2022 at 22:14
• Right now your description of what happens after $A$ is not very careful. You've defined $Z = BC = S(A) C^2$ which is weird, and then you don't explain how you're getting from $Z$ to $L$. (I assume $L$ will eventually be just a scalar value since you want it to be the loss in your model?) Skipping steps in this part was totally reasonable because that isn't what you were focusing the question on, but at this point maybe you could edit to improve it and then re-tag me and I could take a look. Commented Dec 21, 2022 at 22:18
• So my question more pertained to the softmax function, and the jacobian of a matrix multiplication product. In my example I probably should have left out C, but L is just the derivative of the softmax function which I defined in my question. Commented Dec 21, 2022 at 22:59
• Thank you so much! Commented Dec 22, 2022 at 18:15