Vector-Jacobian Product

machine-learning calculus linear-algebra

Definition

Vector-Jacobian Product

Let $f:{\mathbb{R}}^n\to {\mathbb{R}}^m$ be a differentiable function, let $\mathbf{x}\in {\mathbb{R}}^n$, and let $J_f(\mathbf{x})\in {\mathbb{R}}^{m\times n}$ be the Jacobian matrix of $f$ at $\mathbf{x}$.

For a vector $\mathbf{v}\in {\mathbb{R}}^m$, the vector-Jacobian product is

$${\mathrm{VJP}}_f(\mathbf{x};\mathbf{v})=J_f(\mathbf{x}{)}^T\mathbf{v}\in {\mathbb{R}}^n.$$

In row-vector notation, the same object is written as ${\mathbf{v}}^T{J}_f(\mathbf{x})$.

The vector $\mathbf{v}$ represents sensitivities at the output of $f$. The VJP combines them into the corresponding sensitivities at the input of $f$.

Dimensions

If $f:{\mathbb{R}}^n\to {\mathbb{R}}^m$, then:

$$J_f(\mathbf{x})\in {\mathbb{R}}^{m\times n},\mathbf{v}\in {\mathbb{R}}^m,J_f(\mathbf{x}{)}^T\mathbf{v}\in {\mathbb{R}}^n.$$

So a VJP takes a vector from the output space and maps it back to the input space.

Chain Rule

Let

$$\mathbf{y}=f(\mathbf{x}),L=L(\mathbf{y}),$$

where $L:{\mathbb{R}}^m\to \mathbb{R}$ is a scalar function. Then the chain rule gives:

$${\nabla }_{\mathbf{x}}L=J_f(\mathbf{x}{)}^T{\nabla }_{\mathbf{y}}L.$$

This is exactly a vector-Jacobian product with seed vector $\mathbf{v}={\nabla }_{\mathbf{y}}L$.

Backward pass

In backpropagation, one usually does not form the full Jacobian explicitly. Instead, one repeatedly applies VJPs, moving a gradient from the output of each layer back to its input.

Special case

If $f:{\mathbb{R}}^n\to \mathbb{R}$ is scalar-valued, then $J_f(\mathbf{x})$ is a $1\times n$ matrix. With seed $v=1$, the VJP becomes:

$$J_f(\mathbf{x}{)}^T\cdot 1=\nabla f(\mathbf{x}).$$

So for scalar outputs, a VJP recovers the usual gradient.

Notation

Convention

The name “vector-Jacobian product” comes from row-vector notation ${\mathbf{v}}^{T}J$. Many machine-learning texts and libraries use column gradients, so the same computation appears as $J^T\mathbf{v}$. The underlying operation is the same.

Comparison

VJP versus JVP

A Jacobian-vector product uses

$$J_f(\mathbf{x})\mathbf{u},$$

so it pushes an input direction $\mathbf{u}$ forward through $f$.

A vector-Jacobian product uses

$$J_f(\mathbf{x}{)}^T\mathbf{v},$$

so it pulls output sensitivities $\mathbf{v}$ backward to the input side.

Forward-mode automatic differentiation is based on JVPs, while reverse-mode automatic differentiation is based on VJPs.

Examples

Two-output function

Let

$$f(x_1,x_2)=\left[\begin{array}{c}{x}_1{x}_2\\ x_1^2+\sin x_2\end{array}\right].$$

Then the Jacobian is

$$J_f(x_1,x_2)=\left[\begin{array}{cc}{x}_2& x_1\\ 2x_1& \cos x_2\end{array}\right].$$

For $\mathbf{v}=\left[\begin{array}{c}{v}_1\\ v_2\end{array}\right]$, the VJP is

$$J_f(x_1,x_2{)}^T\mathbf{v}=\left[\begin{array}{cc}{x}_2& 2x_1\\ x_1& \cos x_2\end{array}\right]\left[\begin{array}{c}{v}_1\\ v_2\end{array}\right]=\left[\begin{array}{c}{x}_2{v}_1+2x_1{v}_2\\ x_1{v}_1+(\cos x_2)v_2\end{array}\right].$$

Each input coordinate receives contributions from every output coordinate, weighted by $\mathbf{v}$.

If we now define a scalar loss

$$L(\mathbf{y})=3y_1-y_2,$$

then

$${\nabla }_{\mathbf{y}}L=\left[\begin{array}{c}3\\ -1\end{array}\right],$$

and therefore

$${\nabla }_{\mathbf{x}}L=J_f(\mathbf{x}{)}^T\left[\begin{array}{c}3\\ -1\end{array}\right]=\left[\begin{array}{c}3x_2-2x_1\\ 3x_1-\cos x_2\end{array}\right].$$

This is the reverse-mode step carried out by a VJP.