RLS

Contents

RLS#

class braintools.optim.RLS(alpha=1.0, *, forgetting=1.0)#

Recursive least squares (RLS) for linear readouts.

For every registered weight \(W\) (shape (n_in, n_out)) RLS keeps an inverse-correlation matrix \(P\) (shape (n_in, n_in), started at \(P_0 = I / \alpha\)). One step() on a batch of inputs \(X\) ((B, n_in)) with error \(E = Y - T\) applies the block update of brainpy 2.x [1] with a forgetting factor \(\lambda\):

\[\begin{split}K &= P X^\top (\lambda I_B + X P X^\top)^{-1}, \\ P &\leftarrow (P - K X P) / \lambda, \\ W &\leftarrow W - K E.\end{split}\]

With \(\lambda = 1\) and outputs = inputs @ W, after any number of steps \(W\) is the ridge solution \(\arg\min_W \|Y_{all} - X_{all} W\|^2 + \alpha \|W - W_0\|^2\) over all samples seen so far, whatever the batch sizes: one step on a batch equals one step per sample in the same order.

Parameters:

  • alpha (float) – Regulariser (> 0); \(P_0 = I / \alpha\). Small alpha (a large \(P_0\)) makes the first samples count fully and gives fast learning. Note that brainpy 2.x used the inverse convention (\(P_0 = \alpha I\)).

  • forgetting (float) – Forgetting factor \(\lambda \in (0, 1]\), applied once per step() (so per batch, not per sample). Below 1, old batches are down-weighted by \(\lambda^{k}\) after k further steps.

Raises:

ValueError – If alpha <= 0 or forgetting is outside (0, 1].

See also

FORCE

RLS applied to the readout of a recurrent network with feedback.

LMS

first-order (gradient-descent) online rule.

ridge_regression

the same solution in closed form, from all data at once.

Notes

step() is jit-able (brainstate.transform.jit) and can be called inside brainstate.transform.for_loop/scan: the weights and the \(P\) matrices are State objects updated in place. Several weights can be registered; then inputs, outputs and targets are dicts keyed like the weights (each weight is an independent readout of its own inputs). Weights and data must be dimensionless. Each step costs \(O(n_{in}^2 B + B^3)\).

References

Examples

>>> import brainstate
>>> import jax
>>> import jax.numpy as jnp
>>> import braintools
>>> w_true = jnp.array([[1.], [-2.], [0.5]])
>>> X = jax.random.normal(jax.random.PRNGKey(0), (200, 3))
>>> T = X @ w_true
>>> w = brainstate.ParamState(jnp.zeros((3, 1)))
>>> opt = braintools.optim.RLS(alpha=1e-3)
>>> _ = opt.register_trainable_weights({'readout': w})
>>> def train(x, t):
...     return opt.step(x[None], None, t[None])
>>> errors = brainstate.transform.for_loop(train, X, T)
>>> errors.shape
(200, 1, 1)
>>> print(jnp.round(w.value[:, 0], 3))
[ 1.  -2.   0.5]
Contents