Skip to content

Issue #1658 - Add SSBroyden and SSBFGS optimizers to optax.contrib #1659

Open
armbrusl wants to merge 3 commits intogoogle-deepmind:mainfrom
armbrusl:main
Open

Issue #1658 - Add SSBroyden and SSBFGS optimizers to optax.contrib #1659
armbrusl wants to merge 3 commits intogoogle-deepmind:mainfrom
armbrusl:main

Conversation

@armbrusl
Copy link
Copy Markdown

Port the Self-Scaled Broyden and Self-Scaled BFGS quasi-Newton optimizers from PyTorch (SciMBA) to JAX/optax. These maintain a dense inverse Hessian approximation with self-scaling updates, complementing the existing L-BFGS implementation for small to medium scale problems.

New public API:

  • optax.contrib.ssbroyden() — SS-Broyden variant
  • optax.contrib.ssbfgs() — SS-BFGS variant
  • optax.contrib.scale_by_ss_quasi_newton() — shared core transform

Both integrate with scale_by_zoom_linesearch (same pattern as optax.lbfgs) and include tests on quadratic and Rosenbrock functions.

Reference: Urbán et al. (2025), Journal of Computational Physics, 523, 113656.

Port the Self-Scaled Broyden and Self-Scaled BFGS quasi-Newton
optimizers from PyTorch (SciMBA) to JAX/optax. These maintain a
dense inverse Hessian approximation with self-scaling updates,
complementing the existing L-BFGS implementation for small to
medium scale problems.

New public API:
- optax.contrib.ssbroyden() — SS-Broyden variant
- optax.contrib.ssbfgs() — SS-BFGS variant
- optax.contrib.scale_by_ss_quasi_newton() — shared core transform

Both integrate with scale_by_zoom_linesearch (same pattern as
optax.lbfgs) and include tests on quadratic and Rosenbrock functions.

Reference: Urbán et al. (2025), Journal of Computational Physics, 523, 113656.
@google-cla
Copy link
Copy Markdown

google-cla Bot commented Apr 10, 2026

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@armbrusl
Copy link
Copy Markdown
Author

I have added the CLA.

@armbrusl armbrusl changed the title Add SSBroyden and SSBFGS optimizers to optax.contrib #1658 Issue #1658 - Add SSBroyden and SSBFGS optimizers to optax.contrib Apr 10, 2026
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.

1 participant