braintrace.train_synthetic_gradient

braintrace.train_synthetic_gradient#

braintrace.train_synthetic_gradient(learner, inputs, *, chunk_size=1, loss_fn=None, optimizer=None, lr=0.01, epochs=1, reset=True, batch_size=None)#

Fit the synthesiser against the learner’s own returned hidden cotangent.

The regression target for M(h^{a_k}) is dL_{>= a_k}/dh^{a_k}, and the learner already produces exactly that: with the hidden states in the differentiation set, brainstate.transform.grad returns the window’s hidden cotangent, which (by the second pass of _update_fn_bwd) already carries the future term. So no new side channel is needed – the target comes out of the public API.

Two properties make the fit honest, and both are enforced here:

  • the model parameters do not move – only the synthesiser’s are handed to the optimiser;

  • the target is detached, so the regression cannot reshape the model’s gradients to make itself easy to predict.

The auxiliary optimiser is left to the caller, as it is in the paper.

Parameters:

  • learner (DNI) – A compiled learner with a synthesiser attached.

  • inputs (array) – A (T, ...) sequence, consumed one window at a time. T must be a multiple of chunk_size; a ragged final window is refused rather than truncated (see Raises).

  • chunk_size (int, optional) – Steps per window. This must match the window size the learner will be driven with: the synthesiser predicts the future at a window boundary, and boundaries move when the window size does. Training on one-step windows and then deploying on longer ones fits the wrong target – a much shorter future – and the result can easily be worse than no synthesiser at all. Default 1.

  • loss_fn (callable, optional) – loss_fn(output) -> scalar. This must be the objective the learner will actually be trained on, and it is the second half of the same trap chunk_size documents. The synthesiser predicts dL_{>= b}/dh^b – a derivative of this loss. Fit it against the default sum-of-squares and then descend on, say, ((out - target) ** 2).mean(), and the injected cotangent is the gradient of a different function at a different scale: not an approximation of the future credit but noise with the shape of one. Measured on the delayed-reward fixture, a mismatched loss_fn left the run worse than leaving DNI off entirely. Default: sum of squares.

  • optimizer (braintools.optim.Optimizer, optional) – Already registered against learner.synthesizer.states_dict(). If None, a plain SGD step with learning rate lr is applied.

  • lr (float, optional) – Learning rate for the built-in SGD step. Ignored when optimizer is given. Default 1e-2.

  • epochs (int, optional) – Passes over inputs.

  • reset (bool, optional) – Whether to re-initialise the model states and the trace before each epoch. Default True.

  • batch_size (int, optional) – Batch size for that re-initialisation. Inferred from the learner’s own hidden states when omitted, which is almost always what you want; pass it only to override. It is not assumed to be 1 – doing so either raises a shape error on a wider learner or, worse, succeeds and fits the synthesiser against the wrong initial states.

Returns:

list of float – The mean squared prediction error per epoch, averaged over every window boundary and the terminal boundary.

Raises:

  • RuntimeError – If no synthesiser is attached to learner.

  • ValueError – If chunk_size is below 1, or if it does not divide len(inputs).

Notes

The window loop is a brainstate.transform.for_loop, so the body – which drives the learner through its custom_vjp and then takes a regression step – is traced once per epoch rather than once per window (AGENTS.md rule 10). Everything that changes across windows is State and threads through automatically: the model’s hidden states, the learner’s trace, the synthesiser’s parameters, and the optimiser’s moments. Epochs remain a Python loop, because reset calls init_all_states, which reallocates state and cannot run under a trace.

Each epoch fits one extra pair beyond the window loop: the terminal state h^T against a target of exactly zero. Deployment injects at window exit states while the loop iterates entry states, and those sets differ at the ends – h^T is injected at but never otherwise trained, and its true future gradient is zero because nothing follows it.