google-deepmind / google-deepmind/optax

Add Self-Scaled Broyden (SSBroyden) and SSBFGS optimizers

Open
#1,658 1 comment 0 reactions 0 assignees View on GitHub
type:feature
Dominant language
Python
Stars
2.3k
Forks
369
Avg merge
10h 15m
Merged PRs (30d)
7

Description

Hi! Id like to propose adding the SSBroyden and SSBFGS optimizers to `optax.contrib`.

These are full-matrix quasi-Newton methods from Urbán et al. (2025) that use a self-scaled inverse Hessian update. The paper demonstrates strong results on PINNs and similar small-to-medium scale problems.

I've ported the [PyTorch implementation from SciMBA](https://www.scimba.org/) to JAX/optax following the existing patterns (`scale_by_lbfgs` / `lbfgs`). The implementation includes:

- `ssbroyden()` and `ssbfgs()` as separate public optimizers
- A shared `scale_by_ss_quasi_newton()` core transformation
- Integration with `scale_by_zoom_linesearch` (same pattern as `optax.lbfgs`)
- Tests on quadratic and Rosenbrock functions

I have a working implementation ready. Would you be open to a PR?

**Reference:**
Urbán, J. F., Stefanou, P., & Pons, J. A. (2025). Unveiling the optimization process of physics informed neural networks: How accurate and competitive can PINNs be? *Journal of Computational Physics*, 523, 113656. [DOI](https://doi.org/10.1016/j.jcp.2024.113656)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.