JVPs in math#

Mathematically, given a function\(f : \mathbb{R}^n \to \mathbb{R}^m\), the Jacobian of\(f\) evaluated at an input point\(x \in \mathbb{R}^n\), denoted\(\partial f(x)\), is often thought of as a matrix in\(\mathbb{R}^m \times \mathbb{R}^n\):

\(\qquad \partial f(x) \in \mathbb{R}^{m \times n}\).

But you can also think of\(\partial f(x)\) as a linear map, which maps the tangent space of the domain of\(f\) at the point\(x\) (which is just another copy of\(\mathbb{R}^n\)) to the tangent space of the codomain of\(f\) at the point\(f(x)\) (a copy of\(\mathbb{R}^m\)):

\(\qquad \partial f(x) : \mathbb{R}^n \to \mathbb{R}^m\).

This map is called thepushforward map of\(f\) at\(x\). The Jacobian matrix is just the matrix for this linear map on a standard basis.

If you don’t commit to one specific input point\(x\), then you can think of the function\(\partial f\) as first taking an input point and returning the Jacobian linear map at that input point:

\(\qquad \partial f : \mathbb{R}^n \to \mathbb{R}^n \to \mathbb{R}^m\).

In particular, you can uncurry things so that given input point\(x \in \mathbb{R}^n\) and a tangent vector\(v \in \mathbb{R}^n\), you get back an output tangent vector in\(\mathbb{R}^m\). We call that mapping, from\((x, v)\) pairs to output tangent vectors, theJacobian-vector product, and write it as:

\(\qquad (x, v) \mapsto \partial f(x) v\)

JVPs in JAX code#

Back in Python code, JAX’sjax.jvp() function models this transformation. Given a Python function that evaluates\(f\), JAX’sjax.jvp() is a way to get a Python function for evaluating\((x, v) \mapsto (f(x), \partial f(x) v)\).

importjaximportjax.numpyasjnpkey=jax.random.key(0)# Initialize random model coefficientskey,W_key,b_key=jax.random.split(key,3)W=jax.random.normal(W_key,(3,))b=jax.random.normal(b_key,())# Define a sigmoid function.defsigmoid(x):return0.5*(jnp.tanh(x/2)+1)# Outputs probability of a label being true.defpredict(W,b,inputs):returnsigmoid(jnp.dot(inputs,W)+b)# Build a toy dataset.inputs=jnp.array([[0.52,1.12,0.77],[0.88,-1.08,0.15],[0.52,0.06,-1.30],[0.74,-2.49,1.39]])# Isolate the function from the weight matrix to the predictionsf=lambdaW:predict(W,b,inputs)key,subkey=jax.random.split(key)v=jax.random.normal(subkey,W.shape)# Push forward the vector `v` along `f` evaluated at `W`y,u=jax.jvp(f,(W,),(v,))

In terms ofHaskell-like type signatures, you could write:

jvp::(a->b)->a->Ta->(b,Tb)

whereTa is used to denote the type of the tangent space fora.

In other words,jvp takes as arguments a function of typea->b, a value of typea, and a tangent vector value of typeTa. It gives back a pair consisting of a value of typeb and an output tangent vector of typeTb.

Thejvp-transformed function is evaluated much like the original function, but paired up with each primal value of typea it pushes along tangent values of typeTa. For each primitive numerical operation that the original function would have applied, thejvp-transformed function executes a “JVP rule” for that primitive that both evaluates the primitive on the primals and applies the primitive’s JVP at those primal values.

That evaluation strategy has some immediate implications about computational complexity. Since we evaluate JVPs as we go, we don’t need to store anything for later, and so the memory cost is independent of the depth of the computation. In addition, the FLOP cost of thejvp-transformed function is about 3x the cost of just evaluating the function (one unit of work for evaluating the original function, for examplesin(x); one unit for linearizing, likecos(x); and one unit for applying the linearized function to a vector, likecos_x*v). Put another way, for a fixed primal point\(x\), we can evaluate\(v \mapsto \partial f(x) \cdot v\) for about the same marginal cost as evaluating\(f\).

That memory complexity sounds pretty compelling! So why don’t we see forward-mode very often in machine learning?

