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

Ivan Bioli; Mikel Mendibe Abarrategi

cs.MSMar 11, 2026

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

Authors: Ivan Bioli, Mikel Mendibe Abarrategi

Organizations: Dipartimento di Matematica, Universit`a degli Studi di Pavia, Via A. Ferrata 5, 27100 Pavia, Italy · Dipartimento di Ingegneria Civile e Architettura, Universit`a degli Studi di Pavia, Via A. Ferrata 3, 27100 Pavia, Italy · University of the Basque Country, 48013 Bilbao, Spain · TECNALIA, Basque Research & Technology Alliance (BRTA), 48160 Derio, Spain

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), as well as L-BFGS, 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 and on PyPI (pip install ssbroyden-optimistix).

Figures & tables

Solver	$\theta_{k}$	$\tau_{k}$	Description
BFGS	$0$	$1$	Classic BFGS
SSBFGS	$0$	computed	Self-Scaled BFGS
DFP	$1$	$1$	Classic DFP
SSDFP	$1$	computed	Self-Scaled DFP
Broyden	computed	$1$	Broyden family (no scaling)
SSBroyden	computed	computed	Self-Scaled Broyden family

Table 1: Solvers implemented as special cases of the Self-Scaled Broyden family.

Figure 1: Convergence of quasi-Newton solvers on the 3D Poisson PINN problem. The self-scaled variants (SSBFGS, SSBroyden) achieve lower errors in fewer iterations compared to the standard BFGS and Broyden methods.

Quasi-Newton Methods

Quasi-Newton methods approximate curvature from gradient and step information to accelerate optimization without forming full Hessians.

See all 15 papers

asdex: Automatic Sparse Differentiation in JAX
Oct 8, 2026Adrian Hill, Guillaume DalleAutomatic Differentiation Structured Sparsity
BIFOLD – Berlin Institute for the Foundations of Learning and Data, Berlin, Germany · Machine Learning Group, Technical University of Berlin, Berlin, Germany · LVMT, ENPC, Institut Polytechnique de Paris, Univ Gustave Eiffel, Marne-la-Vallée, France
Quasi-Quadratic Gradient: A New Direction for Accelerating the BFGS Method in Quasi-Newton Optimization
Apr 27, 2026John ChiangSecond-Order Optimization Quasi-Newton Methods
SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients
Mar 9, 2026Anselm Paulus, A. René Geist, Vít Musil +3Straight-Through Estimator Differentiable Programming
University of Tübingen, Germany · Max Planck Institute for Intelligent Systems, Tübingen, Germany · Masaryk University, Czechia +1

Oct 8, 2026J/K move · Enter open · S save

asdex: Automatic Sparse Differentiation in JAX

Adrian Hill, Guillaume Dalle

BIFOLD – Berlin Institute for the Foundations of Learning and Data, Berlin, Germany · Machine Learning Group, Technical University of Berlin, Berlin, Germany · LVMT, ENPC, Institut Polytechnique de Paris, Univ Gustave Eiffel, Marne-la-Vallée, France

Automatic Differentiation Structured Sparsity

Open paper AskPDF arXiv

Many tasks in scientific computing and machine learning require the Jacobian or Hessian matrix of a function. Automatic differentiation (AD) computes these derivatives to machine precision, but materializing a dense

m \times n

Jacobian requires

n

forward-mode or

m

reverse-mode AD passes, one per column or row. For a large class of functions, each output depends on only a few inputs, making the derivative matrix sparse. Automatic sparse differentiation (ASD) exploits this structure in four steps: detection of the input-agnostic sparsity pattern, coloring of a graph to group columns or rows that can share an AD pass, compressed differentiation to compute a compressed derivative matrix with one AD pass per color, and finally decompression into the original sparsity pattern. The number of colors, and hence of AD passes, is often independent of the problem dimension: a banded Jacobian with

b

contiguous bands, for instance, only ever requires

b

colors, regardless of its size. asdex offers the first standalone ASD toolkit in the popular JAX ecosystem. With

\texttt{asdex.jacobian}

and

\texttt{asdex.hessian}

, it provides sparse drop-in replacements for

\texttt{jax.jacobian}

and

\texttt{jax.hessian}

Abstract

Quasi-Newton Methods

Explore similar work

asdex: Automatic Sparse Differentiation in JAX

Quasi-Quadratic Gradient: A New Direction for Accelerating the BFGS Method in Quasi-Newton Optimization

SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients