Skip invalid L-BFGS curvature updates

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

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(

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(