Metadata-Version: 2.5
Name: aura-optax
Version: 0.1.0
Summary: AURA (Angular Update Rate Adaptation): a per-parameter step-size multiplier for Adam and Muon on complex-valued and real-valued parameters, as Optax gradient transformations
Project-URL: Repository, https://github.com/enricoballini/aura-optax
Project-URL: Issues, https://github.com/enricoballini/aura-optax/issues
Author: Enrico Ballini
License-Expression: MIT
License-File: LICENSE
Keywords: adam,complex-valued neural networks,jax,muon,optax,optimizer,step size
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Requires-Dist: jax>=0.5.3
Requires-Dist: optax>=0.2.7
Provides-Extra: test
Requires-Dist: numpy; extra == 'test'
Requires-Dist: pytest>=7; extra == 'test'
Description-Content-Type: text/markdown

# aura-optax

AURA (Angular Update Rate Adaptation) for [Optax](https://github.com/google-deepmind/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.GradientTransformation`s:

- `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

```bash
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](https://docs.jax.dev/en/latest/installation.html).

## 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:

```python
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:

```python
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:

```python
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`:

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

## Method
TODO

## License

MIT; see [LICENSE](LICENSE).
