OSTLFeedforward

OSTLFeedforward#

class braintrace.OSTLFeedforward#

OSTL ‘without-H’ regime — feedforward / no recurrent Jacobian.

The ‘without-H’ regime drops the hidden-to-hidden Jacobian \(\mathbf{D}^t\), so the temporal term of the eligibility trace vanishes and only the instantaneous (spatial) contribution survives:

\[\boldsymbol{\epsilon}^t \approx \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 .\]

This is the appropriate approximation for feed-forward SNNs in the OSTL construction [1]. It is realized by delegating to pp_prop (the input-output factorized trace) with a negligible decay, so the trace does not accumulate across time.

Parameters:

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

  • decay_or_rank (float or int, default 1e-6) – Exponential-smoothing factor of the IO-dim trace. The tiny default makes the temporal contribution negligible, matching the ‘without-H’ regime. A float must lie in [0, 1); an int is read as an approximation rank.

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

  • **kwargs (Any) – Additional options forwarded to pp_prop, including vjp_method, fast_solve, control_flow, config, and random_feedback_key.

Notes

The default 1e-6 is negligible, not zero, so in axis terms the coordinate is temporal_recursion=('scalar_leak', 'jacobian') with a tiny coefficient — both recursion terms are structurally present. The exact temporal_recursion='none' coordinate is decay_or_rank=0.0 (or equivalently decay_or_rank=1, since rank 1 maps to decay 0). The default is left at 1e-6 because changing it would move this preset’s gradients.

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.OSTLFeedforward, 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

OSTLFeedforward.__init__(model, decay_or_rank=1e-06, name=None, **kwargs)#