Core Concepts

Core Concepts#

Welcome to BrainTrace! This notebook introduces the core concepts you need to understand before using the library for online learning in recurrent and spiking neural networks.

BrainTrace is built on JAX and brainstate, providing memory-efficient online learning through eligibility trace propagation.

1. What is Online Learning?#

Training recurrent neural networks (RNNs) typically relies on Backpropagation Through Time (BPTT). BPTT unrolls the full computation graph over all time steps before computing gradients:

  • BPTT stores the entire computation graph across \(T\) time steps, requiring \(O(T)\) memory.

  • As sequence length grows, memory usage becomes a bottleneck.

Online learning takes a different approach:

  • Weights are updated at each time step using eligibility traces that summarize the gradient history.

  • Memory cost is \(O(1)\) per time step (independent of sequence length).

  • Eligibility traces accumulate the information needed for gradient computation incrementally.

BrainTrace implements online learning via JAX custom primitives. Instead of relying on string-matching or special parameter wrappers, BrainTrace identifies which operations participate in online learning by their primitive type at the JAX IR level. This gives a clean, composable, and JIT-friendly design.

2. Architecture Overview#

BrainTrace is organized as a 4-layer system. Each layer builds on the one below it:

+--------------------------------------------------------------+
|  Algorithms    D_RTRL / pp_prop / SnAp / UORO / ThreeFactor   |
|                DNI / EProp / OSTL... (or an ETraceConfig)     |
|                trace update + custom_vjp for jax.grad         |
+--------------------------------------------------------------+
|  Executor      ETraceGraphExecutor                            |
|                forward pass + Jacobian computation            |
+--------------------------------------------------------------+
|  Compiler      compile_etrace_graph()                         |
|                jaxpr walk -> find primitives -> connect to     |
|                hidden states                                  |
+--------------------------------------------------------------+
|  Primitives    braintrace.matmul / element_wise / conv        |
|  & Functions   JAX custom primitives (thin markers)           |
+--------------------------------------------------------------+

How it works:

  1. Primitives & Functions (bottom layer): You call braintrace.matmul(x, w) in your model. Under the hood, this binds a JAX custom primitive that acts as a marker — the actual computation is standard JAX (x @ w).

  2. Compiler: When you call braintrace.compile(...), BrainTrace walks the JAX intermediate representation (jaxpr), finds all ETP primitives, and connects each one to its associated hidden states and parameters. (compile_graph() is the lower-level entry point this uses.)

  3. Executor: During the forward pass, the executor computes the model output and the Jacobians needed for eligibility trace updates.

  4. Algorithms (top layer): the chosen algorithm uses the executor outputs to maintain eligibility traces and provide correct gradients via custom_vjp. The named algorithms are thin factories over coordinates in the ETraceConfig axis space, so a rule with no name is as constructible as one with a name.

3. Key Concept: Primitive-Based Parameter Selection#

The central design idea of BrainTrace is that the operation you use determines whether a parameter participates in online learning:

What you write

Effect

braintrace.matmul(x, w)

w is included in online learning (eligibility traces are maintained)

x @ w (regular JAX matmul)

w is excluded from online learning (only instantaneous gradients)

There is no need for special parameter classes. All weights are plain brainstate.ParamState. The choice of operation is what matters.

Here is a concrete example:

import os
os.environ.setdefault("JAX_PLATFORMS", "cpu")

import jax
import jax.numpy as jnp
import brainstate
import braintrace

class SimpleRNN(brainstate.nn.Module):
    def __init__(self, n_in, n_rec, n_out):
        super().__init__()
        self.w_in = brainstate.ParamState(brainstate.random.randn(n_in, n_rec) * 0.01)
        self.w_rec = brainstate.ParamState(brainstate.random.randn(n_rec, n_rec) * 0.01)
        self.w_out = brainstate.ParamState(brainstate.random.randn(n_rec, n_out) * 0.01)
        self.h = brainstate.ShortTermState(jnp.zeros(n_rec))

    def update(self, x):
        # Regular matmul: w_in excluded from online learning
        inp = x @ self.w_in.value

        # ETP matmul: w_rec included in online learning
        rec = braintrace.matmul(self.h.value, self.w_rec.value)

        self.h.value = jax.nn.tanh(inp + rec)

        # Regular matmul: w_out excluded from online learning
        return self.h.value @ self.w_out.value

