You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
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
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #1708.
Summary
s^T y.weights_memory.lbfgsandscale_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_memorynegative, 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.pyruff check .flake8 --select=E9,F63,F7,F82,E225,E251 --show-source --statisticspylintmodule, test, and quote-consistency checkspyrefly check optax/_src/transform.py optax/_src/alias.pypytest -n auto --pyargs optaxmake html -C docsmake doctest -C docs