OSTLRecurrent

OSTLRecurrent#

class braintrace.OSTLRecurrent#

OSTL ‘with-H’ regime — single-layer factorization, RTRL-exact only for block-diagonal hidden-to-hidden Jacobians.

OSTL [1] derives an online rule by cleanly separating the gradient into a temporal eligibility trace and a spatial learning signal. The ‘with-H’ regime retains the hidden-to-hidden Jacobian, so the trace carries the full temporal term:

\[\boldsymbol{\epsilon}^t = \mathbf{D}^t\,\boldsymbol{\epsilon}^{t-1} + \operatorname{diag}(\mathbf{D}_f^t)\otimes \mathbf{x}^t , \qquad \nabla_{\boldsymbol{\theta}}\mathcal{L} = \sum_t \frac{\partial \mathcal{L}^t}{\partial \mathbf{h}^t} \circ \boldsymbol{\epsilon}^t ,\]

where \(\mathbf{D}^t\) is the hidden-to-hidden Jacobian, \(\mathbf{D}_f^t\) the state-to-output Jacobian, and \(\mathbf{x}^t\) the presynaptic input. This is exactly the per-parameter D-RTRL trace (memory \(O(P\cdot H)\)), so the class delegates entirely to ParamDimVjpAlgorithm.

Accuracy caveat. The D-RTRL trace machinery underlying this class only ever retains the per-position block-diagonal of \(\mathbf{D}^t\) (HiddenGroup.diagonal_jacobian(), via block_diagonal_last_dim): cross-hidden-unit terms \(\partial h^t_p / \partial h^{t-1}_q\) for \(p \ne q\) are retained only insofar as they flow through a traced ETP weight (e.g. a recurrent matmul()); any other hidden-to-hidden mixing (e.g. a hand-written convolution/roll/mixing term not expressed as an ETP op) is not captured. Consequently the rule is gradient-equivalent to BPTT only when the hidden-to-hidden Jacobian is (effectively) block-diagonal in this sense — for a single recurrent layer whose only cross-unit coupling is the traced ETP recurrent weight, the two coincide to machine precision. If some other part of the model couples hidden units directly (bypassing the traced weight), the two diverge; see TestOSTLRecurrentVsBPTT in ostl_test.py for a worked example of both regimes.

Parameters:

  • model (brainstate.nn.Module) – The recurrent SNN whose weights are trained online.

  • name (str, optional) – Name of the algorithm instance.

  • vjp_method ({‘single-step’, ‘multi-step’}, optional) – VJP window used to compute the learning signal.

  • fast_solve (bool, optional) – Whether to use closed-form per-primitive contractions when available.

  • trace_dtype (dtype, optional) – Storage dtype for supported fast-path eligibility traces.

  • chunked_trace (bool, optional) – Whether multi-step inputs use the closed-form chunked trace roll.

  • control_flow (ControlFlowPolicy, optional) – Control-flow canonicalization policy used during graph compilation.

  • config (ETraceConfig, optional) – Learning-rule coordinates. None uses the OSTL recurrent preset.

  • random_feedback_key (jax.Array, optional) – Key for fixed random-feedback projections requested by config.

  • snap_max_jacobian_elements (int, optional) – Maximum permitted size of each SnAp widened block Jacobian.

Examples

>>> import brainstate
>>> import braintrace
>>> import jax.numpy as jnp
>>>
>>> class Net(brainstate.nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.cell = braintrace.nn.ValinaRNNCell(1, 20, activation='tanh')
...         self.out = braintrace.nn.Linear(20, 1)
...     def update(self, x):
...         return x >> self.cell >> self.out
>>>
>>> model = Net()
>>> x0 = brainstate.random.randn(1)
>>> # one call: initialise states, build the trace graph, return a learner
>>> learner = braintrace.compile(model, braintrace.OSTLRecurrent, x0)
>>> y = learner(x0)
>>>
>>> # etrace_grad drives the sequence and accumulates the online gradients
>>> xs = brainstate.random.randn(10, 1)   # (T, ...)
>>> ys = brainstate.random.randn(10, 1)
>>> def step_loss(x, y):
...     return jnp.mean((learner(x) - y) ** 2)
>>> grads, losses = learner.etrace_grad(xs, ys, step_fn=step_loss, return_value=True)

References

OSTLRecurrent.__init__(model, name=None, vjp_method='single-step', fast_solve=True, trace_dtype=None, chunked_trace=True, control_flow=None, config=None, random_feedback_key=None, snap_max_jacobian_elements=16777216)#