fixed point iteration

fixed point in 1-D

Rewrite $f(x)=0$ into the form $x=g(x)$ and then label left hand side as $x_{n+1}$ and right hand side as $x_n$.

e.g. $x^2-x-1=0$ becomes $x_{n+1}=1+\frac{1}{x_n}$.

• converges when $|g^{\prime }(x)|<1$
• generally try and reduce the degree of the polynomial
• newton’s method is a special case of fixed point iteration where $g(x)=x-\frac{f(x)}{f^{\prime }(x)}$

Oscar Veliz video 1 and 2.

generalised fixed point iteration

Naive forward iteration, where we iterate until $z_{n+1}$ stays sufficiently close to $z_n$:

$$z_{n+1}=f(\theta ,z_n)$$

def fwd_solver(f, z_init):
	z_prev, z = z_init, f(z_init)
	while jnp.linalg.norm(z_prev - z) > 1e-5:
		z_prev, z = z, f(z)
	return z

gradients

If we use a fixed-point iteration in a model where we need the gradients, such as an implicit layer, naively using automatic differentiation would require the gradients at each intermediate step of the solver.

Instead, we can use the implicit function theorem to define the backward pass only at the solution point.

For the fixed point iterator, the Jacobian becomes (as we now have $z^{\ast }(\theta )=f(\theta ,z^{\ast }(\theta ))$ instead of $0$)

$\partial z^{\ast }({\theta }_0)=[I-{\partial }_{1}f({\theta }_0,z_0){]}^{-1}{\partial }_{0}f({\theta }_0,z_0).$

Jacobian-vector product

$\partial z^{\ast }({\theta }_0)v=[I-{\partial }_{1}f({\theta }_0,z_0){]}^{-1}{\partial }_{0}f({\theta }_0,z_0)v$

We want to avoid the matrix inversion by using the adjoint method.

1. Calculate $u={\partial }_{0}f({\theta }_0,z_0)v$ using jax.jvp.
2. The final solution is now $w=[I-{\partial }_{1}f({\theta }_0,z_0){]}^{-1}u$. So $w=u+{\partial }_{1}f({\theta }_0,z_0)w$ and we can use a fixed point iterator (or other solver) to solve the linear system for $w$.

vector-Jacobian product

$w^T\partial z^{\ast }({\theta }_0)=w^T[I-{\partial }_{1}f({\theta }_0,z_0){]}^{-1}{\partial }_{0}f({\theta }_0,z_0)$

1. Calculate the adjoint $u^T=w^T[I-{\partial }_{1}f({\theta }_0,z_0){]}^{-1}$ as $u^T=w^T+u^T{\partial }_{1}f({\theta }_0,z_0)$.
2. Compute $u^T{\partial }_{0}f({\theta }_0,z_0)$ using jax.vjp.

The jax code for the VJP:

from functools import partial
 
@partial(jax.custom_vjp, nondiff_argnums=(0, 1))
def fixed_point_layer(solver, f, params, x):
	z_star = solver(lambda z: f(params, x, z), z_init=jnp.zeros_like(x))
	return z_star
 
def fixed_point_layer_fwd(solver, f, params, x):
	z_star = fixed_point_layer(solver, f, params, x)
	return z_star, (params, x, z_star)
 
def fixed_point_layer_bwd(solver, f, res, z_star_bar):
	params, x, z_star = res
	_, vjp_a = jax.vjp(lambda params, x: f(params, x, z_star), params, x)
	_, vjp_z = jax.vjp(lambda z: f(params, x, z), z_star)
	return vjp_a(solver(lambda u: vjp_z(u)[0] + z_star_bar, z_init=jnp.zeros_like(z_star)))
 
fixed_point_layer.defvjp(fixed_point_layer_fwd, fixed_point_layer_bwd)