Feature request: Adahessian optimizer

[viralvgupta]

I’d like to request adding AdaHessian to Optax, and I’m happy to implement it.

AdaHessian is a second‑order optimizer that uses a Hessian diagonal approximation for adaptive updates. I already have a working JAX version here and would adapt it to Optax’s optimizer structure:

Paper: https://arxiv.org/abs/2006.00719
PyTorch reference: https://github.com/amirgholami/adahessian
My JAX implementation: https://github.com/viralvgupta/adahessian-jax
If you’re open to it, I can start a PR and align with whatever Optax design conventions you prefer. Thanks!

[rdyro]

I think adding adahessian could be a good idea, though it'd be good to be mindful of how the hessian vector product oracle is be implemented.

Before starting a PR can you sketch out what your proposed design is?

[viralvgupta]

Thanks for the feedback! Here’s a concise design sketch focused on the HVP oracle + Optax integration:

Main concept (HVP oracle)
AdaHessian is Adam‑like but replaces gradient‑variance with a curvature estimate. We estimate the Hessian diagonal using Hutchinson’s method: sample Rademacher vectors v, compute H v via jvp(grad(obj_fn), v), and form (H v) ⊙ v. This yields an unbiased diagonal estimate without forming the full Hessian. We can optionally average multiple samples (n_samples), but default is 1 for cost (even pytorch reference implementation by Amir Gholami uses numSamples=1)

Optax mapping
hessian_diagonal_fn: pluggable GradientTransformationExtraArgs (default Hutchinson) that consumes obj_fn + params at update time, similar to contrib.sophia.

PyTorch precedent
The reference PyTorch AdaHessian uses a single Hutchinson probe per step:
v = 2*randint_like(p)-1 and autograd.grad(..., grad_outputs=v) to get HVPs.

If this is aligned, I can proceed with a PR following the contrib patterns.

[viralvgupta]

@rdyro I am sorry I "closed" the issue by mistake. Can you pls reopen the issue.

[rdyro]

Oh, nice, I hadn't realized Sophia already defines this pattern, it'd be great to follow it!

I think this would be a great contributions, can you aim to place it in contrib, please?

[viralvgupta]

Thanks for the quick response. Yes I am aiming to create the PR in couple of days. I want to explore if certain components can be abstracted like shared Hutchinson estimator utility - can be reused across many Hessian-based optimizers that need a diagonal Hessian estimate via stochastic probing.

[viralvgupta]

Created a PR #1604 for the optimizer.