In the model above:

  • w_in and w_out use standard x @ w — they receive only instantaneous gradients (no temporal credit assignment through eligibility traces).

  • w_rec uses braintrace.matmul(h, w_rec) — the compiler will automatically maintain eligibility traces for this weight, enabling gradient computation that accounts for temporal dependencies.

All three weights are the same type (brainstate.ParamState). The operation is the only difference.

4. Using braintrace.nn Modules#

While you can use primitives directly (as shown above), BrainTrace provides pre-built layers in the braintrace.nn module that already use ETP primitives internally. These are drop-in replacements for standard brainstate.nn layers:

Module

Description

braintrace.nn.Linear

Dense linear layer using braintrace.matmul

braintrace.nn.SignedWLinear

Linear layer with sign-constrained weights (E/I networks)

braintrace.nn.ScaledWSLinear

Weight-standardized linear layer

braintrace.nn.SparseLinear

Linear layer with sparse connectivity (uses sparse_matmul)

braintrace.nn.LoRA

Low-rank adapter layer (uses lora_matmul)

braintrace.nn.Conv1d / Conv2d / Conv3d

Convolutional layers using braintrace.conv

braintrace.nn.GRUCell / LSTMCell / ValinaRNNCell

Recurrent cells with ETP-aware gates

braintrace.nn.LeakyRateReadout

Rate-coded SNN readout

braintrace.nn.BatchNorm1d / LayerNorm

Normalisation layers

For the matching low-level API, the user-facing primitive functions are:

  • braintrace.matmul(x, w, bias=None) – dense matrix multiplication

  • braintrace.element_wise(weight, weight_fn=...) – element-wise weight ops (gating, learnable thresholds)

  • braintrace.conv(x, kernel, bias, ...) – convolution

  • braintrace.sparse_matmul(x, weight_data, *, sparse_mat, bias=None) – sparse matmul

  • braintrace.lora_matmul(x, B, A, *, alpha=1.0, bias=None) – LoRA decomposition

Use the braintrace.nn layers when you can; reach for the primitive functions when you need a custom layer that participates in online learning.

class GRUNet(brainstate.nn.Module):
    def __init__(self, n_in, n_rec, n_out):
        super().__init__()
        self.rnn = braintrace.nn.GRUCell(n_in, n_rec)
        self.readout = braintrace.nn.Linear(n_rec, n_out)

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

The GRUCell internally uses braintrace.matmul for its weight operations, so all its recurrent parameters automatically participate in online learning. The Linear readout also uses ETP primitives, but the compiler will detect that it is not connected to any hidden state and handle it appropriately.

5. Online Learning in 3 Steps#

Using BrainTrace for online learning follows a simple three-step workflow:

  1. Define the model using braintrace.nn modules or manual ETP primitives.

  2. Compile with braintrace.compile — one call initialises states, builds the eligibility-trace graph, and returns a ready learner.

  3. Train with standard JAX gradient computation — eligibility traces are updated inside the wrapped model call.

Here is the complete workflow:

# Step 1: Define a reproducible model and one supervised example
brainstate.random.seed(11)
model = GRUNet(3, 8, 1)
input_sample = jnp.array([[0.2, -0.1, 0.3]])
target = jnp.array([[0.5]])

# Step 2: Compile — initialises states, builds the trace graph, returns a ready learner.
# The example input carries the batch axis: shape (batch_size, n_in) = (1, 3).
learner = braintrace.compile(model, braintrace.D_RTRL, input_sample, batch_size=1)

