FORCE

Contents

FORCE#

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

FORCE learning: RLS on the readout of a recurrent network with feedback.

FORCE (first-order reduced and controlled error) learning [1] trains the linear readout \(z = W^\top r\) of a chaotic recurrent network whose output is fed back into the network. The readout is updated by RLS at every (or every few) time steps, with a strong initial learning rate (\(P_0 = I / \alpha\), \(\alpha\) of order 1), which keeps the error small from the start and suppresses the chaos. The update rule is exactly RLS; this class provides the FORCE defaults and the protocol documented below.

Parameters:

  • alpha (float) – Regulariser; \(P_0 = I / \alpha\). Sussillo and Abbott use values from 1 to 100 (smaller means faster learning).

  • forgetting (float) – Forgetting factor \(\lambda \in (0, 1]\), as in RLS.

Raises:

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

See also

RLS

the underlying update and the step signature.

Notes

Per time step: compute the rates \(r = \tanh(x)\) and the readout \(z = W^\top r\), advance the network with \(z\) fed back, then call step(r, z, target). Wrap the time loop in brainstate.transform.for_loop so that it compiles to a single XLA loop.

References

Examples

Train a chaotic rate network (g = 1.5) to produce a sine wave, then let it run on its own feedback:

>>> import brainstate
>>> import jax
>>> import jax.numpy as jnp
>>> import braintools
>>> N, g, dt, tau = 200, 1.5, 1.0, 10.0
>>> k1, k2, k3 = jax.random.split(jax.random.PRNGKey(3), 3)
>>> J = g * jax.random.normal(k1, (N, N)) / jnp.sqrt(N)
>>> w_fb = jax.random.uniform(k2, (N,), minval=-1., maxval=1.)
>>> x = brainstate.HiddenState(0.5 * jax.random.normal(k3, (N,)))
>>> w = brainstate.ParamState(jnp.zeros((N, 1)))
>>> opt = braintools.optim.FORCE(alpha=1.0)
>>> _ = opt.register_trainable_weights({'readout': w})
>>>
>>> def run(target, train):
...     r = jnp.tanh(x.value)
...     z = r @ w.value
...     x.value = x.value + dt / tau * (-x.value + J @ r + w_fb * z[0])
...     if train:
...         opt.step(r[None], z[None], target[None])
...     return z[0]
>>>
>>> t = jnp.arange(5000) * dt
>>> target = jnp.sin(2 * jnp.pi * t / 100.)[:, None]
>>> _ = brainstate.transform.for_loop(lambda y: run(y, True), target[:4000])
>>> z = brainstate.transform.for_loop(lambda y: run(y, False), target[4000:])
>>> bool(jnp.mean((z - target[4000:, 0]) ** 2) < 0.05)
True
Contents