Custom vjp vs jvp for newton's method

[mrtupek]

As one step in my objective function I am solving a nonlinear system with Newton's method. I'm trying to be able to compute both gradients and hessians of the resulting objective. Following the jax online examples for custom vjp/jvp (one of which has a sign error in it, by the way), I can get both vjp and jvp to work on a simple example (simpler than the training example, I'd argue). I have one main question:

Why can't we use forward mode differentation to a custom vjp?

Here is the vjp:

from jax import grad, jacfwd, hessian, jacrev, custom_vjp, custom_jvp
from jax.lax import while_loop
from jax.config import config
config.update("jax_enable_x64", True)

f = lambda x : x*x - 5
df = grad(f)

@custom_vjp
def newton(x0):
    x = while_loop(lambda xp: f(xp) > 1e-8,
                   lambda xp: xp - f(xp) / df(xp),
                   x0)
    return x

def newton_f(x0):
    x = newton(x0)
    return x, x 

def newton_b(x, g):
    return (g/df(x),)

newton.defvjp(newton_f, newton_b)

x0 = 6.
print(np.sqrt(5.), newton(x0))
print(grad(np.sqrt)(5.), grad(newton)(x0))
print(hessian(np.sqrt)(5.), jacrev(grad(newton))(x0))
#print(hessian(np.sqrt)(5.), hessian(newton)(x0))

If I uncomment this last line I get:

TypeError: can't apply forward-mode autodiff (jvp) to a custom_vjp function.

I have also implemented a jvp for this example (I can include if helpful). This works fine, but my understanding is that in reverse mode it will have to recompute the newton iterations on the backward pass. I'd ideally want a way to save info from the forward pass to use in the reverse pass (so as not to redo these calculations and only apply the linearized 'adjoint'), but also have that function work in forward differentiation so I can call hessian or use hessian-vector products. Is there a way to setup something like this?

[mattjj]

Thanks for the question!

The reason is that a custom_vjp defines how a function should behave only under reverse-mode autodiff, using a custom rule. It leaves undefined how a function should behave under forward-mode autodiff.

If you define a custom_jvp then you can use both forward- and reverse-mode autodiff because even though your rule only specifies how the function should behave under forward-mode autodiff, JAX will attempt to automatically transpose the forward-mode rule (see the pdf on this page for a description of how automatic transposition relates to autodiff).

We could in principle attempt automatic transposition of a custom_vjp rule. We haven't yet only because that'd be a bit weird: defining a custom_vjp directly is in effect defining both a custom_jvp and a transpose rule at once (as well as a partial evaluation rule). So if you want both modes it makes most sense to define a custom_jvp.

I think with your example you can define a custom_jvp and get everything you want. Reverse mode won't need to unroll the iterations or anything like that. Depending on whether you want to use a direct or implicit solver for the tangent system, it might be helpful to use @shoyer's lax.custom_linear_solve. (If you post your custom_jvp then we can check it together!)

Alternatively, if you want to define only a custom_vjp, you can still compute Hessians or HVPs, just using reverse-over-reverse rather than forward-over-reverse. Something like this:

def hessian(f):
  return jax.jacrev(jax.jacrev(f))

def hvp(f, x, v):
    return grad(lambda x: jnp.vdot(grad(f)(x), v))(x)

The autodiff cookbook has some discussion about these options.

What do you think?

[mrtupek]

Thanks a lot for the quick reply! The jvp is here:

import jax.numpy as np
from jax import grad, jacfwd, hessian, jacrev, custom_vjp, custom_jvp
from jax.lax import while_loop
from jax.config import config
config.update("jax_enable_x64", True)

val = 6.

f = lambda x : x*x - val
df = grad(f)

@custom_jvp
def newton(x0):
    x = while_loop(lambda xp: abs(f(xp)) > 1e-8, # bug fixed here from previous version, need abs for residual!
                   lambda xp: xp - f(xp) / df(xp),
                   x0)
    return x

@newton.defjvp
def newton_jvp(xt, vt):
    x, = xt
    v, = vt
    primal_out = newton(x)
    tangent_out = v / df(primal_out)
    return primal_out, tangent_out

x0 = 3.
print(np.sqrt(val), newton(x0))
print(grad(np.sqrt)(val), grad(newton)(x0))
print(hessian(np.sqrt)(val), hessian(newton)(x0))

It works in forward and reverse. However, I'm worried about the performance in the reverse pass of having to redo the nonlinear newton iterations. It is potentially much cheaper to save information (like the vjp lets you do) to only solve a linear system on the reverse pass. This is very likely a case of premature optimization on my part. But this is typically how people do adjoint-based sensitives for dynamic PDE-constrained inverse problems, so it was how I thought to do it first.

Would it make sense to have an interface where one can specify both the vjp and jvp?

[shoyer]

I have been intending to write a JAX version of scipy.optimize.root that uses Newton's method, but right now I just have a prototype. I'll see if I can dig that up.

As far as performance gains go, the most obvious improvement to your code is that you can (and should) calculate both the forward and reverse passes inside your loop together, i.e., call value_and_grad(f)(x) instead of f(x) and grad(f)(x) separately:

@custom_jvp
def newton(x0):
    def body(x):
        y, dy = jax.value_and_grad(f)(x)
        return x - y / dy
    x = while_loop(lambda xp: abs(f(xp)) > 1e-8, # bug fixed here from previous version, need abs for residual!
                   body,
                   x0)
    return x

I'm pretty sure your code would not in fact recompute the non-linear Newton iterations in the backwards pass. You only call newton once, inside the JVP rule.

[mrtupek]

Fantastic! I've confirmed that it only calls newton one, even if that function is called within a longer reverse-mode differentiation chain. I'll admit to being a bit confused about how that works, but it seems to work great. I've changed to use value_and_grad, thanks. One final question:

What is the value of implementing custom vjp instead of jvp?

Based on what I've seen here one should always just implement the custom_jvp interface. What situation is implementing a custom_vjp preferred?

[shoyer]

Based on what I've seen here one should always just implement the custom_jvp interface. What situation is implementing a custom_vjp preferred?

Correct, custom_jvp is generally the preferred solution.

custom_vjp only makes sense for cases where you want to directly control the implementation of the reverse pass, e.g.,

• you are wrapping an external auto-diff system that only supports reverse-mode
• you want to do non-linear transformations on cotangents like gradient clipping (which can't be implemented via JVP rules)