LMS

Contents

LMS#

class braintools.optim.LMS(lr=0.1)#

Least mean squares (LMS, Widrow-Hoff) rule for linear readouts.

One step() on a batch of inputs \(X\) ((B, n_in)) with error \(E = Y - T\) ((B, n_out)) applies the brainpy 2.x update [1]

\[W \leftarrow W - \eta \sum_{b=1}^{B} x_b e_b^\top = W - \eta X^\top E,\]

a stochastic-gradient step on \(\tfrac12 \|E\|^2\) (summed, not averaged, over the batch).

Parameters:

lr (float) – Learning rate \(\eta \ge 0\). With outputs = inputs @ W the rule is stable roughly for \(0 < \eta < 2 / \lambda_{max}\), with \(\lambda_{max}\) the largest eigenvalue of the per-step input correlation \(X^\top X\) (|x|^2 for a single sample).

Raises:

ValueError – If lr is negative.

See also

RLS

second-order online rule (much faster convergence).

Notes

step() is jit-able and usable inside brainstate.transform.for_loop; see RLS for the handling of several registered weights. The optimizer keeps no state besides the weights.

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), (2000, 3))
>>> w = brainstate.ParamState(jnp.zeros((3, 1)))
>>> opt = braintools.optim.LMS(lr=0.02)
>>> _ = opt.register_trainable_weights({'readout': w})
>>> def train(x):
...     return opt.step(x, None, x @ w_true)
>>> _ = brainstate.transform.for_loop(train, X)
>>> print(jnp.round(w.value[:, 0], 3))
[ 1.  -2.   0.5]
Contents