Skip to content

Skip invalid L-BFGS curvature updates - #1749

Open
Trighap52 wants to merge 1 commit into
google-deepmind:mainfrom
Trighap52:issue-1708-lbfgs-skip-updates
Open

Skip invalid L-BFGS curvature updates#1749
Trighap52 wants to merge 1 commit into
google-deepmind:mainfrom
Trighap52:issue-1708-lbfgs-skip-updates

Conversation

@Trighap52

Copy link
Copy Markdown

Fixes #1708.

Summary

  • Skip L-BFGS memory updates when the curvature pair has non-positive s^T y.
  • Avoid writing negative inverse-curvature weights into weights_memory.
  • Add regression coverage for the backtracking line-search reproducer from the issue.
  • Document the skip behavior for lbfgs and scale_by_lbfgs.

Why

When backtracking/Armijo line search is used, L-BFGS can observe a parameter/gradient pair that does not satisfy the positive-curvature condition. Storing that pair makes weights_memory negative, which can break the inverse Hessian approximation. This change leaves such slots as zero so the two-loop recursion ignores them.

Tests

  • pre-commit run --files optax/_src/transform.py optax/_src/alias.py optax/_src/lbfgs_test.py
  • ruff check .
  • flake8 --select=E9,F63,F7,F82,E225,E251 --show-source --statistics
  • pylint module, test, and quote-consistency checks
  • pyrefly check optax/_src/transform.py optax/_src/alias.py
  • pytest -n auto --pyargs optax
  • make html -C docs
  • make doctest -C docs

@Trighap52
Trighap52 marked this pull request as ready for review August 6, 2026 20:44
@Trighap52

Copy link
Copy Markdown
Author

I looked into the failing jobs.

Pytest 3.10 on ubuntu-latest jax=newest is related to this PR. It fails in:

LBFGSTest::test_against_plain_implementation3
problem_name = 'himmelblau', scale_init_precond = False
ACTUAL:  [234.10304, 179.07068]
DESIRED: [234.02055, 179.15878]

I reproduced this locally with jax==0.6.2. The L-BFGS curvature guard itself is working; the mismatch appears to be a compiled lax.while_loop float32 numerical divergence in the existing plain-implementation comparison after the first tiny positive curvature pair.

Would you prefer that I open a small follow-up PR to relax the tolerance for this comparison while keeping the new regression test strict about preventing negative weights_memory?

The Build and check types with pyrefly failure appears unrelated to this PR. The reported errors are in untouched upstream files such as optax/_src/linear_algebra.py, optax/_src/linesearch.py, optax/perturbations/_make_pert.py, and optax/projections/_projections.py; none are in the files changed here.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

LBFGS - Add functionality to skip updates

1 participant