LBFGS - Add functionality to skip updates

[inversecrime]

Would it be possible to add the ability to skip updates to the lbfgs algorithm, in case a certain curvature condition does not hold? You can find the relevant information on page 537 of the book "Numerical Optimization" by Nocedal and Wright.

While it isn't great to skip updates, it prevents the algorithm from breaking when the Wolfe conditions are not satisfied (sometimes, one needs to use backtracking line search due to computational restrictions).

This would mean to add about 20 lines to the existing optax code - I'm happy to provide my current implementation.

[rdyro]

Do you mean the zoom line search, the default line search in L-BFGS?

I think it already has fallback for it can't satisfy Wolfe conditions, what change did you have in mind?

[inversecrime]

I'm specifically talking about backtracking/Armijo line search which does not require any gradient evaluations (which might be relevant for functions where calculating the gradient is very expensive). For example, it could be implemented like this:

Or just with a straightforward check if $s^T y > 0$.

Where it could be applied:

import jax
import jax.numpy as jnp
import optax as tx

def f(x):
    return jnp.sum(jnp.abs(100 - x**3))

optimizer = tx.lbfgs(scale_init_precond=False, linesearch=tx.scale_by_backtracking_linesearch(max_backtracking_steps=1000, slope_rtol=0.7))

x = jnp.ones(2)
state = optimizer.init(x)
for _ in range(3):
    (y, g) = jax.value_and_grad(f)(x)
    (updates, state) = optimizer.update(g, state, x, value=y, grad=g, value_fn=f)
    x = tx.apply_updates(x, updates)
    print(x, f(x))

print("weights:", state[0].weights_memory)

As one can see, the lbfgs-memory contains negative weights, which destroys the algorithm.