ThreeFactor

ThreeFactor#

class braintrace.ThreeFactor#

Reward-modulated eligibility-trace learning.

The coordinate is D_RTRL’s with learning_signal='modulatory': a per-parameter trace, diagonal recurrence scope, and a learning signal supplied by the caller rather than by reverse-AD.

\[\frac{\partial L}{\partial \theta} \;\leftarrow\; \sum_g \mathrm{expand}(m) \cdot \varepsilon_g\]

Replace, not multiply. The modulator is the learning signal; it does not scale dL/dh. Multiplying would give a four-factor rule, and would make the degenerate check – set the modulator to dL/dh and recover symmetric exactly – impossible to satisfy, leaving the axis with no coordinate at which it reduces to the rule it generalises.

One array, expanded to every group. A scalar reward is valid for any model whatever its HiddenGroup count. There is deliberately no per-group sequence spelling: binding the signal to the hidden-group decomposition, whose size is a property of the compiled graph rather than of the task, is what made OSTTP non-general. Expansion follows expand_modulator_to_group – shape driven, never group indexed.

Single-step only, and it raises otherwise. Under multi-step, _solve_weight_gradients adds the within-window reverse-AD gradient of the ETP parameters on top of the trace contraction, so replacing the boundary signal would leave that in-window half unmodulated: a hybrid that is part three-factor rule and part plain loss gradient. Single-step routes every ETP contribution through the replaced signal, and makes the modulator per step, which is what a neuromodulator is. A consequence worth stating: because the window is one step, update_schedule has nothing to schedule and stays out of this preset.

Parameters:

model (brainstate.nn.Module) – The one-step model.
name (str, optional) – Node name.
vjp_method (str, optional) – Must be 'single-step' (the default); anything else raises.
modulator (array_like or Quantity, optional) – The initial standing modulator, equivalent to assigning ETraceVjpAlgorithm.modulator after construction. A scalar, an array shaped like a group’s varshape, or an array broadcastable to (*varshape, num_state). Leaving it None is fine as long as one is supplied before the first update(); there is no fallback to symmetric.
fast_solve (bool, optional) – Whether registered closed-form kernels may be used. Default True.
trace_dtype (DTypeLike, optional) – Reduced trace precision, as in D_RTRL.
chunked_trace (bool, optional) – Whether to roll the trace in chunks. Default True.
control_flow (ControlFlowPolicy, optional) – Control-flow canonicalization policy.
snap_max_jacobian_elements (int, optional) – Passed through; unused at recurrence_scope='diagonal'.

Examples

>>> import brainstate, braintrace, jax.numpy as jnp
>>> class Net(brainstate.nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.w = brainstate.ParamState(0.1 * jnp.ones((4, 4)))
...         self.h = brainstate.HiddenState(jnp.zeros((1, 4)))
...     def update(self, x):
...         self.h.value = jnp.tanh(x + braintrace.matmul(self.h.value, self.w.value))
...         return self.h.value
>>> model = Net()
>>> brainstate.nn.init_all_states(model, batch_size=1)
>>> learner = braintrace.ThreeFactor(model)
>>> learner.compile_graph(jnp.zeros((1, 4)))
>>> learner.init_etrace_state()

Then drive it with a reward per step, either per call or as a standing value:

>>> out = learner.update(jnp.zeros((1, 4)), modulator=0.5)
>>> learner.modulator = -1.0        # standing, until reassigned

The keyword takes precedence for the call it appears on.

Over a sequence, the per-step reward is simply a second sequence: etrace_grad slices every sequence in lockstep and hands the slices to step_fn positionally, and step_fn – not the driver – owns the model call, so there is nowhere the modulator has to be threaded through.

>>> xs = jnp.zeros((10, 1, 4))
>>> rewards = jnp.linspace(-1.0, 1.0, 10)
>>> ys = jnp.zeros((10, 1, 4))
>>> def step_loss(x, reward, y):
...     return jnp.mean((learner.update(x, modulator=reward) - y) ** 2)
>>> grads, losses = learner.etrace_grad(
...     xs, rewards, ys, step_fn=step_loss, return_value=True)

Notes

Under single-step, every plain (non-ETP) parameter’s gradient is exactly zero (F-33) – not merely truncated. Since this preset is single-step by construction, a model whose parameters are partly plain will train only its ETP-routed parameters here. Route the parameters you intend to modulate through an ETP primitive (braintrace.matmul and friends).

The modulator is a genuine data dependency of the traced computation, not a lazily-read attribute: it is read synchronously at the top of update() (see _get_update_aux), because an outer transform may stage the forward trace and invoke the backward rule only after update() has returned.

ThreeFactor

Contents

ThreeFactor#

Modeling

Infrastructure

Compilation

Vibe coding