IODimVjpAlgorithm

class braintrace.IODimVjpAlgorithm

Online gradient algorithm with diagonal approximation and input-output-dimension complexity.

This algorithm computes the gradients of the weights with the diagonal approximation and the input-output dimensional complexity. It implements the input/output-factorized estimator of Wang et al. [1] on the RTRL foundation of Williams and Zipser [2].

Parameters:

• model (brainstate.nn.Module) – The model function, which receives the input arguments and returns the model output.

• decay_or_rank (float, int, or tuple of two floats or ints) – Parameterization of the exponential smoothing factor. A float in \((0, 1)\) is used directly as the decay. A positive integer \(r\) is converted to \(\alpha=(r-1)/(r+1)\). The integer form does not allocate \(r\) separate trace factors; both forms use the same input/output-factorized trace structure. A two-item tuple configures the input- and output-side factors separately.

• name (str, optional) – Name of the eligibility-trace algorithm.

• vjp_method (str, optional) – The method for computing the VJP. It should be either "single-step" or "multi-step".

  • "single-step": the VJP is computed at the current time step, i.e., \(\partial L^t/\partial h^t\).

  • "multi-step": the VJP is computed at multiple time steps, i.e., \(\partial L^t/\partial h^{t-k}\), where \(k\) is determined by the data input.

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

• control_flow (ControlFlowPolicy, optional) – Policy governing control-flow canonicalization (cond if-conversion, scan unrolling, structured scan descent, …) during graph compilation. None (default) uses ControlFlowPolicy().

• config (ETraceConfig, optional) – Learning-rule coordinates. None uses the input/output-factorized preset derived from decay_or_rank.

• random_feedback_key (jax.Array, optional) – Key used to initialize fixed random-feedback projections when the selected config requests them.

Notes

The learning rule is

\[\begin{split}\begin{aligned} & \boldsymbol{\epsilon}^t \approx \boldsymbol{\epsilon}_{\mathbf{f}}^t \otimes \boldsymbol{\epsilon}_{\mathbf{x}}^t \\ & \boldsymbol{\epsilon}_{\mathbf{x}}^t=\alpha \boldsymbol{\epsilon}_{\mathbf{x}}^{t-1}+\mathbf{x}^t \\ & \boldsymbol{\epsilon}_{\mathbf{f}}^t=\alpha \operatorname{diag}\left(\mathbf{D}^t\right) \circ \boldsymbol{\epsilon}_{\mathbf{f}}^{t-1}+(1-\alpha) \operatorname{diag}\left(\mathbf{D}_f^t\right) \\ & \nabla_{\boldsymbol{\theta}} \mathcal{L}=\sum_{t^{\prime} \in \mathcal{T}} \frac{\partial \mathcal{L}^{t^{\prime}}}{\partial \mathbf{h}^{t^{\prime}}} \circ \boldsymbol{\epsilon}^{t^{\prime}} \end{aligned}\end{split}\]

where \(\boldsymbol{\epsilon}_{\mathbf{x}}^t\) is the input-side trace, \(\boldsymbol{\epsilon}_{\mathbf{f}}^t\) the output-side trace, \(\alpha\) the exponential-smoothing factor, \(\mathbf{D}^t\) the hidden-to-hidden Jacobian, \(\mathbf{D}_f^t\) the state-to-output Jacobian, and \(\mathbf{x}^t\) the presynaptic input.

\(\mathbf{D}_f^t\) is read off by ETraceVjpGraphExecutor._compute_hid2weight_jacobian from a single all-ones-tangent jax.jvp of the y -> hidden map; see that method’s docstring for when this is exact (elementwise maps) versus an approximation (non-elementwise maps, e.g. a normalization layer between the weight op and the neuron) — the same approximation is shared with ParamDimVjpAlgorithm (D_RTRL).

The full per-parameter D-RTRL trace \(\boldsymbol{\epsilon}^t \in \mathbb{R}^{I\times O}\) is approximated by the outer product of two exponentially-smoothed vectors — one over the input dimension and one over the output dimension. Storing the two factors instead of the full matrix drops the memory from \(O(I\cdot O)\) to \(O(I+O)\) per layer. The decay \(\alpha\) (or its integer parameterization) controls how much temporal history the factored trace retains; the bias of the exponential estimator is corrected at solve time.

This algorithm has \(O(BI+BO)\) memory complexity and \(O(BIO)\) computational complexity, where \(I\) and \(O\) are the number of input and output dimensions, and \(B\) the batch size. In particular, for a linear transformation layer, the weight gradients are computed with \(O(Bn)\) memory complexity and \(O(Bn^2)\) computational complexity, where \(n\) is the number of hidden dimensions.

For more details, please see the ES-D-RTRL algorithm presented in our manuscript.

Examples

>>> import brainstate
>>> import braintrace
>>> import jax.numpy as jnp
>>>
>>> class RNN(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 = RNN()
>>> x0 = brainstate.random.randn(1)
>>> # one call: initialise states, build the trace graph, return a learner
>>> learner = braintrace.compile(model, braintrace.pp_prop, x0, decay_or_rank=0.9)  # or rank: decay_or_rank=19
>>> y = learner(x0)             # forward pass + eligibility-trace update
>>>
>>> # 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

[1]

Wang, C., Dong, X., Ji, Z., Xiao, M., Jiang, J., Liu, X., Huan, Y., & Wu, S. (2026). “Model-agnostic linear-memory online learning in spiking neural networks.” Nature Communications. https://doi.org/10.1038/s41467-026-68453-w (preprint: bioRxiv 2024.09.24.614728)

[2]

Williams, R. J., & Zipser, D. (1989). “A Learning Algorithm for Continually Running Fully Recurrent Neural Networks” (RTRL). Neural Computation, 1(2), 270-280. https://doi.org/10.1162/neco.1989.1.2.270

property decay

The shared exponential-smoothing decay factor.

Returns:

float – The single decay, when both sides carry the same one — which is the case for every shipped preset.

Raises:

AttributeError – If the two sides differ, in which case there is no single decay to return; read decay_x / decay_f instead.

decay_f

The f-side (Jacobian-propagated output trace) decay factor. Also the one the warm-up bias correction is indexed by.

decay_x

The x-side (presynaptic input trace) exponential-smoothing decay factor.

get_etrace_of(weight)

Get the eligibility trace of the given weight.

Parameters:

weight (brainstate.ParamState or Path) – The weight whose eligibility trace is requested, given either as a brainstate.ParamState instance or as its path in the model.

Returns:

• etrace_xs (dict) – The input-side eligibility traces keyed by the weight-input variable.

• etrace_dfs (dict) – The output-side eligibility traces keyed by (y_var, hidden-group index).

Raises:

ValueError – If no eligibility trace is found for the given weight.

init_etrace_state(*args, **kwargs)

Initialize the eligibility trace states of the etrace algorithm.

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

reset_state(batch_size=None, **kwargs)

Reset the eligibility trace states.

Parameters:

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

IODimVjpAlgorithm.__init__(model, decay_or_rank, name=None, vjp_method='single-step', fast_solve=True, control_flow=None, config=None, random_feedback_key=None)