DNI

DNI#

class braintrace.DNI#

Decoupled Neural Interfaces: a learned estimate of the truncated future.

The coordinate is D_RTRL’s with learning_signal='bootstrapped', and unlike the other two P4 presets it is multi-step: the whole point is a window with an exit, and a one-step window has almost no truncated future to estimate.

What DNI fixes, precisely. Index windows [a_k, b_k) with b_k = a_{k+1} and let l_t be the loss of the step that writes h^{t+1}. The injected estimate is

\[g_k \;\approx\; \frac{\partial \sum_{t \ge b_k} l_t}{\partial h^{b_k}}\]

– strictly future, half-open, so the exit step’s own loss lies inside window k and is not counted twice.

It reaches the plain parameters, the inputs and the other states, where the sum over windows then telescopes to the exact gradient. It deliberately does not reach the ETP parameters or the boundary learning signal: their cross-window credit is already carried by the eligibility trace, and adding the estimate there would count the same path a second time. So DNI does not make the ETP gradients better or worse – with a synthesiser attached they are bit-identical to the plain run – it gives the plain parameters the credit the trace already gives the ETP ones.

Parameters:

model (brainstate.nn.Module) – The one-step model.
synthesizer (SyntheticGradient, optional) – The estimator. May be attached later via attach_synthesizer(); a window that runs without one raises rather than silently degrading to the truncated rule.
name (str, optional) – Node name.
vjp_method (str, optional) – Must be 'multi-step' (the default).
fast_solve (bool, optional) – Whether registered closed-form kernels may be used. Default True.
trace_dtype (DTypeLike, optional) – Reduced trace precision.
chunked_trace (bool, optional) – Whether to roll the trace in chunks. Default True.
control_flow (ControlFlowPolicy, optional) – Control-flow canonicalization policy.
snap_max_jacobian_elements (int, optional) – Passed through; unused at recurrence_scope='diagonal'.

Examples

>>> import brainstate, braintrace, jax.numpy as jnp
>>> model = ...
>>> learner = braintrace.DNI(model)
>>> learner.compile_graph(braintrace.MultiStepData(xs))
>>> learner.init_etrace_state()
>>> learner.attach_synthesizer(
...     braintrace.SyntheticGradient(learner.group_signal_shapes()))

Notes

The synthesiser’s parameter values are threaded into the traced computation as an explicit argument rather than closed over, so jax.custom_vjp never captures a tracer it might be asked to differentiate, and the estimate is wrapped in stop_gradient so the online loss cannot train the synthesiser through the wrong path. Fit it with train_synthetic_gradient().

DNI

Contents

DNI#

Modeling

Infrastructure

Compilation

Vibe coding