Vector Jacobian product in automatic differentiation

Question

my questions is related to this post Higher Order of Vectorization in Backpropagation in Neural Network @shimao

I don't really get the following claim (I know how the chain rule works and what is the essence of how reverse mode autodifferentiation ): "In reality, it's not necessary to compute the jacobian in order to perform backpropagation. All that is needed is the "vector jacobian product", or VJP."

As far as I know, you need to compute the Jacobian at each step in order to perform the gradient accumulation. For example suppose we are in the step of computing the derivative of the following operation (very typical in neural networks):

$$ t = Wz, \,\,\, z\in \mathbb{R}^{m\times 1}, t \in \mathbb{R}^{n \times 1}, W\in\mathbb{R}^{n \times m} $$

Then we now that the vector derivative (the Jacobian) is given by:

$$ \frac{\partial t}{\partial z} = W $$

which is then plug into the chain rule (assume $\frac{\partial y}{\partial o}$ is the already computed gradient), to compute the final derivative given by the chain rule:

$$ \frac{\partial y}{\partial o} \frac{\partial t}{\partial z} $$

Hence, I dont really understand that sentence because that sentence says that we dont need to compute the Jacobian to perform the back propagation as we only need a Jacobian vector product; but clearly for this step we do need to obtein the Jacobian of the transformation between $t$ and $z$. Can someone provide a bit of intuition of what is exactly refering to?. In general I have seen that reverse-mode autodiff can be efficiently implemented through vector Jacobian products (and that doesn't require to compute the Jacobian), but when I think about it I realize that we do need to compute the Jacobian's at each step

Answer 1

I myself is new to this topic so I can't be sure. But I will tell you my intuition.

Regardless whether it is forward or backward mode AD, to evaluate the gradient is vector calculus that can be expressed as matrix-vector multiplications.

If you have an expression for the gradient: g=ABCDe, where A,B,C,D are matrices and e is an vector, in the computer program it is common that you have a way to obtain D*e directly with no explicit storage of D in memory.

Such example operation of D is like the negation of the vector, permutation of the vector elements, vector cross-product or convolutions. The core operation can be represented by much much less variables than the full expression of D as a matrix.

From the product of D*e, then you can apply C,B,A consecutively (without explicit expression) to obtain g.

This concept appears often in linear algebra. When you only require the product of a matrix and a vector, but not the matrix explicitly, it is way easier to implement the computation in program and have much lower memory requirement.

It just happens that this concept applies to AD calculation.