diff --git a/optax/_src/alias.py b/optax/_src/alias.py index 4b6ff23d2..8e3220368 100644 --- a/optax/_src/alias.py +++ b/optax/_src/alias.py @@ -2703,7 +2703,9 @@ def lbfgs( preconditioning matrix subject to some secant condition, see references for more details. Computing :math:`P_k u_k` can be done by a sequence of vector operations using past differences of parameters and gradients stored in - a memory bufffer. + a memory buffer. + Pairs whose curvature :math:`\delta u_k^\top \delta w_k` is not positive are + skipped to preserve a positive-definite inverse Hessian approximation. The present function just outputs the LBFGS direction :math:`P_k u_k`. It can be chained with a linesearch ensuring sufficient decrease and low diff --git a/optax/_src/lbfgs_test.py b/optax/_src/lbfgs_test.py index a6b11f763..42397c596 100644 --- a/optax/_src/lbfgs_test.py +++ b/optax/_src/lbfgs_test.py @@ -207,16 +207,28 @@ def _plain_lbfgs( else: identity_scale = jnp.vdot(dus[-1], dws[-1]) identity_scale /= jnp.sum(dus[-1] ** 2) + identity_scale = jnp.where( + jnp.isnan(identity_scale), 1.0, identity_scale + ) else: identity_scale = 1.0 - direction = -_plain_preconditioning(dws, dus, g, identity_scale) + valid_pairs = [ + (dw, du) for dw, du in zip(dws, dus) if jnp.vdot(dw, du) > 0.0 + ] + valid_dws = [dw for dw, _ in valid_pairs] + valid_dus = [du for _, du in valid_pairs] + direction = -_plain_preconditioning(valid_dws, valid_dus, g, identity_scale) w_old, g_old = w, g w = w + stepsize * direction _, g = value_and_grad_fun(w) - dws.append(w - w_old) - dus.append(g - g_old) + dw = w - w_old + du = g - g_old + if jnp.vdot(dw, du) <= 0.0: + dw, du = jnp.zeros_like(dw), jnp.zeros_like(du) + dws.append(dw) + dus.append(du) if len(dws) > memory_size: dws = dws[1:] # Pop left. @@ -604,6 +616,33 @@ def fun(x): sol, _ = _run_opt(opt, fun, init_params=jnp.ones(n), tol=tol) test_utils.assert_trees_all_close(sol, jnp.zeros(n), atol=tol, rtol=tol) + def test_lbfgs_skips_nonpositive_curvature_update(self): + def fun(x): + return jnp.sum(jnp.abs(100 - x**3)) + + opt = alias.lbfgs( + scale_init_precond=False, + linesearch=_linesearch.scale_by_backtracking_linesearch( + max_backtracking_steps=1000, slope_rtol=0.7 + ), + ) + + params = jnp.ones(2) + state = opt.init(params) + + for _ in range(3): + value, grad = jax.value_and_grad(fun)(params) + updates, state = opt.update( + grad, state, params, value=value, grad=grad, value_fn=fun + ) + params = update.apply_updates(params, updates) + + lbfgs_state = state[0] + self.assertTrue(bool(jnp.all(lbfgs_state.weights_memory >= 0.0))) + test_utils.assert_trees_all_close( + lbfgs_state.weights_memory[0], jnp.array(0.0) + ) + @parameterized.product( linesearch=[ _linesearch.scale_by_backtracking_linesearch( diff --git a/optax/_src/transform.py b/optax/_src/transform.py index ec2e8a2b5..7f06df827 100644 --- a/optax/_src/transform.py +++ b/optax/_src/transform.py @@ -1680,6 +1680,8 @@ def scale_by_lbfgs( for more details. Computing :math:`P_k u_k` can be done by a sequence of vector operations using past differences of parameters and gradients stored in a memory buffer. + Pairs whose curvature :math:`\delta u_k^\top \delta w_k` is not positive are + skipped to preserve a positive-definite inverse Hessian approximation. The present function just outputs the LBFGS direction :math:`P_k u_k`. It can be chained with a linesearch ensuring sufficient decrease and low @@ -1758,15 +1760,18 @@ def update_fn( vdot_diff_params_updates = optax.tree.real( optax.tree.vdot(diff_updates, diff_params) ) - weight = jnp.where( - vdot_diff_params_updates == 0.0, 0.0, 1.0 / vdot_diff_params_updates + update_memory = jnp.logical_and( + jnp.greater(state.count, 0), jnp.greater(vdot_diff_params_updates, 0.0) ) + safe_vdot_diff_params_updates = jnp.where( + update_memory, vdot_diff_params_updates, 1.0 + ) + weight = 1.0 / safe_vdot_diff_params_updates # params_diff, updates_diff, weight depend on differences of parameters - # that are not defined at the first iteration. Hence we keep them at 0 if - # state.count = 0. + # that are not defined at the first iteration. We also skip memory updates + # when the curvature condition is not satisfied. diff_params, diff_updates, weight = jax.tree.map( - # pyrefly: ignore[unsupported-operation] - lambda x: jnp.where(state.count > 0, x, jnp.zeros_like(x)), + lambda x: jnp.where(update_memory, x, jnp.zeros_like(x)), (diff_params, diff_updates, weight), ) diff_params_memory, diff_updates_memory, weights_memory = jax.tree.map(