e-prop: local eligibility and learning signals

e-prop: local eligibility and learning signals#

e-prop separates a synapse-local eligibility trace from a learning signal broadcast by the readout. The separation is the scientific content of the method; a generic loss curve alone would not demonstrate it.

1. Principle#

For synapse \(i\to j\), e-prop forms

\[\frac{d\mathcal{L}}{dW_{ji}} \approx \sum_t L_j^t\,\bar e_{ji}^t,\qquad \bar e_{ji}^t=\kappa\bar e_{ji}^{t-1}+e_{ji}^t.\]

\(e_{ji}^t\) is local to the neuron/synapse dynamics. \(L_j^t\) carries task information from the readout. Symmetric feedback uses the actual reverse-AD learning-signal direction. Random feedback uses a fixed projection and removes weight transport only partially in the current hook contract; it must not be described as numerically equivalent to symmetric feedback.

2. Applicability and exclusions#

Use e-prop for recurrent LIF/ALIF-style networks when local traces and an explicit learning-signal approximation are the intended inductive bias.

Do not infer exact BPTT gradients for arbitrary recurrent coupling. The \(\kappa\) filter changes temporal credit assignment, while random feedback adds a separate approximation and requires an explicit key.

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

brainstate.random.seed(21)


class SmallEPropSNN(brainstate.nn.Module):
    def __init__(self, n_in=2, n_rec=6):
        super().__init__()
        self.n_rec = n_rec
        self.w_in = brainstate.ParamState(
            0.3 * brainstate.random.randn(n_in, n_rec)
        )
        self.w_rec = brainstate.ParamState(
            0.1 * brainstate.random.randn(n_rec, n_rec)
        )
        self.w_out = brainstate.ParamState(
            0.2 * brainstate.random.randn(n_rec, 1)
        )
        self.surrogate = braintools.surrogate.ReluGrad()

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

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

    def update(self, x):
        drive = (
            braintrace.matmul(x, self.w_in.value)
            + braintrace.matmul(self.z.value, self.w_rec.value)
        )
        voltage = 0.9 * self.v.value + drive - self.z.value
        spike = self.surrogate(voltage - 0.45)
        self.v.value = voltage
        self.z.value = spike
        return spike @ self.w_out.value

3. Compile and expose the filter coordinate#

The sequence driver advances the SNN once per call. This example uses symmetric feedback so the effect under inspection is the \(\kappa\) eligibility filter, not a simultaneous random-feedback approximation.

inputs = brainstate.random.bernoulli(
    0.35, size=(16, 2)
).astype(jnp.float32)
targets = jnp.full((16, 1), 0.25)
mask = (jnp.arange(16) >= 4).astype(jnp.float32)

model = SmallEPropSNN()
brainstate.nn.init_all_states(model)
learner = braintrace.compile(
    model,
    braintrace.EProp,
    inputs[0],
    feedback="symmetric",
    kappa_filter_decay=0.8,
)


def loss_step(x, target):
    return jnp.mean((learner(x) - target) ** 2)


brainstate.nn.reset_all_states(model)
learner.reset_state()
grads, loss = learner.etrace_grad(
    inputs,
    targets,
    step_fn=loss_step,
    mask=mask,
    loss_output="scalar",
    return_value=True,
)
gradient_norm = jnp.sqrt(sum(jnp.sum(g * g) for g in grads.values()))
etraces = learner.get_etrace_of(model.w_rec)
trace_norm = jnp.sqrt(sum(
    jnp.sum(leaf * leaf) for leaf in jax.tree.leaves(etraces)
))

print("trace filter:", learner.config.trace_filter)
print("kappa:", learner.config.kappa)
print("learning signal:", learner.config.learning_signal)
print(f"masked loss: {float(loss):.4f}")
print("finite nonzero trace:", bool(jnp.isfinite(trace_norm) & (trace_norm > 0)))
print("finite nonzero gradient:", bool(jnp.isfinite(gradient_norm) & (gradient_norm > 0)))

trace filter: kappa
kappa: 0.8
learning signal: symmetric
masked loss: 0.0586
finite nonzero trace: True
finite nonzero gradient: True

4. Interpretation and limits#

The printed configuration verifies which learning-rule coordinates were compiled. The nonzero recurrent-weight trace shows that the local eligibility state was advanced, while the finite gradient checks its contraction through the public sequence path. Neither check isolates the numerical effect of \(\kappa\). A causal comparison would hold model, sequence, feedback, and parameters fixed while varying only \(\kappa\), then use a finite-window oracle appropriate to the learning-rule claim.

Reference and API#