FixedRandomFeedback

class braintrace.FixedRandomFeedback

Frozen random feedback matrix with a stop-gradient guard.

The feedback matrix \(B \in \mathbb{R}^{n_{\mathrm{target}} \times n_{\mathrm{layer}}}\) is sampled once at construction and frozen via jax.lax.stop_gradient. It backs EProp’s random-feedback mode, and is the intended home for any rule that replaces the symmetric learning signal with a fixed random projection.

Parameters:

• n_target (int) – Number of target dimensions (the row count of B).

• n_layer (int) – Number of layer dimensions (the column count of B).

• key (jax.Array) – A PRNG key used to sample the feedback matrix. Obtain one from brainstate.random.split_key().

• init_scale (float, optional) – Standard-deviation scaling applied to the sampled normal entries. Default is 0.1.

Variables:

• B (jax.Array) – The frozen feedback matrix of shape (n_target, n_layer).

• n_target (int) – Number of target dimensions.

• n_layer (int) – Number of layer dimensions.

Examples

>>> import jax
>>> import brainstate
>>> import braintrace
>>>
>>> brainstate.random.seed(0)
>>> fb = braintrace.FixedRandomFeedback(2, 3, brainstate.random.split_key())
>>> print(fb.B.shape)
(2, 3)
>>> y = jax.numpy.ones(2)
>>> print(fb.project(y).shape)
(3,)

project(y_target)

Project the target onto the frozen feedback matrix.

Parameters:

y_target (jax.Array) – The target tensor to project. Both batched and unbatched layouts are handled.

Returns:

jax.Array – The projection y_target @ B with B frozen.

FixedRandomFeedback.__init__(n_target, n_layer, key, init_scale=0.1)