Difference between VJP and JVP functions in JAX

[GoktugGuvercin]

Hello JAX Community;

I am confused about VJP and JVP functions introduced by JAX. While jax.jvp() for jacobian-vector product directly computes the multiplication of jacobian matrix evaluated at given input arguments by the tangents, and as a result produces the corresponding cotangent vector, why does jax.vjp() for vector-jacobian product prefer to return a function that performs that multiplication in reverse order ?

[YouJiacheng]

Because jax.vjp can compute the function output without extra cost, and you usually need the function output to determine the cotangent.
i.e. the input of jacobian-vector product is known in prior, while the input of vector-jacobian product is not.
Why jax.jvp directly computes the product? Because a function that can perform jacobian-vector product of f(g(·)) at point x need to recompute or save g(x).
jax.linearize return a jvp_fun that can perform jacobian-vector product, at the cost of saving intermediate result of the function(in the case that the function is a composite of many functions).
https://jax.readthedocs.io/en/latest/_autosummary/jax.linearize.html

[GoktugGuvercin]

I guess I can be still confused at "jax.vjp" can compute the output without extra cost".

Do you want to mean that since cotangent is known (it will be input argument to vjp) in vector jacobian product, function output is computable explicitly without extra cost, but since we try to compute cotangent in vjp, so function output needs to be recomputed.

[GoktugGuvercin]

As a result, both of them (vjp and jvp) take the function and its input arguments (primals) and it computes the output of that function over those arguments. How does vjp no need extra cost ?

[YouJiacheng]

jax.vjp can compute the f(x) without extra cost because compute vjp_fun of f at point x need a forward run of f

[mattjj]

Have you already checked out the Autodiff cookbook?

I think there are two parts to this question:

1. why does jax.vjp return a function that performs the multiplication in reverse order?
2. why the asymmetry between the signatures of jax.jvp and jax.vjp, where in particular the latter returns a function and the former does not?

For the first part: that's just the mathematical definition of what a VJP computes, along with how the chain rule works (i.e. the rule for how differentiation interacts with function composition). If you have a function f = g ∘ h, where ∘ means composition, i.e. f(x) = g(h(x)), then ∂ f(x) = ∂g(h(x)) ∘ ∂h(x), where the ∘ is composing linear maps or, if you like, multiplying matrices. So if we want to compute a vector-Jacobian product v ↦ v' ∂ f(x) then we can decompose it in terms of the constituent functions as ((v' ∂g(h(x))) ∂h(x)), which looks like evaluating a sequence of two simpler VJPs in the reverse order from how we would evaluate f(x) for some x.

As for the asymmetry between the jvp and vjp signatures, that is indeed curious, and maybe even unnatural from a mathematical point of view. It's motivated by computational rather than mathematical concerns, just as @YouJiacheng already pointed out: one can evaluate a JVP using less memory.

In a bit more detail, the "dual" to jax.vjp is really jax.linearize, which returns a function just like jax.vjp. Sometimes we write type signatures like this:

jvp      : (a -> b) -> (a, T a) -> (b, T b)
linearize: (a -> b) -> a -> (b, T a -o T b)
vjp      : (a -> b) -> a -> (b, CT b -o CT a)

Even if type signatures aren't your thing, you can see that the linearize and vjp ones are much more similar to each other than to that of jvp.

The relationship between jvp and linearize can be summarized like this: for any x and x_dot of the right shapes, we have

y, y_dot = jax.jvp(f, (x,), (x_dot,))

y2, f_lin = jax.linearize(f, x)
y_dot2 = f_lin(x_dot)

assert jnp.all(y == y2)  # or at least all close, up to floating point differences
assert jnp.all(y_dot == y_dot2)

In other words, we've just computed the same thing in two ways. But the jax.jvp call did it in less memory. See the autodiff cookbook for why, though there are some clues at the bottom of this post too. (On the other hand, if we have a bunch of different vectors like x_dot that we want to push through this computation, maybe it's cheaper in terms of FLOPs to produce f_lin and call that many times instead of re-running the jax.jvp every time, even though it'll cost more memory.)

The relationship with vjp is that for any y_bar of the right shape we further have

y3, f_vjp = jax.vjp(f, x)
x_bar, = f_vjp(y_bar)

assert jnp.all(y == y3)
assert jnp.vdot(x_bar, x_dot) == jnp.vdot(y_bar, y_dot)

In math notation, the last line is saying we've computed a quantity like u' ∂f(x) v in two ways, once like (u' ∂f(x)) v and the other time like u' (∂f(x) v).

Maybe some last clues are to look at how each of these would look for a primitive operation like sin, and also for generic function composition. Here's some pseudocode:

def jvp(sin)(x, xdot):
  y = sin(x)
  ydot = cos(x) * xdot
  return y, ydot

def lin(sin)(x):
  y = sin(x)
  cos_x = cos(x)
  sin_lin = lambda xdot: cos_x * xdot
  return y, sin_lin

def vjp(sin)(x):
  y = sin(x)
  cos_x = cos(x)
  sin_vjp = lambda ybar: ybar * cos_x
  return y, sin_vjp

From this you might be able to see jvp requires less memory. With jvp(sin) we computed cos(x) for use in a multiplication, but then we were able to throw it away before returning the result. Whereas for both lin(sin) and vjp(sin) the value of cos_x gets saved and returned as part of the result (in the closure of the returned function).

And for generic function composition:

def jvp(g ∘ h)(x, xdot):
  y, ydot = jvp(h)(x, xdot)
  z, zdot = jvp(g)(y, ydot)
  return z, zdot

def lin(g ∘ h)(x):
  y, h_lin = lin(h)(x)
  z, g_lin = lin(g)(y)
  return z, lambda xdot: g_lin(h_lin(xdot))

def vjp(g ∘ h)(x):
  y, h_vjp = vjp(h)(x)
  z, g_vjp = vjp(g)(y)
  return z, lambda ybar: h_vjp(g_vjp(ybar))

Hope these clues help!

[mattjj]

Yes, you nailed it. The "dot" and "bar" convention has existed for a long time in autodiff (like here).

I just use it as a naming clue, without necessarily having something precise in mind. But to try to unpack it, it'd go something like you said: if we have a variable named x, use the name x_dot to refer to its corresponding tangent variable in forward-mode (mathematically a tangent vector), and for reverse-mode use the name x_bar to refer to its corresponding cotangent value in reverse-mode (mathematically something like a representer vector for the cotangent vector). I say "a representer vector for" because x_bar usually just refers to an array value, usually of the same type as x_dot (assuming these objects correspond to finite-dimensional real vector spaces), rather than a scalar-valued function on array values.

There's some additional commentary in this "implicit layers" tutorial, starting under the heading "The two transformations of autodiff: JVPs and VJPs", particularly around the lines "JVPs answer questions like" and "VJPs answer questions like". I found being precise about the cotangent view to be helpful in providing a simple, mechanical derivation of how to differentiate through ODE solvers, without having to import a lot of extraneous mathematical machinery.

Hope that helps!

[mattjj]

Just to underscore something: x, x_dot, and x_bar typically all have the same type, particularly when we think of the type of x as corresponding to a real finite-dimensional vector space.

[AntixK]

Thank you very much!!

[frhack]

vjp : (a -> b) -> a -> (b, T b -o T a)

May be there is a typo, do you mean CT ?

vjp : (a -> b) -> a -> (b, CT b -o CT a)

[mattjj]

Yes, thanks! Edited above to fix it.