Possible bug: acprop explodes during initial steps

[carlosgmartin]

There appears to be a bug in acprop that causes it to explode during its initial steps in some situations:

Code

import jax
import optax
from jax import lax, numpy as jnp
from matplotlib import pyplot as plt, rcParams


def rosenbrock_function(params, a=1, b=100):
    x, y = params
    return jnp.square(a - x) + b * jnp.square(y - jnp.square(x))


def main():
    lr = 1e-4
    iters = 1000

    params = jnp.zeros(2)

    fig, ax = plt.subplots(constrained_layout=True)

    for label, opt in [
        ("sgd", optax.sgd(lr)),
        ("adam", optax.adam(lr)),
        ("adabelief", optax.adabelief(lr)),
        ("acprop", optax.contrib.acprop(lr)),
    ]:

        def update(state, _):
            params, opt_state = state
            value, grads = jax.value_and_grad(rosenbrock_function)(params)
            updates, opt_state = opt.update(grads, opt_state, params)
            params = optax.apply_updates(params, updates)
            return (params, opt_state), value

        opt_state = opt.init(params)
        _, values = lax.scan(update, (params, opt_state), length=iters)
        ax.plot(values, label=label)

    ax.legend()
    ax.set(xlabel="iteration", ylabel="value", yscale="log")
    ax.grid(True, "both", alpha=0.1)
    rcParams["savefig.dpi"] = 300
    plt.show()


if __name__ == "__main__":
    main()

Zooming in:

I'm currently investigating this.