Self-Scaled Broyden Family of Quasi-Newton Methods in JAX

TL;DR

This paper provides a JAX-compatible implementation of the Self-Scaled Broyden family of quasi-Newton optimization methods, including BFGS, DFP, and Broyden variants, to facilitate their adoption in the JAX community.

Contribution

It documents and shares an implementation of various quasi-Newton methods in JAX, enhancing accessibility and usability for practitioners.

Findings

Implementation includes BFGS, DFP, Broyden, and their Self-Scaled variants.

Includes a Zoom line search satisfying strong Wolfe conditions.

Aims to ease adoption within the JAX community.

Abstract

We present a JAX implementation of the Self-Scaled Broyden family of quasi-Newton methods, fully compatible with JAX and building on the Optimistix~\cite{rader_optimistix_2024} optimisation library. The implementation includes BFGS, DFP, Broyden and their Self-Scaled variants(SSBFGS, SSDFP, SSBroyden), together with a Zoom line search satisfying the strong Wolfe conditions. This is a short technical note, not a research paper, as it does not claim any novel contribution; its purpose is to document the implementation and ease the adoption of these optimisers within the JAX community. The code is available at https://github.com/IvanBioli/ssbroyden_optimistix.git.