# Reverse-Mode Automatic Differentiation with respect to a Matrix: How to "Matrix Multiply" 4D Tensors?

This is a follow up question I have on this excellent answer: https://stats.stackexchange.com/a/235758/307400. I will save me writing down any details about reverse-mode automatic differentiation, the given answer gives a nice introduction. I also found this resource very helpful: https://marksaroufim.medium.com/automatic-differentiation-step-by-step-24240f97a6e6

I am implementing Automatic Differentiation of functions on matrices: $$f: \mathbb{R}^{k\times l} \rightarrow \mathbb{R}^{n\times m}$$. As mentioned in the answer linked above, if I calculate the partial derivative of such a function $$f: \mathbb{R}^{k\times l} \rightarrow \mathbb{R}^{n\times m}$$ with respect to another function $$g: \mathbb{R}^{r\times s} \rightarrow \mathbb{R}^{p\times q}$$, as we do in reverse-mode differentiation, I get a 4D array. The array has the shape $$\frac{\partial f}{\partial g}: n\times m \times p \times q$$. So far so good.

The formula for computing the adjoint $$\overline{v_i}$$ given the adjoints of $$v_i$$'s parents in the computation graph is $$\overline{v_i} = \sum_{j\in parents(i)} \overline{v_j} \frac{\partial v_j}{\partial v_i}$$ The shapes are:

• $$\overline{v_p}: p\times q \times n \times m$$
• $$\frac{\partial v_p}{\partial v_i}: n\times m \times k \times l$$

So here we have to perform "matrix multiplication" of two 4D arrays. How does this work?

I have come up with an approach to multiply these 4D arrays in the given context that makes some intuitive sense to me, but I have no idea if it is theoretically sound:
First thing I noticed is that we want to multiply over both dimensions $$n$$ and $$m$$: We multiply an array of shape $$p\times q \times n \times m$$ with another one of shape $$n\times m \times k \times l$$ and want to obtain an array of shape $$p\times q \times k \times l$$, as we want to obtain another adjoint, that is a partial derivative of the final "output" function that I left implicit here but has range $$\mathbb{R}^{p\times q}$$ with respect to $$v_i$$ with range $$\mathbb{R}^{k\times l}$$. We want to "multiply out" the dimensions $$n\times m$$.

What I did then was to consider each two dimensions together as one, implying that matrices are collapsed into vectors, e.g. by iterating them in row major order. This simplifies the shapes in the multiplication to $$\underbrace{p\times q}_{a} \times \underbrace{n\times m}_{b} \;\;\cdot\;\; \underbrace{n\times m}_{b} \times \underbrace{k\times l}_{c}$$ Now I just did ordinary matrix multiplication with these new indices (using $$x\in \{1, ..., a\}$$; $$y \in \{1, ..., b\}$$ and $$z \in \{1, ..., c\}$$): $$\left(\overline{v_p} \frac{\partial v_p}{\partial v_i} \right)_{x,z} = \sum_{y=1}^{b} \bigg(\;\overline{v_p}\;\bigg)_{x,y} \bigg(\frac{\partial v_p}{\partial v_i}\bigg)_{y,z}$$

This sums over the two dimensions we want to get rid of, but is this also theoretically sound?

I know this a bit of an older question, but I came across it while trying to understand the same thing, so let me try to explain how I came to think of this, which is not with tensors first, but maps and operators.

I can't quite follow what you're suggesting in terms of tensor products, but in general I think that the calculation of the backwards pass will depend on the details of what your function is and may requires some custom derivation. You may be doing this correctly, but I'm not sure that a tensor product or matrix multiplication is always the most straightforward interpretation/implementation.

The answers you link to give a very practical point of view on how to perform reverse-mode AD, but there's a more abstract way of looking at it as well, in terms of vector-Jacobian products. I think the JAX documentation has a good explanation of this. Let me explain that on a vector-valued function before getting to matrix-valued functions

### Vector-valued function

Let's say we have a function $$f: \mathbb{R}^n \rightarrow \mathbb{R}^m$$. The Jacobian (or tangent linear map) of this function evaluated at some "primal" point $$x_0 \in \mathbb{R}^n$$ is a linear map from the input space to the output space, so $$\mathcal{J}: \mathbb{R}^n \rightarrow \mathbb{R}^m$$. Since this is a linear operator, we could represent it numerically using a matrix $$J \in \mathbb{R}^{m \times n}$$.