# Step 3: Differentiate, update, and compare clean before/after states
import braintools

weights = model.states(brainstate.ParamState)
optimizer = braintools.optim.SGD(lr=0.1)
optimizer.register_trainable_weights(weights)

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

input_sequence = input_sample[None, ...]  # (time=1, batch=1, features=3)
target_sequence = target[None, ...]

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

def evaluate_online():
    reset_online_state()
    predictions = learner.etrace_evolve(input_sequence, return_outputs=True)
    prediction = predictions[-1]
    return jnp.mean((prediction - target) ** 2), prediction

initial_loss, initial_prediction = evaluate_online()
parameters_before = jax.tree.map(lambda value: value.copy(), weights.to_dict_values())
reset_online_state()
grads, _, _ = learner.etrace_grad(
    input_sequence, target_sequence, step_fn=step_loss,
    has_aux=True, return_value=True,
)
optimizer.update(grads)
final_loss, final_prediction = evaluate_online()

parameter_differences = jax.tree.map(
    lambda after, before: after - before,
    weights.to_dict_values(),
    parameters_before,
)
parameter_change = jnp.sqrt(sum(
    jnp.sum(delta ** 2) for delta in jax.tree.leaves(parameter_differences)
))

print('initial loss:', float(initial_loss))
print('final loss:', float(final_loss))
print('initial prediction:', initial_prediction)
print('final prediction:', final_prediction)
print('parameter change:', float(parameter_change))

initial loss: 0.24410992860794067
final loss: 0.2132834494113922
initial prediction: [[0.00592517]]
final prediction: [[0.03817381]]
parameter change: 0.05644570291042328

What happens under the hood:

  • braintrace.compile(model, braintrace.D_RTRL, x0, batch_size=1) initialises all model states, traces the model through JAX, identifies all ETP primitives, and builds the eligibility trace computation graph — all in one call.

  • learner.etrace_evolve(...) drives a sequence without a loss gradient while advancing both hidden state and eligibility traces.

  • learner.etrace_grad(..., step_fn=step_loss) drives the supervised sequence and accumulates online gradients; step_loss owns the single learner call for each step.

Code comparison with BPTT#

A fair comparison holds model structure, initialization, inputs, targets, loss, and reset policy fixed. The online path differentiates one step at a time and accumulates gradients during the forward drive, which learner.etrace_grad performs. The BPTT path first constructs the full sequence loss and differentiates through the complete for_loop.

Online learning

BPTT

learner.etrace_grad(..., step_fn=step_loss) drives the sequence

grad(sequence_loss) outside brainstate.transform.for_loop

Eligibility state carries temporal information forward

Reverse mode traverses the unrolled sequence

No stored trajectory for the online update

Stores or rematerializes the sequence trajectory

Matching forward losses establish that the comparison starts from the same model and state. They do not imply that every online estimator must equal the BPTT gradient: equality depends on the algorithm, VJP mode, comparison window, and the regime guaranteed by its mathematics.

brainstate.random.seed(23)
online_model = GRUNet(3, 8, 1)
brainstate.random.seed(23)
bptt_model = GRUNet(3, 8, 1)

sequence_inputs = jnp.array([
    [[0.2, -0.1, 0.3]],
    [[0.1, 0.0, -0.2]],
    [[-0.3, 0.2, 0.1]],
])
sequence_targets = jnp.array([[[0.5]], [[0.1]], [[-0.2]]])
online_learner = braintrace.compile(
    online_model, braintrace.D_RTRL, sequence_inputs[0], batch_size=1
)
bptt_weights = bptt_model.states(brainstate.ParamState)
brainstate.nn.init_all_states(bptt_model, batch_size=1)

def online_step_loss(x, y):
    prediction = online_learner(x)
    return jnp.mean((prediction - y) ** 2)

