OSTL: with-H and without-H regimes

OSTL: with-H and without-H regimes#

Online Spatio-Temporal Learning is not one interchangeable preset. BrainTrace exposes the recurrent with-H and feedforward without-H regimes as separate classes because they retain different temporal terms.

1. Principle#

The recurrent regime propagates

\[\epsilon_t = H_t\epsilon_{t-1}+F_t,\]

where \(H_t=\partial h_t/\partial h_{t-1}\) is the hidden-to-hidden Jacobian and \(F_t\) is the instantaneous parameter contribution. The without-H regime drops \(H_t\epsilon_{t-1}\) and keeps the spatial/instantaneous contribution. In BrainTrace, OSTLRecurrent uses a coupled per-parameter trace; OSTLFeedforward uses an IO-factorized trace with negligible temporal decay.

2. Model-algorithm matching#

  • With-H: use for a recurrent layer when the retained block structure matches the model’s hidden coupling.

  • Without-H: use for feedforward SNN dynamics where no recurrent Jacobian should be propagated.

  • Do not apply the without-H rule to a recurrent model merely because it is cheaper; that changes the estimator’s mathematical target.

import brainstate
import braintrace
import jax.numpy as jnp

brainstate.random.seed(31)


class OSTLModel(brainstate.nn.Module):
    def __init__(self, recurrent):
        super().__init__()
        self.recurrent = recurrent
        self.w_in = brainstate.ParamState(
            0.25 * brainstate.random.randn(2, 5)
        )
        if recurrent:
            self.w_rec = brainstate.ParamState(
                0.1 * brainstate.random.randn(5, 5)
            )
        self.w_out = brainstate.ParamState(
            0.2 * brainstate.random.randn(5, 1)
        )

    def init_state(self, **kwargs):
        self.h = brainstate.HiddenState(jnp.zeros(5))

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

    def update(self, x):
        drive = braintrace.matmul(x, self.w_in.value)
        if self.recurrent:
            drive = drive + braintrace.matmul(self.h.value, self.w_rec.value)
        self.h.value = jnp.tanh(drive)
        return self.h.value @ self.w_out.value

3. Compile each regime on a matching model#

Both examples use the same public compile and sequence interfaces. They are not a head-to-head accuracy benchmark because the model structures intentionally differ.

inputs = brainstate.random.randn(10, 2)
targets = jnp.zeros((10, 1))


def compile_and_measure(algorithm, recurrent):
    model = OSTLModel(recurrent=recurrent)
    brainstate.nn.init_all_states(model)
    learner = braintrace.compile(model, algorithm, inputs[0])

    def step_loss(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=step_loss,
        loss_output="scalar",
        return_value=True,
    )
    norm = jnp.sqrt(sum(jnp.sum(g * g) for g in grads.values()))
    return learner, loss, norm


with_h, recurrent_loss, recurrent_norm = compile_and_measure(
    braintrace.OSTLRecurrent, recurrent=True
)
without_h, feedforward_loss, feedforward_norm = compile_and_measure(
    braintrace.OSTLFeedforward, recurrent=False
)

print("with-H recurrence scope:", with_h.config.recurrence_scope)
print("without-H factorization:", without_h.config.trace_factorization)
print("without-H decay:", without_h.config.decay)
print("finite recurrent gradient:", bool(jnp.isfinite(recurrent_norm)))
print("finite feedforward gradient:", bool(jnp.isfinite(feedforward_norm)))

with-H recurrence scope: coupled
without-H factorization: io_factorized
without-H decay: (1e-06, 1e-06)
finite recurrent gradient: True
finite feedforward gradient: True

4. Evidence boundary#

The output verifies that the two public presets compile to different learning-rule coordinates and produce finite gradients on structurally matched models. It does not show that either is exact for a deep network. With-H is gradient-equivalent to BPTT only under the documented block-structured hidden Jacobian conditions; without-H deliberately removes temporal recurrence.

Reference and API#