RandomProjectionVjpAlgorithm

class braintrace.RandomProjectionVjpAlgorithm

Rank-1 random-projection eligibility trace — the UORO engine.

Per hidden group g the influence of parameter j on hidden unit u = (position, state) is carried by one rank-1 pair

\[\tilde\varepsilon_g[j, u] \approx \tilde s_g[u] \, \tilde\theta_g[j]\]

with \(\tilde s_g\) shaped (*varshape, num_state) and \(\tilde\theta_g\) parameter-shaped — no hidden axis, no trailing state axis. One \(\tilde s\) per group; one \(\tilde\theta\) per (group, ParamState path), so a weight consumed by two relations of one group keeps a single parameter factor and the two projections are summed into it.

With \(\nu_g\) a Rademacher draw of \(\tilde s_g\)’s shape and \(J_f\) the instantaneous term, the step is

\[\begin{split}\mathrm{proj}_g &= \nu_g^\top J_f \\ \rho_0 &= \sqrt{(\|\tilde\theta_g\| + \epsilon) / (\|D_g \tilde s_g\| + \epsilon)} \\ \rho_1 &= \sqrt{(\|\mathrm{proj}_g\| + \epsilon) / (\|\nu_g\| + \epsilon)} \\ \tilde s_g &\leftarrow \rho_0 (D_g \tilde s_g) + \rho_1 \nu_g \\ \tilde\theta_g &\leftarrow \tilde\theta_g / \rho_0 + \mathrm{proj}_g / \rho_1\end{split}\]

and the gradient contributed at a window boundary is \(\sum_g (\text{signal}_g \cdot \tilde s_g)\, \tilde\theta_g\) — a scalar per group times a parameter-shaped array.

Unbiased for what. Conditionally on earlier draws, the updated outer product is a b^T + c d^T + (rho0/rho1) a d^T + (rho1/rho0) c b^T with a = D s_tilde, b = theta_tilde, c = nu, d = nu^T J_f. The first term is D eps_tilde (rho0 cancels), the second is nu nu^T J_f, and the two cross terms are odd in the current draw while rho1 is even, so any negation-symmetric draw law with E[nu nu^T] == I gives E[eps_tilde_new] == D eps_tilde + J_f exactly. Induction over time then gives an unbiased estimate of the exact within-group influence recursion that the compiled transition defines.

That is the whole claim, and it is narrower than “unbiased”. Three approximations are untouched:

1. cross-group coupling — the compiler splits hidden states into groups and drops inter-group terms; one factor pair per group cannot carry them;

2. the instantaneous term’s tail — df comes from a single all-ones JVP of the y -> hidden map, exact only for a position-preserving elementwise tail (finding F-31);

3. each primitive’s own solve regime — the projector is built from the framework’s own per-primitive rules, so it inherits them exactly.

On a single-group model with a position-preserving elementwise tail — the class where that recursion is itself exact — this is therefore unbiased for BPTT. Elsewhere it is unbiased for the block-local recursion and biased against BPTT, like every other coordinate in this repository. It is not restricted to that class.

Memory, honestly. The carrier is |theta| + B * P * S elements against B * |theta| * S for the anchored per-parameter trace: on tanh_rnn(3, 4) that is 20 against 16. UORO is not a memory win over the anchored trace — it is a bias win at comparable carrier size, and a memory win against saturated SnAp-n (64) or the full influence matrix (64), which is what unbiasedness would otherwise cost. And O(|theta| + P S) describes carrier storage, not peak memory: the full transition Jacobian is an O((P S)^2) transient per step (follow-up F-32 tracks a matrix-free D @ s product). The fused stepper keeps that transient per-step rather than stacked over the window.

Parameters:

• model (brainstate.nn.Module) – The one-step model.

• name (str, optional) – Node name.

• vjp_method (str, optional) – 'single-step' or 'multi-step'. Default 'multi-step': the chunked finite-window path requires it, and it is the path the estimator is validated on.

• fast_solve (bool, optional) – Whether registered closed-form kernels may be used when building the projector. Default True.

• control_flow (ControlFlowPolicy, optional) – Control-flow canonicalization policy.

• config (ETraceConfig, optional) – The learning-rule coordinate. Must have trace_factorization='random_projection'.

• projection_key (int or jax.Array, optional) – Seed for the projection stream. Default 42. reset_state re-derives the stream from it, so a reset run repeats bit-for-bit.

• projection_eps (float, optional) – Guard added to every norm before the ratio. Default 1e-12. At the first step both factors are zero, so both ratios are 0/0; with projection_eps=0 the carrier is NaN at every window length. Because the guard perturbs the normalisers, the exactness pins carry a tolerance above float64 epsilon — see uoro_test.py.

• random_feedback_key (jax.Array, optional) – Passed through to the base class for learning_signal='random_feedback'.

• snap_max_jacobian_elements (int, optional) – Passed through; unused at recurrence_scope='coupled'.

See also

braintrace.UORO

the preset that fixes the coordinate.

braintrace.SnAp

the deterministic scale whose saturating end this estimates.

get_etrace_of(weight)

The parameter-side factors associated with weight.

Parameters:

weight (brainstate.ParamState or Path) – The weight whose factors are requested.

Returns:

dict – {(group index, path): theta_tilde} for that weight.

Raises:

ValueError – If the weight has no random-projection factor.

init_etrace_state(*args, **kwargs)

Allocate the rank-1 factors, the projection key and the step counter.

reset_state(batch_size=None, **kwargs)

Reset the factors, the step counter and the projection stream.

The stream is re-derived from projection_key rather than advanced, so a reset run reproduces the previous one bit-for-bit.

RandomProjectionVjpAlgorithm.__init__(model, name=None, vjp_method='multi-step', fast_solve=True, control_flow=None, config=None, projection_key=42, projection_eps=1e-12, random_feedback_key=None, snap_max_jacobian_elements=16777216)