pp_prop

pp_prop#

class braintrace.pp_prop#

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

pp_prop is the canonical name for the input-output-dimension eligibility trace algorithm implemented by IODimVjpAlgorithm. It computes the gradients of the weights with the diagonal approximation and the input-output dimensional complexity introduced by Wang et al. [1], based on the RTRL construction of Williams and Zipser [2].

This subclass inherits all behavior from IODimVjpAlgorithm without modification; it exists to provide the canonical pp_prop name.

Parameters:

  • model (brainstate.nn.Module) – Recurrent model whose ETP-routed parameters are trained online.

  • decay_or_rank (float, int, or tuple of two floats or ints) – Exponential-smoothing decay, integer rank parameterization, or separate input/output-side settings.

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

  • vjp_method ({‘single-step’, ‘multi-step’}, optional) – VJP window used to compute the learning signal.

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

  • control_flow (ControlFlowPolicy, optional) – Control-flow canonicalization policy used during graph compilation.

  • config (ETraceConfig, optional) – Learning-rule coordinates. None derives the pp-prop preset from decay_or_rank.

  • random_feedback_key (jax.Array, optional) – Key for fixed random-feedback projections requested by config.

See also

IODimVjpAlgorithm

Input/output-factorized engine implementing this preset.

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}\]

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

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