OSTTP

class braintrace.OSTTP(model, B_list, target_timing='per-step', name=None, vjp_method='single-step', fast_solve=True, **kwargs)

Online Spatio-Temporal Learning with Target Projection.

OSTTP reuses the OSTL / D-RTRL per-parameter eligibility trace but replaces the back-propagated learning signal with a direct random target projection (DRTP):

\[\boldsymbol{\epsilon}^t \approx \mathbf{D}^t\,\boldsymbol{\epsilon}^{t-1} + \operatorname{diag}(\mathbf{D}_f^t)\otimes \mathbf{x}^t , \qquad L_l^t = y^{*\,t}\, B_l , \qquad \nabla_{W}\mathcal{L} = \sum_t L^t \circ \boldsymbol{\epsilon}^t ,\]

where \(y^{*\,t}\) is the task target at time \(t\), \(B_l \in \mathbb{R}^{n_\text{target}\times n_l}\) is a fixed random feedback matrix for HiddenGroup \(l\) (frozen via stop_gradient), \(\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.

How it works. The eligibility trace carries the temporal credit exactly as in OSTLRecurrent (‘with-H’), but the spatial credit normally obtained by back-propagating \(\partial \mathcal{L}/\partial h\) is replaced by a frozen random projection of the target. Because the projection matrices \(B_l\) are fixed, there is no weight transport and no backward pass — the rule is fully forward and update-unlocked in both space and time.

Parameters:

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

• B_list (Sequence[Array]) – One feedback matrix per HiddenGroup, each of shape (n_target, n_l). Frozen via stop_gradient at construction; the count and trailing dimension are validated against the compiled graph.

• target_timing (str) – 'per-step' requires y_target at every update() call. 'sequence-end' zeros the learning signal on intermediate steps (the trace still accumulates) and applies the projection only when y_target is supplied.

• name (str | None) – Name of the algorithm instance.

• vjp_method (str) – Forwarded verbatim to ParamDimVjpAlgorithm.

• fast_solve (bool) – Forwarded verbatim to ParamDimVjpAlgorithm.

Raises:

ValueError – If target_timing is invalid; if len(B_list) differs from the number of HiddenGroups; if a matrix’s trailing dimension does not match its HiddenGroup width; or if target_timing='per-step' and y_target is omitted from an update() call.

Examples

>>> import brainstate
>>> import jax
>>> import braintrace
>>>
>>> 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 (n_target, n_l) feedback matrix per HiddenGroup (here n_l = 20)
>>> B = jax.random.normal(jax.random.PRNGKey(0), (1, 20))
>>> # ``compile`` initialises states + builds the trace graph in one call
>>> learner = braintrace.compile(model, braintrace.OSTTP, x0, B_list=[B])
>>> y = learner.update(x0, y_target=brainstate.random.randn(1))

References

[1]

Ortner, T., Pes, L., Gentinetta, J., Frenkel, C., & Pantazi, A. (2023). “Online Spatio-Temporal Learning with Target Projection.” 2023 IEEE 5th International Conference on Artificial Intelligence Circuits and Systems (AICAS), 1-5. https://doi.org/10.1109/AICAS57966.2023.10168623 (arXiv:2304.05124)

[2]

Frenkel, C., Lefebvre, M., & Bol, D. (2021). “Learning Without Feedback: Fixed Random Learning Signals Allow for Feedforward Training of Deep Neural Networks” (DRTP). Frontiers in Neuroscience, 15, 629892. https://doi.org/10.3389/fnins.2021.629892

compile_graph(*args)

Compile the eligibility trace graph of the relationship between etrace weights, states and operators.

The compilation process includes:

• building the etrace graph

• separating the states

• initializing the etrace states

Parameters:

*args (Any) – The input arguments.

Return type:

None

update(x, y_target=None)

Call super().update(x) after stashing y_target for _get_update_aux.

The stash is read back synchronously by _get_update_aux() (called at the very start of the base class’s update(), before any custom_vjp tracing happens) and threaded from there into _true_update_fun as a genuine argument. This is required because outer transforms (e.g. brainstate.transform.grad) may stage the forward trace and only invoke the custom_vjp fwd/bwd rules after this Python-level call has already returned – by then any instance-attribute stash read directly inside _update_fn_fwd or _update_fn_bwd would already see None. Threading the value as a real traced argument sidesteps that timing problem entirely.

Return type:

Any