Fix incorrect (1 - b2) coefficient in scale_by_adan

[codewithfourtix]

Summary

scale_by_adan scales the gradient-difference term by 1 - b2 in two
places, but the algorithm documented in optax.adan scales it by
(1 - beta_2), which equals b2 in optax's parameterization. With the
default b2=0.92 the optimizer applied a weight of 0.08 instead of
0.92, diverging from the documented equations from the second step onward.

Details

optax parameterizes Adan's moment EMAs via update_moment(..., decay=b2),
i.e. v_t = b2 * v_{t-1} + (1 - b2) * (g_t - g_{t-1}). Matching this to the
documented equation v_t = (1 - beta_2) v_{t-1} + beta_2 (g_t - g_{t-1})
forces beta_2 = 1 - b2. Therefore the documented coefficient (1 - beta_2)
that appears in:

• the squared accumulator n_t = ... (g_t + (1 - beta_2)(g_t - g_{t-1}))^2
• the update u_t = m_t + (1 - beta_2) v_t

equals b2, not 1 - b2. This also matches Algorithm 1 of
Xie et al., 2022 (https://arxiv.org/abs/2208.06677), where the term is
(1 - beta_2).

The first step is unaffected (the gradient difference is zero), which is why
the existing "loss decreases" tests did not catch this.

Changes

• scale_by_adan: use b2 instead of 1 - b2 for the gradient-difference
  term in both the n accumulator and the update direction.
• Add test_adan_matches_documented_update, which checks scale_by_adan
  against an independent reference implementation of the documented equations
  over several steps.

Compatibility

This changes the numerical behavior of optax.adan to match its documented
algorithm and the reference paper. Existing code using optax.adan will now
follow the intended Adan update.

[rdyro]

Thanks, you seem to have caught a genuine problem with the optax implementation!

[rdyro]

I'd prefer optax stick closer to the math and we keep 1 - b2, but update the constructor defaults.

[rdyro]

See comment above, use 0.02, etc.

[rdyro]

You'll need to update the docstring in adan docstring in alias.py too:

Document: api/generated/optax.adan
----------------------------------
**********************************************************************
File "../optax/_src/alias.py", line ?, in default
Failed example:
    for _ in range(5):
     grad = jax.grad(f)(params)
     updates, opt_state = solver.update(grad, opt_state, params)
     params = optax.apply_updates(params, updates)
     print('Objective function: {:.2E}'.format(f(params)))
Expected:
    Objective function: 1.28E+01
    Objective function: 1.17E+01
    Objective function: 1.07E+01
    Objective function: 9.68E+00
    Objective function: 8.76E+00
Got:
    Objective function: 1.28E+01
    Objective function: 1.17E+01
    Objective function: 1.07E+01
    Objective function: 9.69E+00
    Objective function: 8.77E+00

[codewithfourtix]

@rdyro Done, kept the 1 - b2 form and switched the defaults to the paper's betas (b1=0.02, b2=0.08, b3=0.01), flipping the moment/bias-correction decays to 1 - b_i so the EMAs stay consistent. Behavior is unchanged, the parameterization now matches the paper. Updated the test to the new defaults and the docstring example. Let me know if you'd like anything else.