google-deepmind / google-deepmind/optax
Add Self-Scaled Broyden (SSBroyden) and SSBFGS optimizers
- 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
Assessment
This issue has not been assessed yet.