Jacobian-Vector Product

machine-learning calculus linear-algebra

Definition

Jacobian-Vector 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{u}\in {\mathbb{R}}^n$, the Jacobian-vector product is

$${\mathrm{JVP}}_f(\mathbf{x};\mathbf{u})=J_f(\mathbf{x})\mathbf{u}\in {\mathbb{R}}^m.$$

The vector $\mathbf{u}$ represents a direction in the input space. The JVP gives the corresponding first-order change in the output of $f$.

Dimensions

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

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

So a JVP takes a direction from the input space and pushes it forward to the output space.

Directional derivative

Let

$$\gamma (t)=\mathbf{x}+t\mathbf{u}.$$

Then, by the chain rule,

$$\frac{d}{dt}f(\gamma (t)){|}_{t=0}=J_f(\mathbf{x})\mathbf{u}.$$

So the JVP is the directional derivative of $f$ at $\mathbf{x}$ in direction $\mathbf{u}$.

Special case

If $f:{\mathbb{R}}^n\to \mathbb{R}$ is scalar-valued, then $J_f(\mathbf{x})$ is a $1\times n$ matrix. Therefore:

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

So for scalar outputs, a JVP is the directional derivative obtained by taking the gradient and dotting it with the direction vector.

Comparison

JVP versus VJP

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.

Implementations

Forward-mode evaluation

A JVP can be computed without forming the full Jacobian explicitly. The usual method is to evaluate the function once while carrying, for each intermediate value, both:

• its ordinary value, called the primal;
• its first-order change in the chosen direction, called the tangent.

The input is seeded with the direction vector, and constants are seeded with zero.

struct Dual {
    primal
    tangent
}

function jvp(f, x, u)
    // x is the input point, u is the input direction
    seed each input variable x_i as Dual(x_i, u_i)
    seed each constant c as Dual(c, 0)
    evaluate f using the propagation rules below
    return the tangent part of the final output

Primitive operations propagate tangents by the ordinary differentiation rules:

add((a, da), (b, db)) = (a + b, da + db)
sub((a, da), (b, db)) = (a - b, da - db)
mul((a, da), (b, db)) = (a * b, da * b + a * db)
div((a, da), (b, db)) = (a / b, (da * b - a * db) / (b * b))
sin((a, da))         = (sin(a), cos(a) * da)

Python sketch:

from __future__ import annotations
 
from dataclasses import dataclass
from math import cos, sin
from typing import Callable
 
 
@dataclass(frozen=True)
class Dual:
    primal: float
    tangent: float
 
 
def add(x: Dual, y: Dual) -> Dual:
    return Dual(x.primal + y.primal, x.tangent + y.tangent)
 
 
def sub(x: Dual, y: Dual) -> Dual:
    return Dual(x.primal - y.primal, x.tangent - y.tangent)
 
 
def mul(x: Dual, y: Dual) -> Dual:
    return Dual(
        x.primal * y.primal,
        x.tangent * y.primal + x.primal * y.tangent,
    )
 
 
def div(x: Dual, y: Dual) -> Dual:
    return Dual(
        x.primal / y.primal,
        (x.tangent * y.primal - x.primal * y.tangent) / (y.primal * y.primal),
    )
 
 
def sin_dual(x: Dual) -> Dual:
    return Dual(sin(x.primal), cos(x.primal) * x.tangent)
 
 
def jvp(f: Callable[[Dual], Dual], x: float, u: float) -> float:
    seeded_x = Dual(x, u)
    return f(seeded_x).tangent

For tensor-valued code, the same idea is used with matching shapes: the tangent associated with a tensor has the same shape as the tensor itself. The final tangent is the Jacobian-vector product.

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{u}=\left[\begin{array}{c}{u}_1\\ u_2\end{array}\right]$, the JVP is

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

This gives the first-order change of both output coordinates when the input is perturbed in direction $\mathbf{u}$.

If we choose

$$\mathbf{u}=\left[\begin{array}{c}1\\ 0\end{array}\right],$$

then

$$J_f(x_1,x_2)\mathbf{u}=\left[\begin{array}{c}{x}_2\\ 2x_1\end{array}\right],$$

which is exactly the rate of change of $f$ when only $x_1$ is varied.