aura-optax

AURA (Angular Update Rate Adaptation) for Optax.

AURA scales the update direction of a base optimizer by a per-parameter multiplier γ, which increases when consecutive directions agree in the complex plane and decreases when they do not. It is designed for complex-valued neural networks and applies to real-valued parameters as well. The package provides three optax.GradientTransformations:

• aura_optax.scale_by_aura: AURA's multiplier alone, to be chained after the direction of any Optax optimizer;
• aura_optax.adam_aura: the Adam direction, scaled by AURA's multiplier, with decoupled weight decay;
• aura_optax.muon_aura: the Muon direction on weight matrices and the Adam direction on all other parameters, scaled by AURA's multiplier, with decoupled weight decay.

Installation

pip install aura-optax

The package requires Python ≥ 3.10, JAX and Optax ≥ 0.2.7. For GPU support, install the JAX build for your platform first, following the JAX installation guide.

Usage

scale_by_aura is a gradient transformation that multiplies each update by γ. It is placed in an optax.chain after the transformation that produces the direction (optax.scale_by_adam, optax.contrib.scale_by_muon, ...) and before the learning rate:

import optax

import aura_optax

optimizer = optax.chain(
    optax.scale_by_adam(),  # the direction of any Optax optimizer
    aura_optax.scale_by_aura(),  # AURA's multiplier
    optax.add_decayed_weights(1e-4),  # optional decoupled weight decay
    optax.scale_by_learning_rate(1e-2),
)

The complete training loop below fits a complex-valued linear model by least squares with this optimizer:

import jax
import jax.numpy as jnp
import optax

import aura_optax

key_x, key_w, key_b, key_init = jax.random.split(jax.random.key(0), 4)
x = jax.random.normal(key_x, (256, 8), dtype=jnp.complex64)
y = (
    x @ jax.random.normal(key_w, (4, 8), dtype=jnp.complex64).T
    + jax.random.normal(key_b, (4,), dtype=jnp.complex64)
)
params = {
    "w": 0.1 * jax.random.normal(key_init, (4, 8), dtype=jnp.complex64),
    "b": jnp.zeros(4, dtype=jnp.complex64),
}


def loss_fn(params):
    residual = x @ params["w"].T + params["b"] - y
    return jnp.mean(jnp.real(residual * jnp.conj(residual)))


optimizer = optax.chain(
    optax.scale_by_adam(),
    aura_optax.scale_by_aura(),
    optax.add_decayed_weights(1e-4),
    optax.scale_by_learning_rate(1e-2),
)
opt_state = optimizer.init(params)


@jax.jit
def train_step(params, opt_state):
    loss, grads = jax.value_and_grad(loss_fn)(params)
    grads = jax.tree.map(jnp.conj, grads)  # see the note below
    updates, opt_state = optimizer.update(grads, opt_state, params)
    return optax.apply_updates(params, updates), opt_state, loss


for step in range(1000):
    params, opt_state, loss = train_step(params, opt_state)
print(f"final loss: {loss:.3e}")

Complex gradients. For a real-valued loss L of complex parameters w = x + iy, jax.grad returns ∂L/∂x − i ∂L/∂y, which is the complex conjugate of the gradient g = ∂L/∂x + i ∂L/∂y whose negative is the steepest-descent direction. The gradients must therefore be conjugated before optimizer.update, as in the example above. The same conjugation is required by every Optax optimizer applied to complex parameters, and it leaves real-valued parameters unchanged.

Ready-made optimizers. adam_aura is the chain above; muon_aura applies AURA to the Muon direction on the 2-D parameters and to the Adam direction on all others, with its own gate values. Both take the learning rate as a scalar or an Optax schedule:

optimizer = aura_optax.adam_aura(learning_rate=1e-2)
optimizer = aura_optax.muon_aura(learning_rate=1e-2)

Monitoring the multiplier. AURA's state is the entry of opt_state at the position of scale_by_aura in the chain: opt_state[1] in the chain above and in adam_aura, opt_state[2] in muon_aura:

gamma = opt_state[1].multiplier  # same tree structure as params

Method

TODO

License

MIT; see LICENSE.