D-RTRL: parameter-dimensional online traces

D-RTRL propagates an eligibility trace forward with the model rather than retaining a sequence-wide reverse-mode graph. This chapter identifies the approximation, trains a compact recurrent model, and limits each conclusion to what the example actually tests.

1. Principle and retained terms

For parameter \(\theta\) and hidden state \(h_t\), RTRL propagates

\[G_t = \frac{\partial h_t}{\partial \theta} = \frac{\partial h_t}{\partial h_{t-1}}G_{t-1} + \frac{\partial h_t}{\partial \theta}.\]

D-RTRL stores a parameter-shaped trace but retains only the per-position diagonal/block-diagonal part of the hidden-to-hidden Jacobian. The instantaneous parameter term remains. The online gradient estimate contracts the current learning signal with that trace. It therefore avoids BPTT activation storage, but it is not generally element-wise equal to BPTT when hidden positions mix.

2. Applicability

Use D-RTRL when a recurrent ETP model needs forward-only updates and its per-parameter trace is affordable, especially when cross-position hidden coupling is limited or an explicitly diagonal approximation is acceptable.

Do not use this example to claim universal BPTT equivalence, biological plausibility, or favorable scaling for very wide recurrent layers. A decreasing loss is an optimization smoke check, not a gradient oracle.

3. Setup

import brainstate
import braintools
import braintrace
import jax.numpy as jnp
import matplotlib.pyplot as plt

brainstate.random.seed(7)

4. A compact recurrent task

The MiniGRU supplies genuine recurrent state and ETP-aware parameterized operations. The fixed sequence keeps the optimization check reproducible; it does not remove the estimator’s approximation error.

class SequenceModel(brainstate.nn.Module):
    def __init__(self):
        super().__init__()
        self.rnn = braintrace.nn.MiniGRU(in_size=1, out_size=6)
        self.readout = braintrace.nn.Linear(6, 1)

    def update(self, x):
        return self.readout(self.rnn(x))


model = SequenceModel()
inputs = jnp.linspace(-1.0, 1.0, 12).reshape(12, 1, 1)
targets = 0.7 * inputs + 0.2

# The Linear readout is intentionally non-temporal; Section 5 explains
# the compiler diagnostic emitted for it.
learner = braintrace.compile(
    model,
    braintrace.D_RTRL,
    inputs[0],
    batch_size=1,
)
weights = model.states(brainstate.ParamState)
optimizer = braintools.optim.SGD(lr=0.08)
optimizer.register_trainable_weights(weights);

D:\BrainTrace\braintrace\_compiler\hid_param_op.py:969: UserWarning: ETP primitive etp_mm (weight=('readout', 'weight')) has no connected hidden states. It will be treated as a non-temporal parameter.
  _emit_no_relation_diag(

5. Online update and evidence

etrace_grad owns the temporal loop, advances the model exactly once per time step, and accumulates the estimated gradient. The final assertion is deliberately narrow: the fixed-seed objective should decrease.

def reset_sequence():
    brainstate.nn.reset_all_states(model, batch_size=1)
    learner.reset_state(batch_size=1)


def evaluate():
    reset_sequence()
    predictions = learner.etrace_evolve(inputs, return_outputs=True)
    return jnp.mean((predictions - targets) ** 2)


def local_loss(x, target):
    prediction = learner(x)
    return jnp.mean((prediction - target) ** 2)


def train_epoch(_):
    reset_sequence()
    # etrace_grad owns the loop, the accumulation and the reduction; local_loss
    # owns the model call. 'mean' divides by the total mask weight -- here T --
    # which is exactly the hand-written `grads / inputs.shape[0]` it replaces.
    grads, step_losses = learner.etrace_grad(
        inputs, targets, step_fn=local_loss,
        reduction='mean', return_value=True,
    )
    optimizer.update(grads)
    return step_losses.mean()


initial_loss = evaluate()
training_losses = brainstate.transform.for_loop(
    train_epoch, jnp.arange(25)
)
final_loss = evaluate()

print(f"initial loss: {float(initial_loss):.4f}")
for epoch in (0, 12, 24):
    print(f"epoch {epoch:2d} loss: {float(training_losses[epoch]):.4f}")
print(f"final loss: {float(final_loss):.4f}")
print("loss decreased:", bool(final_loss < initial_loss))

with plt.style.context("default"), plt.rc_context({
    "figure.facecolor": "white",
    "axes.facecolor": "white",
    "savefig.facecolor": "white",
}):
    fig, ax = plt.subplots(figsize=(7, 3.8))
    ax.plot(training_losses, color="#2563eb")
    ax.set(xlabel="Training epoch", ylabel="Mean online loss")
    ax.set_title("D-RTRL mini-GRU training loss")
    ax.grid(True, alpha=0.3)
    fig.tight_layout()

initial loss: 0.0802
epoch  0 loss: 0.0802
epoch 12 loss: 0.0100
epoch 24 loss: 0.0086
final loss: 0.0086
loss decreased: True

6. Interpretation and limits

• The loss curve supports that this estimator can provide a descent direction on this task; it does not establish equality with BPTT.

• A non-temporal readout may appear in compiler diagnostics because only paths that influence recurrent hidden state carry temporal eligibility traces.

• Memory follows participating parameter and hidden dimensions. pp-prop is the relevant next comparison when a factorized trace is required.

References and API

• Wang et al., “Model-agnostic linear-memory online learning in spiking neural networks,” Nature Communications (2026), doi:10.1038/s41467-026-68453-w.

• Williams and Zipser, “A Learning Algorithm for Continually Running Fully Recurrent Neural Networks,” Neural Computation (1989), doi:10.1162/neco.1989.1.2.270.

• API: braintrace.D_RTRL, braintrace.compile(), braintrace.SequenceDriverMixin.etrace_grad().