pp-prop: input/output-factorized online traces

pp-prop: input/output-factorized online traces#

pp-prop is the public name of BrainTrace’s input/output-factorized estimator (ES_D_RTRL remains a compatibility alias). This example uses a recurrent spiking model because a matched MiniGRU loss curve would not demonstrate the method’s intended linear-memory SNN regime.

1. Principle and approximation#

Instead of storing a parameter-shaped eligibility tensor, pp-prop maintains an input-side factor \(x_t\) and an output/hidden-side factor \(f_t\). Their smoothed outer-product contraction estimates the parameter trace:

\[\widehat{G}_t \approx \bar{x}_t \otimes \bar{f}_t,\qquad \bar{x}_t=\rho\bar{x}_{t-1}+(1-\rho)x_t.\]

The analogous recursion is applied to the output-side factor. This changes the memory dependence, but discards correlations that cannot be represented by the factorization. The decay is part of the estimator, not merely an optimizer hyperparameter.

2. Applicability#

Use pp-prop when a recurrent SNN needs an input/output-factorized trace and the full parameter-dimensional trace is too costly.

Avoid unqualified claims when strong cross-neuron correlations dominate, the decay has not been validated, or exact BPTT gradients are required. Windowed etrace_grad is also inappropriate for the current IO-factorized bias correction; drive one time step at a time as below.

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

brainstate.random.seed(11)

3. Minimal recurrent spiking model#

Each call computes one membrane update and one surrogate spike. The spike is then used once by the recurrent path and once by the readout. This one-call contract matters: calling a neuron twice would advance its state twice while etrace_grad believes only one logical time step elapsed.

class RecurrentLIF(brainstate.nn.Module):
    def __init__(self, n_in=3, n_rec=8, n_out=1):
        super().__init__()
        self.n_rec = n_rec
        self.w_in = brainstate.ParamState(
            0.35 * brainstate.random.randn(n_in, n_rec)
        )
        self.w_rec = brainstate.ParamState(
            0.12 * brainstate.random.randn(n_rec, n_rec)
        )
        self.w_out = brainstate.ParamState(
            0.2 * brainstate.random.randn(n_rec, n_out)
        )
        self.surrogate = braintools.surrogate.ReluGrad()

    def init_state(self, **kwargs):
        self.v = brainstate.HiddenState(jnp.zeros(self.n_rec))
        self.spike = brainstate.HiddenState(jnp.zeros(self.n_rec))

    def reset_state(self, **kwargs):
        self.v.value = jnp.zeros_like(self.v.value)
        self.spike.value = jnp.zeros_like(self.spike.value)

    def update(self, x):
        current = (
            braintrace.matmul(x, self.w_in.value)
            + braintrace.matmul(self.spike.value, self.w_rec.value)
        )
        voltage = 0.85 * self.v.value + current - self.spike.value
        spike = self.surrogate(voltage - 0.5)
        self.v.value = voltage
        self.spike.value = spike
        return voltage @ self.w_out.value

4. Compile one time step and align the objective#

The compilation sample is inputs[0], one feature vector, not the whole sequence. A warm-up mask gates the loss while still advancing both neuron and eligibility states. Evaluation uses the same mask and per-step squared error as training.

n_steps = 24
inputs = brainstate.random.bernoulli(
    0.3, size=(n_steps, 3)
).astype(jnp.float32)
targets = jnp.full((n_steps, 1), 0.4)
loss_mask = (jnp.arange(n_steps) >= 6).astype(jnp.float32)

model = RecurrentLIF()
brainstate.nn.init_all_states(model)
learner = braintrace.compile(
    model,
    braintrace.pp_prop,
    inputs[0],
    decay_or_rank=0.5,
)
optimizer = braintools.optim.Adam(2e-3)
optimizer.register_trainable_weights(learner.param_states)


def reset_sequence():
    brainstate.nn.reset_all_states(model)
    learner.reset_state()


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


def masked_objective():
    reset_sequence()
    outputs = learner.etrace_evolve(inputs, return_outputs=True)
    losses = jnp.mean((outputs - targets) ** 2, axis=1)
    return jnp.sum(losses * loss_mask) / jnp.sum(loss_mask)


def train_epoch(_):
    reset_sequence()
    grads, objective = learner.etrace_grad(
        inputs,
        targets,
        step_fn=step_loss,
        mask=loss_mask,
        reduction="mean",
        loss_output="scalar",
        return_value=True,
    )
    optimizer.update(brainstate.nn.clip_grad_norm(grads, 1.0))
    return objective


initial_loss = masked_objective()
training_losses = brainstate.transform.for_loop(
    train_epoch, jnp.arange(30)
)
final_loss = masked_objective()

print(f"initial masked loss: {float(initial_loss):.4f}")
print(f"final masked loss:   {float(final_loss):.4f}")
print("finite trajectory:", bool(jnp.all(jnp.isfinite(training_losses))))
print("loss decreased:", bool(final_loss < initial_loss))
print("compiled relations:", len(learner.graph.hidden_param_op_relations))

initial masked loss: 0.1687
final masked loss:   0.0467
finite trajectory: True
loss decreased: True
compiled relations: 2

5. What the check establishes#

The run verifies the public compile/sequence API, single-step state evolution, finite factorized updates, a shared train/evaluation window, and descent on one fixed task. It does not measure factorization error against BPTT. That needs a reduced finite-window gradient oracle for the intended model and decay.

References and API#