To answer that, first think about how you could use a JVP to build a full Jacobian matrix. If we apply a JVP to a one-hot tangent vector, it reveals one column of the Jacobian matrix, corresponding to the nonzero entry we fed in. So we can build a full Jacobian one column at a time, and to get each column costs about the same as one function evaluation. That will be efficient for functions with “tall” Jacobians, but inefficient for “wide” Jacobians.

If you’re doing gradient-based optimization in machine learning, you probably want to minimize a loss function from parameters in\(\mathbb{R}^n\) to a scalar loss value in\(\mathbb{R}\). That means the Jacobian of this function is a very wide matrix:\(\partial f(x) \in \mathbb{R}^{1 \times n}\), which we often identify with the Gradient vector\(\nabla f(x) \in \mathbb{R}^n\). Building that matrix one column at a time, with each call taking a similar number of FLOPs to evaluate the original function, sure seems inefficient! In particular, for training neural networks, where\(f\) is a training loss function and\(n\) can be in the millions or billions, this approach just won’t scale.

To do better for functions like this, you just need to use reverse-mode.

VJPs in math#

Let’s again consider a function\(f : \mathbb{R}^n \to \mathbb{R}^m\).Starting from our notation for JVPs, the notation for VJPs is pretty simple:

\(\qquad (x, v) \mapsto v \partial f(x)\),

where\(v\) is an element of the cotangent space of\(f\) at\(x\) (isomorphic to another copy of\(\mathbb{R}^m\)). When being rigorous, we should think of\(v\) as a linear map\(v : \mathbb{R}^m \to \mathbb{R}\), and when we write\(v \partial f(x)\) we mean function composition\(v \circ \partial f(x)\), where the types work out because\(\partial f(x) : \mathbb{R}^n \to \mathbb{R}^m\). But in the common case we can identify\(v\) with a vector in\(\mathbb{R}^m\) and use the two almost interchangeably, just like we might sometimes flip between “column vectors” and “row vectors” without much comment.

With that identification, we can alternatively think of the linear part of a VJP as the transpose (or adjoint conjugate) of the linear part of a JVP:

\(\qquad (x, v) \mapsto \partial f(x)^\mathsf{T} v\).

For a given point\(x\), we can write the signature as

\(\qquad \partial f(x)^\mathsf{T} : \mathbb{R}^m \to \mathbb{R}^n\).

The corresponding map on cotangent spaces is often called thepullbackof\(f\) at\(x\). The key for our purposes is that it goes from something that looks like the output of\(f\) to something that looks like the input of\(f\), just like we might expect from a transposed linear function.

VJPs in JAX code#

Switching from math back to Python, the JAX functionvjp can take a Python function for evaluating\(f\) and give us back a Python function for evaluating the VJP\((x, v) \mapsto (f(x), v^\mathsf{T} \partial f(x))\).

fromjaximportvjp# Isolate the function from the weight matrix to the predictionsf=lambdaW:predict(W,b,inputs)y,vjp_fun=vjp(f,W)key,subkey=jax.random.split(key)u=jax.random.normal(subkey,y.shape)# Pull back the covector `u` along `f` evaluated at `W`v=vjp_fun(u)

In terms ofHaskell-like type signatures, we could write

vjp::(a->b)->a->(b,CTb->CTa)

where we useCTa to denote the type for the cotangent space fora. In words,vjp takes as arguments a function of typea->b and a point of typea, and gives back a pair consisting of a value of typeb and a linear map of typeCTb->CTa.

This is great because it lets us build Jacobian matrices one row at a time, and the FLOP cost for evaluating\((x, v) \mapsto (f(x), v^\mathsf{T} \partial f(x))\) is only about three times the cost of evaluating\(f\). In particular, if we want the gradient of a function\(f : \mathbb{R}^n \to \mathbb{R}\), we can do it in just one call. That’s howjax.grad() is efficient for gradient-based optimization, even for objectives like neural network training loss functions on millions or billions of parameters.

There’s a cost, though the FLOPs are friendly, memory scales with the depth of the computation. Also, the implementation is traditionally more complex than that of forward-mode, though JAX has some tricks up its sleeve (that’s a story for a future notebook!).

For more on how reverse-mode works, check outthis tutorial video from the Deep Learning Summer School in 2017.