SnAp

SnAp#

class braintrace.SnAp#

Sparse n-step Approximation of RTRL.

SnAp-\(n\) [1] sparsifies the full influence matrix carried by RTRL [2]. In the unapproximated construction, RTRL carries \(\mathbf{J}^t_{q,\theta} = \partial h^t_q / \partial \theta\) and rolls it with

\[\mathbf{J}^t = \mathbf{D}^t \mathbf{J}^{t-1} + \mathbf{J}_f^t , \qquad \mathbf{D}^t = \frac{\partial \mathbf{h}^t}{\partial \mathbf{h}^{t-1}}, \qquad (\mathbf{J}_f^t)_{q,\theta} = \frac{\partial h^t_q}{\partial \theta} \bigg|_{\mathbf{h}^{t-1}} ,\]

which costs \(O(P \cdot |\theta|)\) memory. SnAp-\(n\) keeps the same recursion but masks \(\mathbf{J}\) to the sparsity pattern

\[\mathcal{S}_n \;=\; \operatorname{nz}\!\Bigl( \textstyle\bigvee_{k < n} \mathbf{A}^k \Bigr),\]

where \(\mathbf{A}\) is the one-step position-adjacency of the hidden group (position \(p\) influences position \(q\) in one step). Only the retained entries are stored and rolled, so the trace’s trailing state axis widens from \(S\) (states per position) to \(M = K \cdot S\), with \(K = |\mathcal{N}_n(p)|\) the largest retained neighbourhood.

Parameters:

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

  • n (int, default 2) – The SnAp order — how many propagation steps of the instantaneous term are kept. Must be an integer >= 1. n = 1 canonicalises to recurrence_scope='coupled' (OSTLRecurrent’s coordinate); any n at or above a group’s diameter saturates to full within-group RTRL. There is deliberately no “infinity” spelling: saturation is a property of the model, not of the vocabulary.

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

  • vjp_method ({‘single-step’, ‘multi-step’}, default ‘single-step’) – Execution option, forwarded verbatim. The finite-window oracle (chunked_online_param_gradients) needs 'multi-step'.

  • fast_solve (bool, default True) – Execution option, forwarded verbatim: whether to use the per-primitive fast contraction instead of the legacy vmap path.

  • snap_max_jacobian_elements (int, optional) – Ceiling on each hidden group’s widened block Jacobian, P * (K * S) ** 2 elements. Raising it is how a deliberately large neighbourhood is admitted; the default rejects roughly half a gigabyte per operator. Forwarded to the compiler.

  • **kwargs (Any) – Additional options forwarded to ParamDimVjpAlgorithm, including trace_dtype, chunked_trace, control_flow, config, and random_feedback_key.

Variables:

n (int) – The requested order, as passed (before canonicalisation).

Notes

The pattern is computed, not assumed. The compiler analyses the hidden group’s transition jaxpr and derives \(\mathbf{A}\) from the recurrent mixing primitive it finds. Two primitive families yield a precise pattern: dense (etp_mm / etp_mv, all-to-all) and sparse (etp_sp_mm / etp_sp_mv, the structural pattern of sparse_mat, transposed). For every other primitive, more than one mixing equation, any control flow around it, or a non-position-preserving tail, the analysis falls back to the all-to-all pattern and emits a DiagnosticKind.SNAP_PATTERN_CONSERVATIVE diagnostic. A conservative pattern is always correct — it retains a superset of the true influence — but it costs saturated memory, so the diagnostic is worth reading.

The scale is within-group. \(\mathcal{N}_n(p)\) never leaves the hidden group that owns \(p\); cross-group influence is not represented by the per-parameter trace at any coordinate. On a model with exactly one hidden group and an elementwise y -> hidden tail, saturation therefore equals BPTT; on a multi-group model it equals full RTRL within each group, which is not BPTT.

Memory. The trace grows linearly in \(K\); the widened transition operator \(\mathbf{D}_g\) costs \(P \cdot K^2 \cdot S^2\) on top, which is the term that actually bounds usable \(n\) on large groups.

References

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)
>>> learner = braintrace.compile(model, braintrace.SnAp, x0, n=2)
>>> 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)
SnAp.__init__(model, n=2, name=None, vjp_method='single-step', fast_solve=True, snap_max_jacobian_elements=16777216, **kwargs)#