SOFOScan

class braintools.optim.SOFOScan(rnn_cell, 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)

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

Like SOFO, but for a stateful one-step recurrent Module. The model is scanned over the input sequence with the hidden state (“latent”) carried explicitly; forward-mode JVPs propagate the tangents through time automatically (jax.jvp through lax.scan), so the Generalised Gauss-Newton matrix is accumulated over every (timestep, batch) sample and solved once. The resulting direction is applied via the same SGD-style optax update as SOFO.

Parameters:

• rnn_cell (Callable) – Stateful Module called as rnn_cell(latent, inputs) -> (new_latent, output), using brainstate.ParamState objects internally for its trainable weights.

• loss_fn (Callable) – loss_fn(predictions, targets) -> scalar, where both arguments have their leading (time, batch) axes collapsed into a single sample axis.

• 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 Cell(brainstate.nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.wh = brainstate.nn.Linear(4, 4)
...         self.wx = brainstate.nn.Linear(3, 4)
...         self.wo = brainstate.nn.Linear(4, 2)
...     def __call__(self, latent, inp):
...         new_latent = self.wh(latent) + self.wx(inp)
...         return new_latent, self.wo(new_latent)
>>>
>>> cell = Cell()
>>> loss_fn = lambda pred, y: jnp.mean((pred - y) ** 2)
>>> opt = braintools.optim.SOFOScan(cell, loss_fn, lr=1e-2, tangent_size=64)
>>> opt.register_trainable_weights(cell.states(brainstate.ParamState))
>>> xs = jnp.ones((5, 4, 3)); ys = jnp.zeros((5, 4, 2)); z0 = jnp.zeros((4, 4))
>>> loss = opt.step(z0, (xs, ys))

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(z_init, batch)

Compute the recurrent SOFO direction and apply one optimization step.

Parameters:

• z_init (array or pytree) – Initial hidden (“latent”) state for the scan over the sequence.

• batch (tuple) – (inputs_seq, labels_seq) with leading (time, batch) axes.

Returns:

The loss evaluated at the parameters before this update.

Return type:

jax.Array

update(z_init, batch)

Alias of step() taking (z_init, batch) (not gradients).