SOFO

class braintools.optim.SOFO(model, loss_fn, lr=0.001, loss='mse', tangent_size=100, damping=1e-05, momentum=0.0, nesterov=False, weight_decay=0.0, grad_clip_norm=None, grad_clip_value=None, key=None)

Second-Order Forward-mode Optimization (SOFO) optimizer.

SOFO computes its own search direction by sampling random tangent vectors, taking forward-mode JVPs through model and loss_fn, building a Generalised Gauss-Newton matrix in the random subspace, solving a damped linear system, and projecting back to parameter space. The resulting direction is applied via an SGD-style optax update (so learning-rate scheduling, momentum, weight decay, and gradient clipping all work).

Parameters:

• model (Callable) – The network, called as model(inputs) returning predictions. Its trainable parameters are the brainstate.ParamState objects registered via register_trainable_weights().

• loss_fn (Callable) – loss_fn(predictions, targets) -> scalar.

• lr (float) – Learning rate.

• loss (str) – Selects the Generalised Gauss-Newton form.

• tangent_size (int) – Number of random tangents / subspace dimension.

• damping (float) – Damping on the GGN, scaled by the largest singular value.

• momentum (float) – Momentum for the SGD-style update.

• nesterov (bool) – Whether to use Nesterov momentum.

• weight_decay (float) – Decoupled weight decay.

• grad_clip_norm (float | None) – Clip the SOFO direction by global norm before the update.

• grad_clip_value (float | None) – Clip the SOFO direction by value before the update.

• key (jax PRNG key, optional) – Random key for tangent sampling. Defaults to brainstate.random.split_key() each step.

Examples

>>> import brainstate
>>> import braintools
>>> import jax.numpy as jnp
>>>
>>> class MLP(brainstate.nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.l1 = brainstate.nn.Linear(8, 16)
...         self.l2 = brainstate.nn.Linear(16, 3)
...     def __call__(self, x):
...         import jax
...         return self.l2(jax.nn.relu(self.l1(x)))
>>>
>>> model = MLP()
>>> loss_fn = lambda pred, y: jnp.mean((pred - y) ** 2)
>>> opt = braintools.optim.SOFO(model, loss_fn, lr=1e-2, tangent_size=64)
>>> opt.register_trainable_weights(model.states(brainstate.ParamState))
>>> loss = opt.step(jnp.ones((4, 8)), jnp.zeros((4, 3)))

default_tx()

Create default gradient transformation with clipping and weight decay.

register_trainable_weights(param_states)

Register trainable weights and initialize optimizer state.

Parameters:

param_states – A pytree (dict) of brainstate.State objects representing parameters.

step(inputs, targets)

Compute the SOFO direction and apply one optimization step.

Unlike the base optimizer (whose step/update take precomputed gradients), SOFO computes its own search direction internally, so this method takes the model inputs and targets instead.

Parameters:

• inputs (array or pytree) – Inputs passed to model.

• targets (array or pytree) – Targets passed to loss_fn.

Returns:

The loss evaluated at the parameters before this update.

Return type:

jax.Array

update(inputs, targets)

Alias of step() taking (inputs, targets) (not gradients).