Reverse-mode AD is really an algorithm for computing the "vector-transpose Jacobian product (VJP)" $$\mathbf{v}^T J$$ for some vector $$\mathbf{v} \in \mathbb{R}^m$$ (I think this is what you're calling the "adjoints of the parents"). If $$f$$ consists of a computational graph over other functions, this happens by propagating the "seed" vector $$\mathbf{v}$$ backwards from the final nodes. In any case, all we have to do to implement a new function is the VJP for that function, so we can ignore anything else that's happening upstream and downstream, assuming that $$\mathbf{v}$$ will be handed to us. For instance, if we wanted to construct the full Jacobian, we could repeat the backwards pass $$m$$ times using unit basis vectors, each time computing one row of $$J$$.

To be a little more precise, while $$\mathbf{v}$$ has the dimensions of an element of the output space, it is actually an element of the cotangent space, and what we evaluate is the pullback or adjoint map $$J^T \mathbf{v}$$. As an easy example, consider $$f(x) = A \mathbf{x}$$, with $$A \in \mathbb{R}^{m \times n}$$. Clearly $$A$$ is the Jacobian of this function, so the pullback $$\mathcal{J}: \mathbb{R}^m \rightarrow \mathbb{R}^n$$ is evaluated numerically by $$A^T \mathbf{v}$$.

### Matrix-matrix multiplication

For scalar- and vector-valued functions this probably seems irrelevant and abstract, but I think it's helpful for more general cases. For instance, let's look at matrix-matrix multiplication $$f(X) = A X$$, with $$A \in \mathbb{R}^{m \times n}$$ and $$X \in \mathbb{R}^{n \times \ell}$$, so that $$f: \mathbb{R}^{n \times \ell} \rightarrow \mathbb{R}^{m \times \ell}$$. The tangent map is another map $$\mathcal{J}: \mathbb{R}^{n \times \ell} \rightarrow \mathbb{R}^{m \times \ell}$$, which can once again be represented numerically by the matrix $$A$$. The pullback applied to $$V \in \mathbb{R}^{m \times \ell}$$ is just $$A^T V$$. You could view this as a specific 4D tensor product, but I think it's much easier to look at it as a map.

What if we instead want to differentiate $$g(A) = A X$$? We still have $$A \in \mathbb{R}^{m \times n}$$ and $$X \in \mathbb{R}^{n \times \ell}$$, but now $$g: \mathbb{R}^{m \times n} \rightarrow \mathbb{R}^{m \times \ell}$$. This is still a linear function in $$A$$, but now we evaluate the tangent map by right-multiplying by the matrix $$X$$:

$$\mathcal{J}(\cdot) = (\cdot) X$$

The adjoint map is $$\mathcal{J}^\dagger: \mathbb{R}^{m \times \ell} \rightarrow \mathbb{R}^{m \times n}$$, which can be evaluated by

$$\mathcal{J}^\dagger(\cdot) = (\cdot) X^T.$$

In other words, during the backward pass, when you are given the seed $$V \in \mathbb{R}^{m \times \ell}$$ you compute $$V X^T$$.

To make a general matrix-matrix product $$F(A, X) = AX$$ differentiable with respect to both arguments, you evaluate the adjoint map with respect to each input (basically partial differentiation). Now $$F: \mathbb{R}^{m \times n} \times \mathbb{R}^{n \times \ell} \rightarrow \mathbb{R}^{m \times \ell}$$, so the tangent maps with respect to each input are $$\mathcal{J}_A: \mathbb{R}^{m \times n} \rightarrow \mathbb{R}^{m \times \ell}$$ and $$\mathcal{J}_X = \mathbb{R}^{n \times \ell} \rightarrow \mathbb{R}^{m \times \ell}$$, as before. You can derive the operators in a principled way by computing the total derivative:

$$dF = (dA) X + A (dX)$$

This is linear with respect to each input, and the tangent and adjoint maps can be computed as for the single-argument case.

### Nonlinear matrix-valued function

So far all the example I gave are linear, so the tangent and adjoint maps are fairly obvious. Let me give a simple nonlinear example as well.

Let's take $$f(X) = X^T X$$, with $$X \in \mathbb{R}^{m \times n}$$ and $$f: \mathbb{R}^{m \times n} \rightarrow \mathbb{R}^{n \times n}$$. The total derivative of $$f$$ is

$$df = (dX)^T X + X^T (dX),$$

so the tangent map $$\mathcal{J}: \mathbb{R}^{m \times n} \rightarrow \mathbb{R}^{n \times n}$$ can be evaluated with

$$\mathcal{J}(\cdot) = (\cdot)^T X + X^T (\cdot).$$

To compute the adjoint we have to do a little more work and use the definition of the adjoint: for a linear operator $$\mathcal{A}: \mathcal{W} \rightarrow \mathcal{V}$$, the adjoint $$\mathcal{A}^\dagger: \mathcal{V} \rightarrow \mathcal{W}$$ is defined using the inner product $$\langle \cdot, \cdot \rangle$$:

$$\langle \mathcal{A} w, v \rangle = \langle w, \mathcal{A}^\dagger v \rangle, \qquad w \in \mathcal{W}, v \in \mathcal{V}.$$

For matrices, the standard inner product is the trace: $$\langle A, B \rangle = \mathrm{tr}\{ A B^T \}$$. In our case, if we choose a arbitrary elements $$W$$ of the tangent space $$\mathbb{R}^{m \times n}$$ and $$V$$ of the cotangent space $$\mathbb{R}^{n \times n}$$, then

$$\langle \mathcal{J}(W), V \rangle = \mathrm{tr} \left\{ (W^T X + X^T W) V^T \right\}.$$

Using some properties of the trace we can show that

$$\mathrm{tr} \left\{ (W^T X + X^T W) V^T \right\} = \mathrm{tr} \left\{ W (V + V^T) X^T \right\} = \langle W, X (V + V^T) \rangle,$$

from which it follows that the pullback is $$\mathcal{J}^\dagger V = X(V + V^T)$$.

### The general recipe

1. Define what spaces your function maps over, and what the tangent and linear maps will be.
2. Compute the total derivative of the function and extract the tangent map.
3. Use the definition of the adjoint and an appropriate inner product to derive the adjoint map.
4. Implement the pullback as an evaluation of the adjoint map on the "seed".

For a great example of how to derive reverse-mode AD rules for a complicated function I'd recommend checking out the OptNet paper, where they show how to backpropagate through a quadratic program.