EProp

EProp#

class braintrace.EProp(model, feedback='symmetric', kappa_filter_decay=0.0, random_feedback_key=None, name=None, vjp_method='single-step', fast_solve=True, **kwargs)#

Eligibility Propagation (e-prop) for recurrent spiking networks.

E-prop approximates the gradient of a loss \(\mathcal{L}\) with respect to a recurrent weight \(W_{ji}\) by the product of a local eligibility trace and a global learning signal, dropping the temporally non-local terms of BPTT:

\[\frac{d\mathcal{L}}{dW_{ji}} = \sum_t L_j^t \, \bar{e}_{ji}^t ,\]

where

\[e_{ji}^t = \frac{\partial h_j^t}{\partial W_{ji}} \approx D_j^t \, e_{ji}^{t-1} + \big[\operatorname{diag}(D_{f,j}^t)\big]\, x_i^t , \qquad \bar{e}_{ji}^t = \kappa\,\bar{e}_{ji}^{t-1} + e_{ji}^t .\]

Here \(h_j^t\) is the hidden state of neuron \(j\) at time \(t\), \(x_i^t\) the presynaptic input, \(D_j^t\) the hidden-to-hidden (recurrent) Jacobian diagonal, \(D_{f,j}^t\) the state-to-output Jacobian, and \(\kappa \in [0, 1)\) the eligibility-trace low-pass factor. The learning signal is

\[\begin{split}L_j^t = \begin{cases} \ell_j^t = \dfrac{\partial \mathcal{L}}{\partial h_j^t} & \text{(symmetric feedback, standard backprop through readout)} \\[2ex] \big(B\,\hat\ell^t\big)_j, \quad \hat\ell^t = \dfrac{\ell^t}{\lVert \ell^t \rVert + \varepsilon} & \text{(random feedback: a fixed random projection } B\text{)} . \end{cases}\end{split}\]

How it works. The eligibility trace \(e_{ji}^t\) is exactly the per-parameter trace maintained by D_RTRL; it depends only on quantities local to the synapse and is updated forward in time. The learning signal \(L_j^t\) is broadcast from the readout. E-prop is therefore online (no backward pass through time) and uses memory linear in the number of parameters. With kappa_filter_decay > 0 each weight’s eligibility trace is additionally low-pass filtered (elementwise over the trailing per-hidden-state axis, so multi-state HiddenGroups are filtered per state with no cross-state mixing); with feedback='random' the symmetric readout gradient \(\ell^t\) is L2-normalized (removing its dependence on the magnitude of the real readout weights, since reverse-AD only ever exposes \(\ell^t = W_\mathrm{out}^\top \delta^t\), never the pre-readout error \(\delta^t\) itself) and then projected through a frozen random matrix \(B\), removing the biologically implausible weight-transport requirement. Normalization removes the scale dependence on \(W_\mathrm{out}\); a residual directional dependence on the readout weights remains, since full direction- independence would require the pre-readout error \(\delta^t\) explicitly, which the current hook contract does not provide.

Parameters:

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

  • feedback (str) – 'symmetric' uses reverse-AD’s \(\partial \mathcal{L}/\partial h\) (standard backprop through the readout). 'random' replaces the readout gradient with a frozen random projection of its L2-normalized direction (requires random_feedback_key). The projection matrix is square (hidden-dim × hidden-dim): E-prop’s hooks only see \(\partial\mathcal L/\partial h\), which has no visibility into a separate readout layer’s width, so feedback='random' assumes a single, direct readout whose output dimensionality equals the HiddenGroup’s own width.

  • kappa_filter_decay (float) – Eligibility-trace low-pass factor \(\kappa\) (see \(\bar{e}_{ji}^t\) above). If > 0, each trainable weight’s raw eligibility trace is filtered every step, per hidden-state channel. 0 disables filtering (the algorithm then reduces exactly to D_RTRL).

  • random_feedback_key (Array | None) – Seed for the random-feedback matrices. Required when feedback='random'; ignored otherwise.

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

  • vjp_method (str) – Forwarded verbatim to D_RTRL.

  • fast_solve (bool) – Forwarded verbatim to D_RTRL.

Raises:

ValueError – If feedback is not one of {'symmetric', 'random'}, or if feedback='random' is given without random_feedback_key.

Examples

>>> import brainstate
>>> import braintrace
>>>
>>> class RSNN(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 = RSNN()
>>> x0 = brainstate.random.randn(1)
>>> # one call: initialise states, build the trace graph, return a learner
>>> learner = braintrace.compile(model, braintrace.EProp, x0, kappa_filter_decay=0.9)
>>> y = learner(x0)             # forward pass + eligibility-trace update

References

init_etrace_state(*args, **kwargs)[source]#

Initialize the eligibility trace states of the etrace algorithm.

This method is needed after compiling the etrace graph. See compile_graph() for the details.

Return type:

None

reset_state(batch_size=None, **kwargs)[source]#

Reset the eligibility trace states.

Parameters:

batch_size (int | None) – The batch size used to reshape the reset trace states. Default None.

Return type:

None