SyntheticGradient

class braintrace.SyntheticGradient

A per-hidden-group linear synthesiser of the future cotangent.

One affine map per hidden group, taking that group’s concatenated hidden value (*varshape, num_state) to a cotangent of the same shape. Linear with a bias, as in the paper: the target is itself a gradient, so a linear predictor is a reasonable hypothesis class and keeps the auxiliary problem convex.

The parameters are not ETP routed. They are ordinary brainstate.ParamStates that never appear in a hidden_param_op_relation, because the synthesiser is not part of the model’s recurrence: it observes h^exit and predicts, it does not participate in producing h. DNI keeps them out of the compiled graph by construction – it is handed the values functionally rather than closing over the states.

The final layer is zero-initialised, so a freshly constructed synthesiser predicts exactly zero and the learner starts bit-identical to the plain truncated rule. That is a deliberate property: it makes “DNI is off” and “DNI is untrained” the same run, so B1’s no-op criterion is checkable.

Parameters:

• group_shapes (dict of int to tuple) – group index -> (*varshape, num_state), from the compiled graph.

• hidden_width (int, optional) – Unused; kept for signature stability.

• scale (float, optional) – Standard deviation of the input-layer draw. Default 0.0 – see the zero-initialisation note above; a non-zero value makes the synthesiser live from the start, which the B1 negative control needs.

• seed (int, optional) – Seed for the draws.

Examples

>>> import braintrace
>>> synth = braintrace.SyntheticGradient({0: (1, 4, 1)})
>>> values = synth.param_values()
>>> est = synth.apply(values, {0: jnp.zeros((1, 4, 1))})
>>> est[0].shape
(1, 4, 1)

apply(param_values, group_hiddens)

Predict each group’s future cotangent, functionally.

Parameters:

• param_values (dict) – As returned by param_values(). Passed explicitly rather than read off self so that a caller inside jax.custom_vjp never captures a tracer it might later be asked to differentiate.

• group_hiddens (dict of int to array) – group index -> concatenated hidden value.

Returns:

dict of int to jax.Array – group index -> estimated cotangent, shaped like the input.

param_values()

The synthesiser’s parameter values, for the functional call.

Returns:

dict – {'w': {gid: array}, 'b': {gid: array}}.

states_dict()

The synthesiser’s brainstate.ParamStates, keyed for an optimiser.

SyntheticGradient.__init__(group_shapes, hidden_width=None, scale=0.0, seed=0)