Optimistix: modular optimisation in JAX and Equinox

TL;DR

Optimistix presents a modular nonlinear optimisation library in JAX+Equinox built around two abstractions, search and descent, which generalise line search, trust-region, and learning-rate strategies to enable plug-and-play construction of custom optimisers. It supports minimisation, nonlinear least-squares, root-finding, and fixed-point iterations with automatic problem-type conversion and differentiable workflows via the implicit-function-theorem and online treeverse backpropagation. Empirical results show competitive runtimes for standard solvers (notably BFGS) and robust performance for Levenberg–Marquardt, with markedly faster compile-times than JAXopt, illustrating practical benefits for sciML and differentiable computing. By unifying optimisation components into modular primitives and enabling mix-and-match customisers, Optimistix aims to accelerate prototyping and deployment of advanced optimisation methods in differentiable scientific computing.

Abstract

We introduce Optimistix: a nonlinear optimisation library built in JAX and Equinox. Optimistix introduces a novel, modular approach for its minimisers and least-squares solvers. This modularity relies on new practical abstractions for optimisation which we call search and descent, and which generalise classical notions of line search, trust-region, and learning-rate algorithms. It provides high-level APIs and solvers for minimisation, nonlinear least-squares, root-finding, and fixed-point iteration. Optimistix is available at https://github.com/patrick-kidger/optimistix.

Figures (9)

• Figure 1: BFGS runtime comparison
• Figure 2: Levenberg-Marquardt runtime comparison
• Figure 3: BFGS compile time comparison
• Figure 4: Levenberg-Marquardt compile time comparison
• Figure 5: Optimistix vs JAXopt BFGS