ETraceVjpAlgorithm

class braintrace.ETraceVjpAlgorithm

Provide VJP-based eligibility-trace gradient computation.

The term VJP comes from two aspects. First, this module is designed to be compatible with JAX’s VJP mechanism, so the gradient is computed according to the reverse-mode differentiation interface, like jax.grad, jax.vjp, or jax.jacrev. The true update function is defined as a custom VJP function ._true_update_fun(), which receives the inputs, the hidden states, other states, and etrace variables at the last time step, and returns the outputs, the hidden states, other states, and etrace variables at the current time step. Second, the algorithm computes the spatial gradient \(\partial L^t / \partial H^t\) using the standard back-propagation algorithm, which enhances the accuracy and the stability of the gradient computation.

Parameters:

• model (brainstate.nn.Module) – The model function, which receives the input arguments and returns the model output.

• name (str, optional) – The name of the etrace algorithm.

• vjp_method (str, optional) – The method for computing the VJP. It should be either "single-step" or "multi-step". Default is "single-step".

  • "single-step": The VJP is computed at the current time step, i.e., \(\partial L^t/\partial h^t\).

  • "multi-step": The VJP is computed at multiple time steps, i.e., \(\partial L^t/\partial h^{t-k}\), where \(k\) is determined by the data input.

• control_flow (ControlFlowPolicy, optional) – Policy governing control-flow canonicalization (cond if-conversion, scan unrolling, structured scan descent, …) during graph compilation. None (default) uses ControlFlowPolicy().

• config (ETraceConfig, optional) – Learning-rule coordinates. None uses the subclass’s preset coordinate.

• random_feedback_key (jax.Array, optional) – Key used to initialize fixed random-feedback projections when config.learning_signal='random_feedback'.

• snap_max_jacobian_elements (int, optional) – Maximum number of elements permitted in each SnAp widened block Jacobian. The default is 16777216.

Notes

For each subclass (or the instance of an etrace algorithm), the following methods define the custom VJP rule:

• ._update(): update the eligibility trace states and return the outputs, hidden states, other states, and etrace data.

• ._update_fwd(): the forward pass of the custom VJP rule.

• ._update_bwd(): the backward pass of the custom VJP rule.

This class provides a default implementation for the ._update(), ._update_fwd(), and ._update_bwd() methods. To implement a new etrace algorithm, users just need to override the following methods:

• ._solve_weight_gradients(): solve the gradients of the learnable weights / parameters.

• ._update_etrace_data(): update the eligibility trace data.

• ._assign_etrace_data(): assign the eligibility trace data to the states.

• ._get_etrace_data(): get the eligibility trace data.

compile_graph(*args)

Compile the eligibility trace graph of the relationship between etrace weights, states and operators.

The compilation process includes:

• building the etrace graph

• separating the states

• initializing the etrace states

Parameters:

*args – The input arguments.

init_etrace_state(*args, **kwargs)

Allocate the axis-side state.

Concrete here, unlike ETraceAlgorithm.init_etrace_state(), which raises: the lifted axes own state of their own, and every engine must reach it. Engines override this to build their traces and then call super().init_etrace_state(...) as the last statement.

Parameters:

*args, **kwargs – The example inputs, forwarded from compile_graph(). Unused here; engines size their traces from the compiled graph.

update(*args, modulator=None)

Update the model states and the eligibility trace.

The input arguments args here support very complex data structures, including the combination of SingleStepData and MultiStepData.

• SingleStepData: indicating the data at the single time step, \(x_t\).

• MultiStepData: indicating the data at multiple time steps, \([x_{t-k}, ..., x_t]\).

Parameters:

• *args – The input arguments.

• modulator (array_like or Quantity, optional) – The per-call modulatory signal for learning_signal='modulatory', taking precedence over the modulator attribute for this call only. It is not forwarded to the model’s forward call; it reaches the rule through _get_update_aux. Ignored on every other learning_signal.

Returns:

Any – The model output.

Notes

Suppose all inputs have the shape of (10,).

If the input arguments are given by:

x = [jnp.ones((10,)), jnp.zeros((10,))]

Then, two input arguments are considered as the SingleStepData.

If the input arguments are given by:

x = [braintrace.SingleStepData(jnp.ones((10,))),
     braintrace.SingleStepData(jnp.zeros((10,)))]

This is the same as the previous case, they are all considered as the input at the current time step.

If the input arguments are given by:

x = [braintrace.MultiStepData(jnp.ones((5, 10)),
     jnp.zeros((10,)))]

or,

x = [braintrace.MultiStepData(jnp.ones((5, 10)),
     braintrace.SingleStepData(jnp.zeros((10,)))]

Then, the first input argument is considered as the MultiStepData, and its data will be fed into the model within five consecutive steps, and the second input argument will be fed into the model at each time of this five consecutive steps.

ETraceVjpAlgorithm.__init__(model, name=None, vjp_method='single-step', control_flow=None, config=None, random_feedback_key=None, snap_max_jacobian_elements=16777216)