google-deepmind/kfac-jax

Second Order Optimization and Curvature Estimation with K-FAC in JAX.

What it solves

KFAC-JAX provides an easy-to-use implementation of the K-FAC (Kronecker-factored Approximate Curvature) optimizer and curvature estimator. It addresses the challenge of performing second-order optimization of neural networks, which is typically more computationally expensive than first-order methods, by providing scalable approximations of the curvature matrix.

How it works

Built on top of JAX, the library allows researchers to integrate K-FAC optimization into their models. Users must register their specific loss function (e.g., using register_softmax_cross_entropy_loss) so the library can correctly approximate the curvature matrix. The optimizer handles internal staging (via jax.jit or jax.pmap) automatically, meaning users should not manually apply these to the optimizer's step function or the loss function.

Who it’s for

AI researchers who are using JAX and Haiku to train neural networks and want to leverage second-order optimization techniques to improve training efficiency or convergence.

Highlights

  • JAX-native: Built specifically for the library's ecosystem, integrating with Haiku.
  • Scalable Curvature: Implements K-FAC for efficient second-order optimization.
  • Automatic Staging: Manages jit and pmap internally to simplify the user experience.
  • Adaptive Parameters: Supports adaptive learning rates, momentum, and damping.

Related

  • Project
    google-deepmind/optax

    A gradient processing and optimization library for JAX that provides composable building blocks for creating custom optimizers and loss functions.

  • Project
    google/rax

    A JAX-based Learning-to-Rank library providing off-the-shelf ranking losses, metrics, and transformations for optimizing the relative order of items.

  • Project
    patrick-kidger/lineax

    A JAX library for linear solves and linear least squares, providing numerically stable routines for solving linear systems, including support for linear operators and GPU/TPU acceleration.

  • Project
    google-research/scenic

    A JAX and Flax-based library for computer vision research that provides scalable libraries and project templates for developing attention-based models across images, video, and audio.

  • Project
    RobertTLange/evosax

    evosax is a JAX‑based library that implements 30+ Evolution Strategies (e.g., CMA‑ES, OpenAI‑ES, ARS, SimpleGA) with a vectorised `ask‑eval‑tell` API, enabling fast, hardware‑accelerated black‑box optimisation for neuro‑evolution, RL, and hyper‑parameter search.