def online_sequence_gradient(inputs, targets):
    brainstate.nn.reset_all_states(online_model, batch_size=1)
    online_learner.reset_state(batch_size=1)
    # The scan this replaces accumulated the per-step gradients without
    # dividing, which is exactly reduction='sum'. return_value=True keeps the
    # per-step losses the comparison below prints.
    return online_learner.etrace_grad(
        inputs, targets, step_fn=online_step_loss,
        reduction='sum', return_value=True,
    )

def bptt_sequence_loss(inputs, targets):
    brainstate.nn.reset_all_states(bptt_model, batch_size=1)
    predictions = brainstate.transform.for_loop(bptt_model, inputs)
    return jnp.mean((predictions - targets) ** 2)

online_grads, online_step_losses = online_sequence_gradient(
    sequence_inputs, sequence_targets
)
bptt_grads, bptt_loss = brainstate.transform.grad(
    bptt_sequence_loss, bptt_weights, return_value=True
)(sequence_inputs, sequence_targets)
print('online forward loss:', float(online_step_losses.mean()))
print('BPTT forward loss:', float(bptt_loss))

online forward loss: 0.10689683258533478
BPTT forward loss: 0.10689682513475418

6. Available Algorithms#

The table lists the concrete algorithms in the current public API. ES_D_RTRL is a compatibility alias of pp_prop, not a separate method.

Algorithm

Applicable scope

D_RTRL

Parameter-dimensional online gradients for recurrent models; diagonal recurrence is the default scope.

pp_prop

Input/output-factorized traces for memory-constrained recurrent or spiking models; ES_D_RTRL is its compatibility alias.

EProp

SNN learning with symmetric or random-feedback learning signals and an optional kappa trace filter.

OSTLRecurrent

Recurrent OSTL coordinate retaining coupled within-position recurrence.

OSTLFeedforward

Feedforward/without-H OSTL coordinate using an input/output-factorized trace.

SnAp

Sparse n-step RTRL approximation when the recurrent position graph has an exploitable sparsity pattern.

UORO

Unbiased rank-1 random projection of the influence matrix: \(O(P)\) memory, higher variance. Built on RandomProjectionVjpAlgorithm.

ThreeFactor

Modulated learning: an external scalar (or per-group) modulator replaces the backpropagated error. Requires vjp_method='single-step'.

DNI

Decoupled neural interfaces: a learned SyntheticGradient predicts the future loss gradient. Train it with train_synthetic_gradient().

Every one of these is a preset over ETraceConfig, which describes a learning rule as a point in a six-axis space (trace_factorization, temporal_recursion, recurrence_scope, learning_signal, trace_filter, update_schedule). braintrace.compile accepts a config wherever it accepts a name, so you are not restricted to the rows above.

7. Summary#

Here is a quick recap of the core concepts:

Concept

Description

Online learning

Update weights at each time step using eligibility traces, achieving \(O(1)\) memory per step

ETP primitives

braintrace.matmul, braintrace.element_wise, braintrace.conv — JAX custom primitives that mark operations for online learning

Primitive-based selection

Use an ETP primitive to include a weight; use regular JAX ops to exclude it

braintrace.nn

Pre-built layers (Linear, GRUCell, LSTMCell, Conv) that use ETP primitives internally

braintrace.compile

braintrace.compile(model, algo, x0, batch_size=B) — the one-call entry point that initialises states, builds the trace graph, and returns a ready learner

Sequence drivers

learner.etrace_grad(...) accumulates online gradients over a sequence; learner.etrace_evolve(...) advances states and traces without gradients — neither needs a hand-written scan

Algorithm presets

D_RTRL, pp_prop/ES_D_RTRL, EProp, OSTLRecurrent, OSTLFeedforward, SnAp, UORO, ThreeFactor, DNI — see section 6

braintrace.ETraceConfig

The six-axis space every preset above is a point in; pass one to compile for a rule with no preset

8. Next Steps#

Now that you understand the core concepts, explore the following tutorials: