UORO

UORO#

class braintrace.UORO#

Unbiased Online Recurrent Optimization.

The coordinate is ETraceConfig(trace_factorization='random_projection', recurrence_scope='coupled'). Both halves are load-bearing:

  • random_projection replaces the per-parameter influence trace with one rank-1 pair per hidden group, eps_tilde[j, u] ~= s_tilde[u] * theta_tilde[j], re-randomised each step with a Rademacher draw.

  • coupled is required, not merely recommended (matrix rule 11). A rank-1 unbiased estimator of an already-biased recursion would be strictly worse than the biased recursion itself: same asymptotic error, more variance, no memory saved — the anchored per-parameter trace is already the smaller carrier. coupled is the cheapest scope whose transition actually contains the hidden-to-hidden ETP mixing, so it is the coordinate where the projection buys something. Measured: rolling the block-diagonal transition converges cleanly onto the biased trace and never onto the exact one.

What follows is the claim in full, because “unbiased” on its own overstates it: UORO is an unbiased estimator of the exact within-group influence recursion that the compiled transition defines — that is, the saturating end of the SnAp scale, at SnAp-1 memory, in expectation. It does not repair cross-group coupling, the F-31 instantaneous tail, or any primitive’s own solve regime. On a single hidden group with a position-preserving elementwise tail, that recursion is itself exact, so there UORO is unbiased for BPTT; elsewhere it is unbiased for the block-local recursion. It runs on every model either way.

Parameters:

  • model (brainstate.nn.Module) – The one-step model.

  • name (str, optional) – Node name.

  • vjp_method (str, optional) – 'multi-step' (default) or 'single-step'. The finite-window oracle path — and hence every acceptance criterion — uses 'multi-step'.

  • fast_solve (bool, optional) – Whether registered closed-form kernels may be used when building the projector. Default True.

  • control_flow (ControlFlowPolicy, optional) – Control-flow canonicalization policy.

  • projection_key (int or jax.Array, optional) – Seed for the projection stream. Default 42. Two different keys give two different gradients; reset_state re-derives the stream from this value, so a reset run repeats bit-for-bit.

  • projection_eps (float, optional) – Guard added to every norm before its ratio. Default 1e-12; 0.0 makes the first step NaN, since both norms start at zero.

  • random_feedback_key (jax.Array, optional) – Required when combining with learning_signal='random_feedback'.

  • snap_max_jacobian_elements (int, optional) – Passed through; unused at recurrence_scope='coupled'.

Examples

>>> import brainstate, braintrace, jax.numpy as jnp
>>> class Net(brainstate.nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.w = brainstate.ParamState(0.1 * jnp.ones((4, 4)))
...         self.h = brainstate.HiddenState(jnp.ones((1, 4)) * 0.4)
...     def update(self, x):
...         self.h.value = jnp.tanh(x + braintrace.matmul(self.h.value, self.w.value))
...         return self.h.value
>>> model = Net()
>>> brainstate.nn.init_all_states(model, batch_size=1)
>>> learner = braintrace.UORO(model, projection_key=0)
>>> learner.compile_graph(braintrace.MultiStepData(jnp.zeros((1, 1, 4))))
>>> learner.init_etrace_state()

UORO is multi-step by construction, so etrace_grad drives it in window mode: pass chunk_size=k with k >= 2, and step_fn receives a (k, ...) slice, wraps its model input in MultiStepData, and returns a (k,) vector of per-step losses rather than a scalar.

>>> xs = jnp.zeros((10, 1, 4))
>>> ys = jnp.zeros((10, 1, 4))
>>> def window_loss(x, y):          # x, y are (k, 1, 4)
...     out = learner(braintrace.MultiStepData(x))
...     return jnp.mean((out - y) ** 2, axis=(1, 2))    # (k,)
>>> grads, losses = learner.etrace_grad(
...     xs, ys, step_fn=window_loss, chunk_size=5, return_value=True)

Notes

Variance grows with the number of window boundaries. The estimate is unbiased, not low-variance: a single run can be far from the mean, and averaging is the caller’s job. Antithetic sampling does not help here and looks like it should — flipping the sign of the entire draw sequence leaves the estimate bit-identical, because rho1 is even and both factors flip together, so their product is invariant.

See also

braintrace.SnAp

the deterministic scale whose saturating end this estimates, at higher memory.

braintrace.D_RTRL

the diagonal-scope biased trace.

UORO.__init__(model, name=None, vjp_method='multi-step', fast_solve=True, control_flow=None, projection_key=42, projection_eps=1e-12, random_feedback_key=None, snap_max_jacobian_elements=16777216)#