How to clip gradients with L-BFGS

[SNMS95]

I have a rugged design space and would like to clip the gradients to allow L-BFGS as the optimizer. The issue is that I don't know how to achieve this. The issue lies in the value_fn argument that L-BFGS needs. For example

import optax
import jax
import jax.numpy as jnp

def fun(x):
    return jnp.sum(x**2)

x = jnp.array([1.0, 2.0, 3.0])
# optimizer
linesearch = optax.scale_by_zoom_linesearch(max_linesearch_steps=100, verbose=True)
opt = optax.chain(optax.clip_by_global_norm(0.01), 
                  optax.scale(-1.0),
                  optax.scale_by_lbfgs(), 
                  linesearch)
opt_state = opt.init(x)
losses = []
for i in range(20):
    value, grad = jax.value_and_grad(fun)(x)
    updates, opt_state = opt.update(grad, opt_state, x, grad=grad, value=value,
                                    value_fn=fun)
    x = optax.apply_updates(x, updates)
    losses.append(value)

print(x)
print(fun(x))
import matplotlib.pyplot as plt
plt.plot(losses)

Does the value_fn also consider clipped gradients? My understanding is that it does not. It get the grad and value at the starting point (which are clipped) but for further evaluations of the line-search, it sees a unclipped design space.
Is my understanding correct ?
If so, Do I need to write a custom ad rule for the value_fn to properly clip the gradients?

[vroulet]

Hello @SNMS95,

Does the value_fn also consider clipped gradients?

No, it considers the gradients of value_fn.

Is my understanding correct ?

Yes.

Do I need to write a custom ad rule for the value_fn to properly clip the gradients?

This won't solve the issue. As said in #1201, if the gradients you feed LBFGS do not match the gradients of the function you are trying to minimize, LBFGS cannot ensure that the direction it outputs is a descent direction and so the may linesearch fail.

If you could define a function f such that its gradient matches the clipped design, then you could use it but I'm afraid there's no way to derive such a f automatically (in other words autointegration does not exist).

If I was you, I would try to see what goes wrong with the verbose option of the linesearch first (to see if indeed there is a problem with the update direction fed to the linesearch not being a descent direction). Then you may try lbfgs without the linesearch but with a small learning rate. Finally, you may consider other quasi-newton methods that may have been developed recently.