Question regarding example given for implicit differentiation

[SNMS95]

Hi everyone,

I am trying to learn automatic differentiation better. I was looking at the implicit differentiation section of the official docs. The example given is clear and informative. However, the code associated with it is a bit hard to digest for me. Therefore, I wrote an alternate code for the same

import jax.numpy as jnp
from jax import vjp, custom_vjp
from functools import partial
######## Unchanged ############
@partial(custom_vjp, nondiff_argnums=(0,))
def fixed_point(f, a, x_guess): 
  def cond_fun(carry):
    x_prev, x = carry
    return jnp.abs(x_prev - x) > 1e-6

  def body_fun(carry):
    _, x = carry
    return x, f(a, x)

  _, x_star = while_loop(cond_fun, body_fun, (x_guess, f(a, x_guess)))
  return x_star

######## Unchanged ############
def fixed_point_fwd(f, a, x_init):
  x_star = fixed_point(f, a, x_init)
  return x_star, (a, x_star) 

######## Modified ############
def fixed_point_rev(f, res, x_star_bar):
  a, x_star = res
  _, vjp_a = vjp(lambda a: f(a, x_star), a) # B in equation
  _, vjp_x = vjp(lambda x: f(a, x), x_star) # A in equation
  # Finding w using fixed_point
  v = x_star_bar
  params = (vjp_x, v)
  w = fixed_point(rev_f, params, x_star_bar)------------------------------------- WHY start with x_star_bar ?
  # Calculate the final expression
  a_bar = vjp_a(w)[0] # w^T.B
  return a_bar, jnp.zeros_like(x_star)

######## Modified ############  
def rev_f(params, u):
  vjp_x, v = params
  return v + vjp_x(u)[0] #  v + u^T.A ------------------------------------- WHY v and not it's transpose?

fixed_point.defvjp(fixed_point_fwd, fixed_point_rev)

This gives the same result for the first and second order gradient

def newton_sqrt(a):
  update = lambda a, x: 0.5 * (x + a / x)
  return fixed_point(update, a, a)

print(newton_sqrt(2.))
from jax import grad
print(grad(newton_sqrt)(2.))
print(grad(grad(newton_sqrt))(2.))

I have three questions regarding this:

1. Is this implementation correct ?
2. Why do we start the fixed_point in the backward pass with u_init = x_start_bar ?
3. Shouldn't we be adding v.T instead of v in the output of rev_f ?

[anh-tong]

Hi,

This may not completely help you but here is my thought.

1. I believe the code is correct, (my justification for this follows)

2. To understand what x_start_bar is, we may need take a detour. We begin with $x=f(a,x)$ (after solving the fixed point equation). We get the output and denote the ouput $F: = f(a, x)$. This output will be used to compute a loss, let say $\mathcal{L}$. Now we perform backward from the loss to $F$ and now we obtain $\frac{\partial \mathcal L}{\partial F}$. Mathematically, x_star_bar is $\frac{\partial \mathcal L}{\partial F}$ as the initial value for solving fixed point of backward pass (using custom VJP instead of JAX autodiff)
   If you find this explanation is still difficult to follow, I recommend watching this tutorial by @mattjj starting around 11:00 to 25:00.

3. You can look up again JAX doc (this section) which explains the difference between VJP in math and VJP in JAX code.
   In math the notation is
   $$(x, v) \rightarrow v \partial f(x)$$
   Whereas in JAX code, the notation is
   $$(x, v) \rightarrow (f(x), v^\top \partial f(x))$$
   So JAX applies transpose to $v$ before multiplying it with Jacobian.

I hope this helps clarify the concept for you.

[SNMS95]

Thanks a lot. It does help!