Second-order forward-mode optimization of recurrent neural networks for neuroscience
Authors: Youjing Yu
Rui Xia
Qingxi Ma
Máté Lengyel
Guillaume Hennequin
Venue: NeurIPS 2024 spotlight (2024)
Keywords: computational neuroscience
recurrent neural networks
motor control
Primary area: neuroscience_and_cognitive_science
TL;DR: A second-order, memory-efficient optimization method for RNNs that does not require backpropagation, runs in wallclock time comparable to first-order methods, and successfully trains RNNs where Adam fails

Abstract: A common source of anxiety for the computational neuroscience student is the question “will my recurrent neural network (RNN) model finally learn that task?”. Unlike in machine learning where any architectural modification of an RNN (e.g. GRU or LSTM) is acceptable if it speeds up training, the RNN models trained as _models of brain dynamics_ are subject to plausibility constraints that fundamentally exclude the usual machine learning hacks. The “vanilla” RNNs commonly used in computational neuroscience find themselves plagued by ill-conditioned loss surfaces that complicate training and significantly hinder our capacity to investigate the brain dynamics underlying complex tasks. Moreover, some tasks may require very long time horizons which backpropagation cannot handle given typical GPU memory limits. Here, we develop SOFO, a second-order optimizer that efficiently navigates loss surfaces whilst _not_ requiring backpropagation. By relying instead on easily parallelized batched forward-mode differentiation, SOFO enjoys constant memory cost in time. Morever, unlike most second-order optimizers which involve inherently sequential operations, SOFO's effective use of GPU parallelism yields a per-iteration wallclock time essentially on par with first-order gradient-based optimizers. We show vastly superior performance compared to Adam on a number of RNN tasks, including a difficult double-reaching motor task and the learning of an adaptive Kalman filter algorithm trained over a long horizon.

Second-order forward-mode optimization of recurrent neural networks for neuroscience
Youjing Yu
Rui Xia
Qingxi Ma
Máté Lengyel
Guillaume Hennequin
Computational and Biological Learning Lab
Department of Engineering
University of Cambridge, Cambridge, UK
(yy471, rx220, gm218)@cam.acc ac. uk, {m. Im.lenge,, g hhennequin)@eng cam ac.uk
Abstract
Training recurrent neural networks (RNNs) to perform neuroscience tasks can be challenging. Unlike in machine learning where any architectural modification of an RNN (c.g. GRU or LSTM) is acceptable if it facilitates training, the RNN models trained as models of brain dynamics are subject to plausibility constraints that funda- mentally exclude the usual machine learning hacks. The "vanilla" RNNs commonly used in computational neuroscience find themselves plagued by ill-conditioned loss surfaces that complicate training and significantly hinder capacity investigate the brain dynamics underlying complex tasks. Moreover, some tasks may require very long time horizons which backpropagation cannot handle given typical GPU memory limits. Here, we develop SOFO, a second-order optimizer that efficiently navigates loss surfaces whilst not requiring backpropagation. By relying instead on easily parallelized batched forward-man mode differentiation, SOFÓ enjoys constant memory cost in time. Moreover, unlike most second-order optimizers which in- volve inherently sequential operations, SOFO's s effective use of GPU parallelism yields per-iteration wallelock time essentially on par with first-order gradient- based optimizers. We show vastly superior performance compared to Adam on number of RNN tasks, including difficult double-reachiee reaching motor task and the learning of adaptive Kalman filter algorithm trained over long horizon.
1 Introduction
In recent years, trained recurrent neural networks (RNN) have gained increasing adoption as models of brain circuits dynamics [47] I As flexible parametric models of sequential dynamics, RNNs can be trained to perform specific computations or reproduce certain behaviors [3] 50, 20 [3]
and be subsequently probed for insights into the distributed computations that give risc to those RNNs are also often used as expressive models of the latent dynamics underlying
behaviors 2 53
the spatiotemporal structure of neural recordings 39
Despite their flexibility, RNNs are notoriously difficult to train [40] 51 36). In addition to the classic problem of vanishing gradients in temporally extended tasks, the loss surfaces that arise from RNN dynamics complex tasks often exhibit pathological curvature that first-order gradient-based optimization techniques struggle to handle. These problems normally addressed in multiple ways. First, the RNNs used in machine learning are modified from their standard ("vanilla"") formulation to include specific gate variables [18] 8) that largely mitigate the vanishing gradient problem. However, there is little biological support for the existence of such gating mechanisms, such that computational neuroscience studies have typically restricted themselves to vanilla RNNs [46). Second, second-order optimizers such as Hessian-free optimization [32 33 50 or KFAC S4 can be used to dramatically
accelerate training. However, those difficult to scale, often much slower in wall-clock time, and
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

memory hungry. Alternatively, the loss surface can be regularized by introducing artificial temporal "skip connections", and the network is progressively weaned off those biologically meaningless connections during the course of training [46].
Another great obstacle training RNNs lies in the memory complexity of gradient-based optimization via backpropagation through time (BPTT). In certain applications, such as in motor neuroscience, RNNs must be trained on tasks that require the production of smooth behavior requiring the use of small (millisecond) time steps and therefore long time horizons. In other applications such where the parameters of synaptic plasticity rule are optimized to yield a specific learning behavior, the network must be simulated over horizon long enough to span the slow timescale of synaptic modifications. Since backpropagation has a memory cost that scales with the time horizon, its naive application is often prohibitively expensive in these scenarios.
as "learning to learn" 22 I0]
Memory efficiency can be improved in a variety of ways. Checkpointing [42] allows discarding information between sporadic checkpoints in the forward pass, at the expense of having to recompute that information during the backward pass. Alternatively, backpropagation through time can be truncated [5]]: instead of computing single gradient for the entire sequence of size T gradients are computed sequentially for consecutive sub-sequences of size T" << T. Although cutting the chain of backpropagating gradients in this way reduces the memory cost from (T) to O the resulting gradient estimates are biased and neglect long-range temporal dependencies that may be important for the task [35 (38. Yet another way of reducing the memory footprint of BPTT is to formulate the RNN dynamics as an ordinary differential equation (ODE) in continuous time. In these "neural ODE" models, gradients
obtained by solving another (adjoint) ODE with no need for caching intermediate states 7 26). However, this requires very accurate, often adaptive ODE solvers that are highly nrivii to implement and adapt the mini-batch setting.
Finally, training RNNs is known to largely under-exploit the parallelization capabilities of modern (GPU) hardware. RNN gradient computations are inherently sequential forward-backward such that the only way to exploit GPU parallelism is to use large batches (e.g. run a large number of trials of the task in each training iteration). Whilst using large batches may enable faster learning with larger learning rates, it is often the case that increasing the batch size beyond a certain point no longer helps in this respect [25]. Thus, some of the parallel processing power of GPU hardware is often underused.
In summary, training RNNs is plagued with many problems. While there exist piecemeal solutions to each of these challenges, we currently lack a simple method that addresses all of them simultaneously. In this paper, we develop such method, which we call SOFO (Second order Forward mode Optimization). At each training iteration, SOFO reparameterizes the model in a random subspace, and uses batched forward-mode automatic differentiation to efficiently compute exact Generalized Gauss-Newton (GGN) updates in that subspace. This circumvents the need for backpropagation. enabling the training of RNNs over very long horizons. We test the performance of SOFO on benchmark tasks (c.g. learning an adaptive Kalman filter over long horizons, performing a motor reach with biomechanical arm; Section 4.. We show empirically that, despite only exploring small fraction of the parameter space in each iteration, SOFO dramatically accelerates training of RNNs in these complex tasks. More generally, we suggest that SOFO strong candidate for machine learning applications that involve relatively few trainable parameters but memory- -hungry computational graphs(e.g. g. fine-tuning of transformers using low-rank adapters 19. or tuning of hyperparameters in Gaussian process-based models).
general purpose optimizer is
2 Background and related work
2.1 Problem formulation
We consider supervised learning problems where neural network with parameters O E R is given batches of inputs æ, and produces batches of outputs y(x, 0) E R where M is the batch size. The specific case of an RNN with output dimension D simulated over time horizon T would imply N = DT i.e. the network's outputs consist of entire sequences. Training the network involves minimizing stochastic cost function c(0, x)) averaged over the training data in each minibatch, that

---

depends on the network's output y and therefore on the parameters A:
M
N Emi (ymi(0,)) 3mm)
x0,x) E(y(0, x))
A
M
m
t=1
where each (mx(-) function is strongly convex in its first argument (which does not imply that the overall cost c(O, is convex in 0). In RNN applications, runs from to N DT, i.e. it indexes the cartesian product of time bins x D output dimensions; the explicit dependence of lma on thus allows for the possibility of providing teaching signals only at specific times. To eliminate clutter, in most of this paper we will drop the explicit dependence of the network's 's outputs on the inputs æ, focusing instead on their dependence on 0.
2.2 Generalized Gauss-Newton optimization
The Generalized Gauss-Newton (GGN) method prescribes parameter updates of the form:
dc
0++1 (00nnG((,2)))
00
where 17 is a learning rate and G(O, x)+ is the Moore-Penrose pseudo-inverse of the GGN matrix
G(O, m) E RCxP, defined
Oe θrz² v(0,w)
G(O, a J(0,x))
J(0. x).
(3)
Here, 2(0,z) E MN is the Jacobian of the network's s output y w.r.t. the parameters 0. We note that for an underparameterized model with a number of network outputs M. exceeding the number of parameters P, the GGN matrix can be full rank, in which case the pseudoinverse coincides with the inverse. For overparameterized models, it has rank S MN. In both cases, it is common practice to damp the inverse to enhance the stability of the training process; this is done by using G (G + + instead of G+ with some damping parameter %.
2.3
Forward-mode automatic differentiation
Given a computational graph with inputs F²"' (keeping the notation relevant to our use case) and producing outputs y@)). forward-mode AD B allows the directional derivative 0y(0) de v (i.e. a Jacobian- vector product, JVP) in any arbitrary direction v to be automatically computed together with y() itself, with roughly the same computational and memory complexity. This is achieved by initializing the parameters with primal value 0 and tangent value Ô = v, extending the standard mathematical operators to operate such primal/tangent pairs, and running the computation forward in this dual space. For any intermediate node z in the computational graph, the tangent value é that produced has the same shape as and represents the sensitivity of z to small dieeneinniii changes in 0 in the direction of w:
z(0 + Ev) z(0)
θz
c
00
R
R.X
Note that unlike backpropagation, or "reverse-mode AD" forward-mode AD does not require caching any of the intermediate results that lead to the output interest. SOFO exploits the fact that Jacobian-vesd vector products are embarassingly parallelizable, i.e. one can rewrite the standard maths functions to operate not on one tangent per value, but on a whole batch of tangents in parallel. This can be done in JAX (5) out-of-the- box by composing Jacobian-vector products with the vmap primitive. As no such functionality exists in PyTorch yet, we provide our own flexxile implementation of batched JVPs based on OCaml--Toret!!!
code available at abaatggg//ggthhibb.ccm//emnuggee

---

2.4 Subspace optimization methods
variety of methods exist that, like SOFO, optimize parameters in a different subspace in each iteration. A well-known example is coordinate descent, which iteratively minimizes the objective function w.r.t. each parameter, one at a time [55]. More generally, one may optimize in a higher- dimensional subspace than ID, in coordinate system that necessarily axis-aligned. This leads to family of randomized subspace algorithms, including stochastic subspace descent (SSD:27 (7) and its variance-reduced version inspired by stochastic variance-reduced gradient (SVRG:22 (2)) The above methods can be construed as Jacobian sketching methods, whereby a random "sketch" of the Jacobian obtained and used estimate the gradient; the convergence of such sketching algorithms has been proved [27]] 28 under standard assumptions on the loss function (e.g. Polyak-Lojasiewing condition; 24). Similarly, it is possible to sketch not only the Jacobian but also any of the matrices that are normally used as curvature estimates, leading to a family of randomized second-orde optimization algorithms. These include the sketched Newton algorithm 4T randomized subspace Newton [4). randomized subspace Gauss-Newin 6 stochastic dual Newton ascent (43) and stochastic subspace cubic Newton [5). Sketched Hessians are usually obtained by nested forward- and reverse-mode AD, and as such incur the same memory complexity as backpropagation. To obtain rank- K sketch of K --dimensional subspace,
the Jacobian, and therefore compute directional first-order derivatives
one can instead perform K independent forward-mode AD computations [28] 4
Here, extend this of forward-mode AD the sketching of the Generalized Gauss-Newtor matrix, which we show can be performed efficiently on GPUs. This leads to a memory- and compute- efficient algorithm that enjoys the fast convergence properties of second-order method, with the runtime complexity of first-order optimizer.
SOFO
SOFO (Algorithm T based on successively optimizing low-dimensional affine re-parameterizations of the model, randomized in each training iteration. Specifically, instead of updating all P parameters independently at every step as is normally done, SOFO locally reparameterizes the model by writing the cost function as
( - c(0. + +
(5) every
where the K columns of θ form random K -idimensional subspace that is drawn iteration. An exact GGN update A (see below) is then obtained for this momentary lower-order model, leading to new parameters 0c+1 O, 7 OA0* around which the model is again re- parameterized in the next iteration. The gradient and GGN matrix associated with (-(-) are related to
Algorithm SOFO: Second-order Forward Optimisation
1: input: 0g !R
2: hyperparameters: subspace dimension K, learning rate n, relative damping À
3: convention: for any variable E R*, let uppercase Z E R denote the associated batch of K tangent values (e.g. (0, 0), (y. Y), fc, C))
4: for = 2, do
> training iterations
5: sample data minibatch
6:
sample θ E P with Opk ~ (0,
iid.
D subspace randomization
7:
(y.Y) network. Cotorrkoouuuut((tt (0, 0
b under batched forward- AD D under batched forward-mode AD
8:
Cc, e( (v( Y}) (c.f. Equation 1
9:
form sketched GGN matrix G - Y HY
> c.f. Equation
10: compute SVD of Ō = USUT
11: extract 8max = max (S)
max. singular value & SOFO update
0 0, 7) OU + \0mux -TUTC
13:
end for
output: last iterate 0,
note: for RNNs, memory can be saved in steps 7-9 by directly accumulating & and C over time steps.

---

those of c(-) via the chain rule:
Oc,
dc
and
G = = GO.
(6)
020
00
R
RK
Thus, O' is random (K x K) sketch of the full (P P) GGN matrix (c.f. Section 2.4 Inserting
Equation into Equation 6 we obtain
8
(JO))
(JO)
Qy²
which is then used compute the main (damped) SOFO parameter update,
θ dc
0++1 = 0: 17 1
(8)
)
00)
where y is the damping parameter (see below). SOFO uses batched forward- AD to efficiently compute JO and @ de on parallel hardware. As in standard forward AD, we rewrite the standard math operators to operate on primal/tangent pairs (c.f. Section 2.3). However, instead of working with single tangents, we operate directly on entire batches of K tangent values for every primal value (note: these tangent batches have nothing do with data minibatches). Those K tangents can be propagated completely independently of each other through every differentiable operation leading to network outputs and the final cost. Thus, computing such Jacobian facobian-matrix products is an embarrassingly parallel problem that can fully exploit GPU parallelism. The steps detailed in Algorithm D show concretely how JO = Y is obtained as the tangent batch associated with the network outputs y, and how the vector of directional derivatives θ de C' is the tangent batch associated with the final cost c.
30
Given that y is typically large tensor (batch size x time horizon x output dimension), one might worry that the Hessian of {(y) in Equation could be expensive to compute. However, (i) the loss over losses applied individual batch elements and time bins of y.
function {(y) is typically implying a block diakg-aa Hessian and therefore cheap products with Jθ; (ii) for standard losses such the mean squared error or the cross-entropy, the diagonal blocks themselves have structure that affords fast computations [32] One might also worry that y may not even fit in GPU memory. of losses accumulated over time, and it is straightforward to similarly accumulate both the overall cost (and its tangents) and the sketched GGN matrix G without ever having to store network activations over time. Overall, this leads to O(T)) memory cost (Figure 6A).
For RNNs, however, E(-) is typically
Second-order methods are known to require appropriate damping of the curvature matrix Finding good absolute damping parameter ? can be difficult without knowing the overall scale of the curvature matrix or the extent of its ill-conditioning. SOFO's lower-dimensional reparameterization of the model each iteration affords us the explicit computation and inversion of the sketched GGN matrix & (K K). In particular, this gives easy access its singular values. We exploit this here by using a relative damping scheme, setting = A8max where Smax is the maximum singular value of G. and À is a relative damping parameter that we have found easier to tune.
33
Connection to real-time recurrent learning (RTRL) The RTRL algorithm [54] is mathematically equivalent to a particular limit of SOFO: the limit where (i) SOFO doesn't use curvature information (i.e. dropping the inverse matrix term in Equation 8 and (ii) the set of tangent vectors used at each iteration (e in Equation 66 is a full basis equal to the identity matrix. In this limit, batched forward- mode differentiation (efficiently) implements the usual RTRL recursion to propagate exact, entire Jacobians network activity w.r.t. the parameters, in (constant memory) forward mode. For model with P parameters, the memory cost of RTRL will be P/ K times that of SOFO if SOFO uses K random tangents. Given that the ratio K/P used in our experiments is around 1 1%, this would make RTRL- 100 times more memory-intensive than SOFO.
4 Results
We now apply SOFO to a range of RNN-based applications relevant to neuroscience, and show that outperforms Adam in all cases, occasionally finding network solutions for tasks where Adam failed.
5

---

A
B
C
Adam
training paradigm
10
SOFO
example trajectory after training
FGD
first
training loss
0.1
0.01
0.001
0.0001
400
800
1200
time (steps) 32
time (steps) 10*
iteration
Figure I: Sparsely supervised Lorenz attractor. (A) In each trial, a 32-step snippet (colored line) of the Lorenz manifold attractor (grey) is chosen at random. The RNN is initialized in the starting state (blue dot), and the final state (green dot) is the only supervision label provided during training. (B) Learning curves for Adam, SOFO, FGD 4 and a first-order version of SOFO that uses the same number of tangents 28). It is possible to get a smoother learning curve for Adam but at the cost of much slower convergence. (C) Example trajectory produced by the RNN after training with SOFO. In this task, SOFO is only 35% slower than Adam in wallclock time per iteration (not shown).
For all tasks, we manually tuned both SOFO and Adam's 's hyperparameters to get best performance in each case, as assessed by the training loss. All experiments are run on an RTX 2080 GPU.
4.1 Inferring dynamics from sporadic observations
We begin by evaluating SOFO's ability to learn the nonlinear dynamics of a system based on temporally sparse state observations - problem related to current neuroscience efforts to infer latent brain dynamics from neural or behavioral data [12). We use the Lorenz attractor [30] as model nonlinear system with non-trive non-trivial dynamics (Figure D. We train an RNN with a 3-dimensional state space, but flexible one-step dynamics parameterized by two-layer MLP with an inverted bottleneck (Appendix E.. this is similar to the neural ODE model class used in 26). In each trial, the network initialized at a random location on the Lorenz manifold, and is trained to generate a state trajectory that, after 32 time steps, terminates exactly where the Lorenz system would have (Figure [[]). Therefore, the RNN is only supervised using trajectory endpoints, being left unsupervised for most of the trial.
In this task, the training loss decays faster, more smoothly. and to a lower minimum with SOFO than with Adam (Figure[B3). The RNN correctly recovers the Lorenz dynamics (Figure
It is remarkable that despite only optimizing a small effective fraction K/ P/Pon of the parameters at each iteration (in this application, K 128 implies K/P < 5%). SOFO still outperforms method that optimizes all parameters at once ( (Adam). We hypothesized that this owes to the use of second-order preconditioning with the GGN matrix. To test this, we compared SOFO
first-order version of 890 (28) This first-order subspace method also sees the same
iee. parameter updates of the form θ0 T dc small parameter subspace at each iteration, but does much worse than SOFO; in fact, in this case, appears to be useless (Figure TB, orange). With K 1 and no second-order preconditioning, SOFO reduces to "Forward-gradiens descent" 4), which performs even worse (Figure 1 green).
4.2 Learning an adaptive Kalman filter
Next, trained vanilla RNN to perform adaptive Kalman filtering (KF) in non-stationary environment (Figure ( - a task highly relevant to adaptive motor control [6 [7]. In this task, the RNN receives noisy observations of the state of an underlying ID linear dynamical system (LDS), and at each time step must infer the current latent state. The parameters of the linear dynamical system (i.e. the context) are subject to sporadic changes during each trial (details in Appendix (E.2). Thus, the RNN must learn to integrate its inputs to (i) rapidly learn about the current parameters of the L.D be able perform optimal KF, and (ii) detect any contextual switches.
SOFO again outperforms Adam on this task, successfully training the RNN in less than 200 training iterations (Figure 2). Adam is much slower, and converges to suboptimal solution. Networks

---

A
B
C
contextual switches
10
Adam SOFO Kalman Kalman random
LDS 1
LDS
LDS N
traming loss
Sonuuuttn error
states
X4. M6 Mg.
XT
predict X4
observations
yr Ms, ya.ys, yo
observe Vi:
200
400
600
800
1000
50
100
iteration
timestep since switch
Learning an adaptive Kalman filter (KF). (A) Task structure, where 2 denotes the latent state and y/ its momentary observation. Each trial (T' 5000 steps) randomly partitioned into successive contexts (of potentially different durations; Appendix E.2). Each context is characterized by a different, randomly parameterized LDS, which produces latent state trajectories noisily observed by the RNN; context switches are uncued. At each step, the RNN must predict the current state. (B) Training curves for Adam and SOFO, compared the MSE noise floor provided by KF knowing the current context (green, average over 1000 trials), and a baseline showing KF performance given random (and thus wrong) LDS parameters sampled from the task distribution (Kalman random; orange, average over 1000 trials). (C) Within-trial evolution of the mean-squared prediction error for the current context the trained models, as function of time elapsed since the last context switch, averaged over 1500 contexts (barely visible shaded areas 22 s.e.m.). Green and orange baselines; same as in (B).
Figure 2:
trained by SOFO are able to (implicitly) infer the current LDS parameters with good accuracy within 10/20 steps following every context switch, gradually approaching the fundamental limit set by a KF that has full knowledge of the active LDS parameters at all times (Figure _ red). Within-context learning in networks trained by Adam is slower and worse overall (black).
4.3 3-bit flip-flop task
To further demonstrate the ability of SOFO Iearn long temporal dependencies, we turn the 3-bit flip-flop task described in [49] (Figure 3A, B; see Appendix E.3 for details). Moreover, we compare the performance of SOFO not only against Adam, but also FORCE (88] (Figure another second- order optimizer based on recursive least-squares (RLS) which is admittedly the optimal solution in the "reservoir" setting. Whilst FORCE does make more rapid initial progress, SOFO eventually achieves a lower loss at convergence. To make the comparison fair to FORCE, we implemented novel batched version of FORCE-RLS (Appendix (B).
Adam SOFO
A
B
D
evaluation loss
SOFO
FORCE
-
0.1
0.1
1
50
100
150
200
200
400
600
800
iteration
iteration
Figure 3: 3-bit flip-flop task. (A) Network architecture for the classical 3-bit flip flop task: three input channels sporadically provide random bits (#1) random time intervals and each corresponding output channel must hold its activity at the value given by the last provided bit (see B). The network feeds back its own output via random feedback weights (grey). Only the output weights (black) trained. (B) Example outputs for the network trained with SOFO (dashed lines); tight match to target outputs (solid pale) shows successful training. Solid dark lines show the corresponding input bits. (C) Evaluation loss for Adam (black), SOFO (red) and FORCE (green), for a network of size 1000 and for SOFO small (purple) where the network size is reduced to 128 but the all sets of weights (i.e. recurrent weights, biases etc) are trainable. (D) Same as C, for varying number of neurons from 1000 8000 (dark to pale).
7

---

reach
propare
B
noisy vanilla RNN
hold
Adam SOFO
traiming less
Tire per ineration
0.1
task inputs
time
0.01
80
100
soFo
Adam
serbory feedback
wallclock time [min)]
example single-neuron PSTHs
single-trial hand trajectories
F
shoulder torque
A d
M
prepare
soFO Adam
EVI
SOFO
MIIV A
Adam
Figure 4: Memory-guided single-reach task. (A) System schematics and task structure (see text). (B) Training curves for Adam and SOFO (mean = 2 s.c.m. over 3 independent runs). (C) Wallclock time per training iteration. (D-F) Example single-neuron firing rates (D), single-trial hand trajectories (E) and shoulder torques (F) for eight selected reach conditions (color-coded). Gray areas denote preparation epochs.
Interestingly, we find that SOFO works best when the reservoir is non-chaotic initialization (g = 0.5 0.5 instead of g = 1.5 as in [49] - whereas FORCE tends to fail in this setting (as well known). Morover, ref. [II] demonstrated that by restricting the locus of learning to those few task-relevant output weights, FORCE tends to come up with brittle solutions. SOFO, on the other hand, is flexible general purpose -rddr optimizer that be used to tweak any of the network's parameters (recurrent weights, biases, feedback and input weights), not just the output weights. We implemented this form of training and found that SOFO was able to train much smaller network of 128 units and ReLu activations (instead of the usual FORCE tanh) on the 3-bit flip flop task (Figure 3C, purple), with 1000 neurons. In summary, by allowing efficient training of more flexible network models from non-chaotic initializations. SOFO stands nicely complementary to FORCE the RNN training toolbox.
final test means squared
more than 15 times smaller than FORCE's
4.4 Motor tasks
Finally, we also applied SOFO to two different memory-guided motor tasks: a single-reach task, and more challenging double-reach task. In both tasks, the RNN receives task-related inputs (see below) and outputs pair of time-varying torques driving the motion of a two-jointed model arm (Figure FA: (.3). The RNN also receives sensory feedback in the form of joint positions and velocities. In the single-reach task (Figure[4) 4). each trial begins with a variable "idle" phase during which the hand must remain a central spot (where it initialized). A variable-length preparation phase follows, during which the network is presented with inputs representing the z- and y y-coordinates of the reach target but must continue to hold to the central spot. The movement phase is initiated by the withdrawal of "hold" input present during the first two phases (c.f. [50].. The RNN must reach the target within 600 ms, and hold it for 200 ms. In the double-reach task, the two target locations are simultaneously presented during the preparation phase, and must be reached in a sequence (see Section (.c4ion/4.4ffe for details of the loss function). Target locations are randomly sampled in each trial from a distribution of reachable positions around the central spot.

---

A
B
noley vanilla RNN
hold
Adam
SOFO
training
task inputs
time
Figure 5: Memory- guided double-reaching task. (A) Task structure
0.1
model arm
0
50
100
150
(see text). (B) Train- ing curves for Adam and
sensory feedback
walldlock time [min]
SOFO. (C) Example
C
target target 2
SOFO Adam V.S..
shoulder torque
single-trial hand trajec- tories for selected pairs of consecutive reach tar- gets. (D) Example
prepure resch 1
reach 2
time series of shoulder torques.
soro
Adam
We trained stochastic vanilla RNN on both tasks. This is difficult task for noisy RNN to acquire, as process noise in the recurrent dynamics not only causes motor disturbances in the arm, but also corrupts the memory of the target location which is longer visible during the reach. A successful network must learn a strategy that mitigates the effect of ongoing noise in the relevant memory / output subspaces, and integrate sensory feedback to correct for output disturbances.
SOFO rapidly trains successful RNNs to perform both tasks proficiently in a few hundred training iterations (between 30min and 2h of wallclock time), where RNNs trained using Adam just about get the gist of the single-reach, and completely fail to learn double reaching (Figure B--E-F and (B- despite days of manual hyperparameter tuning. In the single-reach task, networks trained with Adam fail to withhold movement until the go cue, and appear unable to robustly memorize the target location and/or correct for motor noise.
Excitingly, the single-neuron activity patterns observed the RNN trained with SOFO have a striking resemblance with those recorded in the motor cortex of non-human primates during similar tasks. Some neurons are active during movement only, producing rich multiphasic activity patterns; others active during preparation only, and many are active during both phases. We leave the analysis of these population activity trajectories to future work; for now, we simply note that SOFO renews our ability to train networks complex motor tasks, thus opening many avenues for future motor neuroscience investigations.
4.5 Memory and compute profiling
Details on SOFO's memory and algorithmic complexity can be found in Appendix C To experi- mentally demonstrate that SOFO memory and compute efficient, we carried out profiling in the context of the Kalman filter learning task (Section 4.2), varying the time horizon T used in each trial. As expected from backpropagation, memory usage for Adam increases almost linearly with the time horizon, eventually exceeding the limit of our GPU for T ~ 30K steps [Figure6o44 black). In contrast, SOPO's use of forward-mi moreeed--eeAAA AD incurs a low and constant memory cost independent of T (red). Remarkably, despite being a second-order method, SOFO is only about twice slower than Adam on wall-clock time per iteration (Figure 63; see also wallclock time comparisons Figure 4 and Figure 5B
Limitations of SOFO
All tasks we have used here can be described using relatively small number of bits - much fewer than the information that could in principle be stored in the model's parameters. Accordingly, simply

---

A
B
RTX 2080
limit
Figure 6: Compute time and memory profiling. (A) GPU memory usage a function of time horizon used in the Kalman filter task (Section 4.2). Gray shading indicate the 12GB memory limit of RTX 2080 Ti. (B) Wallclock time per iteration (mean + s.e.m.), for K 256, T 5000, and batch size 256. GPU memory usage is taken from nvtop.
SOFO Adam
15
10
GPU memory (GB)
val clock time iteration
10
5
0
20000 40000 60000
Adam SOFO
time horizon T
instating the right low-dimensional dynamical motifs in the RNN is sufficient to acquire the desired behaviour [9]. Therefore, it is perhaps not surprising that SOFO performs very well even though it only optimizes fraction of all model parameters in each iteration. We speculate that in richer learning tasks that require the RNN to store a lot of information (e.g. supervised learning on a rich dataset), SOFO might not have as large an advantage over Adam. Indeed, our experiments using SOFO to train an MLP-Mixer [52] on CIFAR-10 shows relatively poor performance at best on par with Adam (Appendix (.3)
Even for simple tasks, and assuming fixed (size of the gradient and GGN sketches), we expect SOFO's performance to gradually deteriorate with the number of parameters in the model. This because the sketch of the GGN matrix gets noisier as the tangent to parameter ratio (K/P) gets smaller- meaning that we explore smaller fraction of the parameter space per iteration. Indeed, when training RNNs with increasing number of neurons S, the convergence rate for SOFO worsens while that for Adam improves (Appendix D.2)
6 Conclusion and Future Work
In this paper, we have shown how sketches of both the Generalized Gauss-Newton matrix and the loss gradient be computed efficiently on parallel hardware, leading to RNN optimizer that outperforms Adam in all the neuroscience-relatee tasks we have studied, at lower memory complexity. By accelerating RNN training (and sometimes enabling successful training altogether), SOFO could greatly facilitate a whole line of neuroscientific inquiry that relies on constructing neural networks that solve behavioral tasks. These trained networks constrained for biological realism and reverse- engineered to increase our understanding of brain computations, SOFO's fast and robust convergence in all the tasks we have tried suggests that it could enable faster analysis-by by-synthesis iterations.
Our current implementation of SOFO could be refined in a number of ways. Adaptive damping (e.g. based the Levenberg-Marquaae heuristic 37 (33 is known make second-order optimization algorithms more robust. We speculate that SOFO might benefit from it - although perhaps not to the same usual degree, given that a random sketch of the GGN matrix is typically cetter-conditioned than the full GGN (it is indeed very unlikely that parameter subspaces randomly sampled at each iteration would exactly contain the top and bottom eigenvectors of the GGN).
In principle, SOFO is a general-purpose purpose optimizer that could be deployed onto any problem outside the realm of recurrent neural networks. We expect SOFO to perform particularly well in models that have relatively few tunable parameters yet involve computational graphs that are too large to allow backpropagation to operate large data batches. A prime example of this is the fine tuning of large language models, e.g. with low-rank adaptors 19 but there are many other potential applications e.g. in the physical sciences.
Acknowledgments
This work was funded by Trinity College (graduate studentship to Y.Y.), EPSRC MediaTek Research UK (iCase studentship R.X.) and the Wellcome Trust (Investigator Award in Science 212262/7/18/Z to M.L.).
10

---

References
[1] Omri Barak. Recurrent neural networks as versatile tools of neuroscience research. Current opinion in neurobiology, 46: 1,:116,220 2017.
[2] Omri Barak, David Sussillo, Ranulfo Romo, Misha Tsodyks, and LF Abbott. From fixed points chaos: three models of delayed discrimination. Progress in neurobiology, 103:214 2224222222222 2013.
[3] Atilim Gunes Baydin, Barak A Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic diffcrentiation in machine lcarning: a survey. Journal of machine learning research, 18(153); 1-43, 2018.
[4] Atilm Günes Baydin, Barak A Pearlmutter, Don Syme, Frank Wood, and Philip Torr. Gradients without backpropagation. arXiv preprint arXiv:2202.0885 08587, 2022.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: compos- able transformations of Python+NumPy programs, 2018. URL http::ttt::tiiieii com/ google/joeeeeeoo
[6] Coralia Cartis, Jaroslav Fowkes, and Zhen Shao. A randomised subspace Gauss-Newton method for nonlinear least-squares. arXiv preprint arXiv:221 2tti::222110072.., 2022.
[7] Ricky TQ Chen. Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
[8] Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoder-decoder for statistical machine translation. arXiv preprint arXiv: 1406. 1078, 2014.
[9] Jasmine Collins, Jascha Sohl- Sinl-Dicsstein, and David Sussillo. Capacity and trainability in recurrent neural networks. arXiv preprint arXiv:1611.090 2016.
[10] Basile Confavreux, Friedemann Zenke, Everton Agnes, Timothy Lillicrap, and Tim Vogels. A meta- learning approach to (re) discover plasticity rules that carve desired function into neural network. Advances in Neural Information Processing Systems, 33: 33:16398-16408, 2020.
[11] Brian DePasquale, Christopher Cueva, Kanaka Rajan, G Sean Escola, and LF Abbott. full-force: target-based method for training recurrent networks. PloS one, 13(2)=e011122, 2018.
[12] Daniel Durstewitz, Georgia Koppe, and Max Ingo Thurm. Reconstructing compatational system dynamics from neural data with recurrent neural networks. Nature Reviews Neuroscience, 24(11):693 2023.
[13] Rodrigo Echeveste, Laurence Aitchison, Guillaume Hennequin, and Máté Lengyel. Cortical-like dynamics in recurrent circuits optimized for sampling- based probabilistic inference. Nature nearoscience, 23(9): 1138-1149, 2020.
[14] Robert Gower, Dmitry Kovalev, Felix Lieder, and Peter Richtárik. Rsn: randomized subspace Newton. Advances in Neural Information Processing Systems, 32, 2019.
[15] Filip Hanzely, Nikita Doikov, Yurii Nesterov, and Peter Richtarik. Stochastic subspace cubic Newton method. In International Conference on Machine Learning, pages 4027-4038, PMLR, 2020.
[16] Masahiko Haruno, Daniel M Wolpert, and Mitsuo Kawato. Mosaic model for sensorimator learning and control. Neural computation, 13( 13(10):2201-2220, 2001.
[17] James B Hcald, Máté Lengyel, and Daniel M Wolpert. Contextual inference underlies the learning of sensorimotor repertoires. Nature, 600(7889):4894499 2021.
[18] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 133 1997.
[19] Edward J Hu. Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation large language models, arXiy preprint arXiv:2106.09655 2021.
[20] Dongsung Huh and Emanuel Todorov. Real- tiletimee motor control using recurrent neural networks. In 2009 IEEE Symposium on Adaptive Dynamic Programming and Reinforcement Learning, pages 42 49, IEEE, 2009.

---

[21] Cole Hurwitz, Akash Srivastava, Kai Xu, Justin Jude. Mathhew Perich, Lee Miller, and Matthias Hennig. Targeted neural dynamical modeling. Advances in Neural Information Processing Systems, 34:29379- 29392, 2021.
[22] Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. Advances in neural information processing systems, 26, 2013.
[23] fa-Chu Kao, Mahdieh S Sadabadi, and Guillaume Hennequin. Optimal anticipatory control as theory of motor preparation: thalamo-corticn cortical circuit model. Neuron, 109(9):1567-15 1581,, 2021.
[24] Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximal-gradient methods under the Dalyak-Lojasiewice condition. In Machine Learning and Knowledge Discovery in Databases:
European Conference, ECML PKDD 2016, Riva del Garda, Italy, September 19-23, 2016, Proceedings, Part 1 16, pages 795 811. Springer, 2016.
[25] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv: 04466, 2016.
[26] Timothy D Kim, Thomas Z Luo, Jonathan W Pillow, and Carlos D Brody. Inferring latent dynamics underlying neural population activity via neural differential equations. In International Conference on Machine Learning, pages 5551-1 PMLR, 2021.
[27] David Kozak, Stephen Becker, Alireza Doostan, and Luis Tenorio. Stochastic subspace descent. arXiv preprint arXiv:1904.01145 2019.
[28] David Kozak, Stephen Becker, Alireza Doostan, and Luis Tenorio. A stochastic subspace approach to gradient-free optimization in high dimensions. Computational Optimization and Applications, 79(2): 339-368, 2021.
[29] Jack Lindsey and Ashok Litwin-Kuman Learning to learn with feedback and local plasticity. Advances in Neural Information Processing Systems, 33:21213-21223, 2020.
[30] Edward N Lorenz. Deterministic nonperiodic flow. Journal of atmospheric sciences, 20<2):330-441, 1963.
[31] Valerio Mante, David Sussillo, Krishna V Shenoy, and William T Newsome. Context-dependent computa- tion by recurrent dynamics in prefrontal cortex. nature, 503(7474):78 84, 2013.
[32] James Martens and Ilya Sutskever. Learning recurrent neural networks with Hessian-free optimization. In Proceedings of the 28th international conference on muchine learning (ICML-I1), pages 1033-1040, 2011.
[33] James Martens and Ilya Sutskever. Training deep and recurrent networks with Hessian-free optimization. In Neural Networks: Tricks of the Trade: Second Edition, pages 479-535. Springer, 2012.
[34] James Martens, Jimmy Ba, and Mau Johnson. Kronecker-factored curvature approximations for recurrent neural networks. International Conference on Learning Representations, 2018.
[35] Luke Metz, Niru Maheswaranathan, Jeremy Nixon, Daniel Freeman, and Jascha Sohl- Dii-kteiik... Under- standing and correcting pathologies in the training of learned optimizers. International Conference on Machine Learning, pages 4556-4565. PMLR, 2019.
[36] Jonas Mikhaeil, Zahra Monfared, and Daniel Durstewitz. On the difficulty of learning chaotic dynamics with RNNs. Advances Neural Information Processing Systems, 35:1 1297-11312, 2022.
[37] Jorge Moré. The proceedings of the biennial Conference held Dundee, June 28-July 1, 1977, pages 105-116. Springer, 2006.
algorithm: implementation and theory. In Numerical analysis:
[38] Yann Ollivier, Corentin Tallec, and Guillaume Charpiat. Training recurrent networks online without backtracking. arXiv preprint arXiv:1507.07660 2015.
[39] Chethan Pandarinath, Daniel J O' O'Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D Stavisky, Jonathan C Kao, Eric M Trautmann, Matthew T Kaufman, Stephen I Ryu, Leigh R Hochberg, al. Inferring trigle-trial neural population dynamics using sequential auto-eneoders. Nature methods, 15(10)880 815, 2018
[40] Razvan Pascanu and Herbert Jaeger. A neurodynamical model for working memory. Neural networks, 24 12):199-207, 2011.
[41] Mert Pilanci and Martin J Wainwright. Newton sketch: A near linear-tima time optimization algorithm with qunear-quadratic convergence. SIAM Journal Optimization, 27(1):205-25 245, 2017.
12

---

[42] James S Plank. Micah Beck, Gerry Kingsley, and Kai Li. Libckpt: Transparent checkpointing under unix. Computer Science Department, 1994.
[43] Zheng Qu, Peter Richtárik, Martin Takác, and Olivier Fercoq, Sdna: Stochastic dual Newton ascent for empirical risk minimization. In International Conference Machine Learning, pages 1823- 1832. PMLR, 2016.
[44] Marine Schimel, Ta-Chu Kao, Kristopher T Jensen, and Guillaume Hennequin. iLQR-VAEE control-based learning of input-driven dynamics with applications to neural data, International Conference Learning Representations, 2022.
[45] Dominik Schmidt, Georgia Koppe, Zahra Monfared, Max Beutelspacher, and Daniel Durstewitz. Identity- ing nonlinear dynamical systems with multiple time scales and long-range dependencies. arXiv preprint arXiv: 1910.0341,, 2019.
[46] Wayne Soo, Vishwa Goudar, and Xiao Xiao-Jing Wang. Training biologically plausible recurrent neural networks on cognitive tasks with long-term dependencies. Advances in Neural Information Processing Systems, 36, 2024.
[47] David Sussillo. Neural circuits as computational dynamical systems. 156 163,163,2220 2014.
Current opinion in neurobiology. 25:
[48] David Sussillo and Larry F Abbott. Generating coherent patterns of activity from chaotic neural networks. Neuron, 63(4):544-555 2009.
[49] David Sussillo and Omri Barak. Opening the black box: low -owddmmnsiooal dynamics in high- -igh-dimensional recurrent neural networks. Neural compittation, 25(3):626 649, 2013.
[50] David Sussillo, Mark M Churchland, Malthew T Kaufman, and Krishna V Shenoy. A neural network that finds naturalistic solution for the production of muscle activity. Nature neuroscience, 18(7):0025-1000 2015.
[51] Ilya Sutskever. Training recurrent neural networks. University of Toronto Toronto, ON, Canada, 2013.
[52] Ilya O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung. Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. Mlp-mixer: An all-mip architecture vision. Advances in neural information processing systems, 34:24261 -22222, 2021.
[53] Elia Turner, Kabir V Dabholkar, and Omri Barak. Charting and navigating the space of solutions for recurrent neural networks. Advances Neural Information Processing Systems, 34:25320-25333, 2021.
[54] Ronald Williams and David Zipser. A learning algorithu for continually running fully recurrent neural networks. Neural computation, 1(2):270-280, 1989.
[55] Stephen Wright. Coordinate descent algorithms. Mathematical programming, 151(1)33-34, 2015.
13

---

A
Additional Experiments
In this section we describe some additional experiments and again compare the performance of SOFO with Adam.
A.1 One-shot classification
We also tested SOFO difficult one-shot classification task (see Appendix E.5 for details). This task relies on three Gaussian distributions of 3D inputs: pi(u) = (wu; e,, where the 3 x identity matrix and denotes its 1t column. In each trial, these three input distributions are randomly assigned a unique class label E {1, 2, 3), and the network begins the trial with a phase of exposure to sequence of three input/output pairs ("support set" in meta- meta-learning jargon), each sampled from different class and presented for Tieam = 10 consecutive time steps. In a second phase, the network is presented with three other fresh inputs from each class ("query set", Tespion 10). in random order, and is asked to classify them according to the same input/class contingencies (Figure 7 With these contingencies heing randomized in every trial, the RNN must the initial exposure phase to learn from the support set, and remember and exploit this information to accurately classify the query set and do all this using its own dynamics with fixed parameters.
SOFO outperforms Adam on this task too, successfully training an RNN to 90% one-shot classification accuracy (later 100%) in 150 training iterations, compared to 1500 iterations for Adam. Contrary to SOFO, Adam's 's training loss exhibits the usual plateau characteristic of first-order optimization of loss functions with pathological curvature (Figure BB).
A.2 Delayed addition
To demonstrate that SOFO also train networks on tasks that require complex temporal dependen- cies, we apply SOFO to the delayed-addditon task (Schmidt et al. [45]; see Appendix f.6 details) with varying sequence lengths (Figure [). There, SOFO again outperforms Adam tasks involving both short sequences (Figure (A and long sequences (Figure BB).
B
input
classifiction accuracy
Figure 7: One-shot
1.2 0.3 0.4
classification task. (A) Task structure (see text). (B) One-shot classification accuracy during the course of training.
output
Adam SOFO chance level
label
D
0.6
0
1
0.4
within Irial: learning exploitation
signal
300
1200
iteration
A
10
T=32
B
10'
T-512
C
101
T-128
Adam
109
109
109
K=256
training loss
10-1
SOFO
Adam
10-
SOFO
Adam
10-1
K 128
K-84
10-2
10-2
10-2
K-
10-3
10-1
10-a
10-4
10-4
10-4
10-5
10-8
10-5
2000
4000
5000
8000 10000
2000
4000
6200
8000 10000
2000
4030 6000 8000 10000
iteration
iteration
iteration
Figure 8: Delayed-addition task. (A) Training losses for time horizon = 32. (B) Same as (A), for T 512. (C) Same as (A-B), for T 128, but varying the number of random parameter tangents (K) in SOFO. In both (A-B), we had set K = 256.
14

---

B Batched-FORCE
We also extend the FORCE learning algorithm 48 to its batched version, which we used for the flip-flop task in Section 4.3 and Figure Recursive least squares (RLS)-based FORCE learning. as presented in [48]. assumes a batch size of I, which does not fully exploit GPU parallelism. The batched FORCE algorithm presented in Algorithm iimplements FORCE learning on whole batch of trials in parallel.
Algorithm Batched FORCE
1: hyperparameters: a; batch size M; number of neurons S; output dimension N; time horizon
T'
2: input: initial inverse covariance matrix Po _ initial output weights Wo TR R and initial
network state E SMM.
3: for m = 0, 1, 2,... do
b iterate until convergence
4:
sample data
with targets {Te e RNNM
5: for = 0, 1, , do
b iterate over the horizon length
6: obtain from 2 and y(t) (Equation 15
7:
Tt ((2))
8:
Y/ Wir.
9:
G = yt 8(())
b calculate error update p b update W
P. = P:-1 11-1-PP-1 Pi- (MI + + Pe- P1-1
Pr-
11:
We+1 = We M e Pt.
end for
13: Po e Pr
output: TT
15: end for
C Memory and runtime complexity for SOFO
This section states the memory and runtime complexity of SOFO when applied to RNN with state dimension (assumed to be larger than the output dimension N), minibatch size M, time horizon T and tangent batch size K.
Memory complexity For a recurrent neural network operating in batched forward-mode AD, what has been computed at time can be immediately discarded time f++ 1, leading to ( complexity where 7 is the time horizon. At each time step, the memory cost is dominated by the largest tangent tensor being manipulated, which will typically be the batch of tangents associated with the network's activity, of shape S M x K. For large networks with S M, the dominant contribution will be from batch of tangents associated with the recurrent weight parameters (θ) of shape S S K. Note that the accumulated GGN sketch G of size K K, makes a negligible contribution to memory usage. Thus, SOFO's memory complexity is O(max(S, M)SK), to be compared to O(S2 + +ST) for standard reverse-mode AD.
Runtime complexity SOFO's per-time-step runtime complexity is dominated by the temporal updating of the batch of activity tangents; for RNNs, this involves batched matrix multiplications which have complexity (SS KK.. Runtime complexity also linear in the time horizon, leading complexity. Whilst this is technically K times more than the O MT)
to an overall D(S2MKK runtime complexity of reverse-mode AD, efficient batching of those K multiplications on GPU is such that, in our experiments, SOFO was only small constant (typically between 1.5 and 3) times slower than Adam terms of wallclock time per training iteration (c.f. Figure B3). Note also that SVD of the GGN sketch (Algorithm) |] line. 10) will incur negligible O ( runtime cost unless K very large (which typically isn" 't the case anyway, because of memory constraints see 'Memory complexity" above).
15

---

Train
Loss
Accuracy
Accuracy
0.7
0.7
saro (0.10%)
1. 1.000
1.000
sara (0.52%)
0.5
0.5
SOFO (1.20%)
Adam
0.550
0.550
0.3
0.3
5.900
1500
6.1
0.3
0.875
0.873
0.850
30
100
130
200
0
50
100
150
200
0
50
100
150
200
0
50
100
150
200
Reration
Iteration
korabon
Iteration
Figure 9: MNIST classification task. Training and test performance (see titles) for Adam (black) and SOFO (colored) with different numbers of tangents K. Percentages refer to the K/ ratio, i.e. number of tangents K to total parameter count P.
D Additional profiling experiments
Here we present additional profiling experiments on SOFO to demonstrate how SOFO performs when we (1) increase the number of tangents K"; (2) increase the number of neurons S and (3) scale up to deep networks.
D.1 Investigating the effect of K in SOFO
The choice K affects the quality of SOFO's gradient and GGN sketches the larger, the better. In general, K should be chosen to be small enough such that the training is not too memory-intensive, but large enough to ensure fast convergence. Our experiments with SOFO have shown that performance mainly depends the ratio of K to the number of parameters P. Strikingly, in RNN settings, this K/ P ratio can be chosen as low as 1% and still yield excellent training performance.
Here we carry out two further experiments to demonstrate the effect of K on SOFO-based training.
MNIST classification We investigated the effect of K on the training of two-layer perceptron (MLP) on MNIST classification. With hidden size 100, and input size x 28 784, the network had a large total parameter count of P = 79510. As we increased the ratio K/P P from 0.16% to 1.26%, training became gradually better, with SOFO eventually outperforming Adam in terms of progress per training iteration (Figure
Delayed addition We also investigated the impact of K in the training of an RNN on the delayed- addition task (Appendix A..)). There, the vanilla RNN with 128 neurons had smaller parameter count P = 17024 (see network structure in Appendix E.6). In this case, even using only K = 32 tangents (K/P ss 0.2%) is sufficient for SOFO to outperform Adam, and increasing K leads to even faster convergence (Figure 8C).
D.2 Training SOFO on large networks
In all previous experiments SOFO has mostly been employed to train small networks (i.c. to 1000 neurons in the flip-llop task in Section 4.3 In Figure 3D, we compared the performance of SOFO and Adam in the context of the flip-flop task, as we increased the number of neurons S from 1000 8000. While increasing S improved the performance of Adam in terms of convergence speed, actually hindered SOFO. In light of the discussion in Appendix D.I. this is perhaps not surprising, as we kept the number of tangents K fixed at 256, such that the K/P ratio decreased quadratically with growing S. In other words, with growing S, SOFO got explore diminishing fraction of the parameter space in each training iteration. Note that SOFO nevertheless outperformed Adam for all network sizes in this benchmark.
16

---

2.0
2.0
SOFO 00 (0.27%)
0.7
SaPD (0 35%5%
SDFD (1.00))
Adem
0.5
0.3
0.4
1.-
12-
0.3
0.3
1.0
1.0
0.2
0.2
Iteration
(4)
Figure 10: Training an MLP-mixer on CIFAR-10 classification, (A) Training loss. Percentages refer to the K/P ratio, i.c. number of tangents K total parameter count P. (B) Evaluation loss. (C) Training accuracy. (D) Evaluation accuracy.
D.3 Training SOFO on deep networks
We have also added comparison between Adam and SOFO MLP-Mixee architecture [52] trained on CIFAR- 10 classification (~ 10y parameters; see details in Appendix E.7 Figure 10). This a fairly easy target for Adam due to the presence of skip connections and layer-norm operations, and with the limits imposed by GPU memory, the best SOFO does here is to nearly match Adam's performance. Note that this deep network setup not where expect SOFO to be the most useful anyways, as this network not deep enough to really challenge reverse-mode AD in terms of memory complexity.
E Details on experiments
In this section we present the details on main experiments, describing the task setup and network structure.
E.1
Details on Figure
Task structure
The task consists learning the dynamics Lorenz attractor with minimal
supervision. The Lorenz attractor consists of three-dimensional state (x) 22; 2z3) evolving
8
④ 10(x2 z1) 22 x (28 13)
21
i
21
-3
3
Data is generated by integrating Equation over with time step of 0.01 using a Runge-Kutta solver (RK4) followed by z-scoring. During training, the network runs for total of Tiorenz 32 time steps but the only label provided is the system's state at the final time step.
Network structure We use simple RNN network with inverted bottleneck structure for this and directly used as the
task. The hidden state z of the recurrent network, which is state readout, evolves
2c+1 Az, + Wo(Cz, b)
(10)
where A, C, W, b are parameters and ò the ReLU function. The size of the inverted bottleneck is chosen to be 400 (i.e. W has shape 400 x 3). We use the squared error (SE) between the model prediction and the label as the loss function for training of the network.
E.2
Details on Figure
Task structure The task consists of learning perform adaptive Kalman filtering on session comprising of multiple linear dynamical systems (LDSs)) (i.e. the environment is non-stationary).
The state transition dynamics of an LDS is given as:
L++1 az, + 0
(II)
17

---

where ny is the hidden state at timestep t. The state emission dynamics is given by:
Y/ = Zt + o
(12)
where yt is the observation at timestep (Figure 2 is the state transition parameter, b is the stationary mean of the process and v is the drift at each time point and obtained as v = (1 a)6. of is the process noise variance and 0 is the observation noise variance. e" and eM are both drawn Li.d. from a standard Gaussian distribution A70, 1). In particular, o² is related the stationary variance of the LDS, which we denote o², via o² = a² )²)
Each session has horizon length of T timesteps and consists of number of LDSs: an LDS survives for the first ix number timesteps, and from Tix beyond at every step there is a 1/d probability of transitioning onto new LDS, where d is denoted as the survival time. Before the start of each LDS, we first sample its parameters (i.e. a, 6, oz) from a common global distribution, before sampling the first data point for this particular LDS from (b,, o²). Hence, cach session is characterized by number of LDSs whose transitions are iid. The values parameters and the global distributions from which the parameters sampled are presented in Table T
Given the data emitted by the LDSs as described above, the task is as follows: at each time step t, the neural network receives a new observation yr, and outputs a prediction of the current latent state dr. The training loss a simple temporal accumulation squared prediction errors. The network hence trained to amortize inference (approximating p(rr(11)))) in this non-stationary Kalman filter setting.
Network structure We use a vanilla RNN with for this task. The momentary hidden state z, of the recurrent network evolves as:
2++ = (1 a)zz + a(o(Cz, + Byr + +e)
(13)
where B, C, e are network parameters and @ is the ReL.U nonlinearity. The estimated current state obtained as:
Oz:
(14)
where O is the readout matrix. The number of neurons 10 (dimension of 2) and a = 0.25. Network parameters are initialized to reproduce the current observation (i.e. ft 11). We use the squared error between the predicted current state d and the actual current state I, averaged over T timesteps the loss function.
Kalman filter and Kalman random baselines The KKlman filter' baseline has access to the parameters of each LDS (i.e. a, b, 0², oz)) and is hence aware of each context change. The 'Kalman random' baseline performs the same updates as the Kalman filter, except for that it has knowledge of b but not of the other parameters o², of). Instead, it assumes random values for these parameters, drawn from the same global distribution the one used to generate the parameters for each LDS (Table T).
E.3 Details on Figure
where three input Task structure We conduct a 3-bit flip-flop task as in Sussillo and Barak 449,
channels sporadically provide bits (#1) random time intervals while remaining at zero other
Table 1: Values of LDS parameters and global distributions for the sampling of LDS parameters in learning adaptive Kalman filter task (Section 4.2) p(x) denotes the distribution that is drawn from and (a, b) denotes the uniform distribution between a and b.
parameters
distribution/values
T
5000
Tax d
100
100
p(b)
u(-3.3)
P(T)
41,
p(ag)
4(0.1,3)
P(a)
1
18

---

times. Each corresponding output channel must sustain its activity (#1) at the value given by the last provided imput bit. In other words, the network must maintain at least 2 8 bit memory any time instance, while ignoring the cross-talk between differing input/output pairs.
Network structure
The hidden state z of the recurrent network evolves as:
-++1 = (1 - az + (<(((Zz, + By + WFn(s))
(15)
where J, B, WrB are network parameters, y is the output and @ is the ReL.U nonlinearity. The output is obtained as:
a
Wo((z)
(16)
where Wi is the readout matrix. Note that to similar to FORCE learning, we fix all weights and only train output weights W. In particular, entries of J are drawn from Gaussian distribution (0 $) where g > 1 initializes the network in the chaotic regime (which is necessary for FORCE learning) but g < initializes the network in the con-chaotic regime. We use S 1000 neurons and mean squared error between the target trace and output trace as the training criterion.
E.4 Details on Figure and Figure
Task structure See main text for an overall description of the task. The details of the two-jointed arm model used in these tasks can be found in [23]. The central spot is about 20cm in front of the shoulder joint. For the single-reach task, the duration of the idle phase is drawn uniformly between 0.1 and 0.5s in each trial, and the duration of the preparation phase is drawn uniformly between 0.2 and 0.8s. Each trial lasts for 2 seconds regardless of whether the target has been successfully reached. During the idle & preparation epochs, the loss function is a quadratic penalty on the squared deviation of the hand position from the central spot and quadratic penalty on the hand velocity. During the movement epoch, there is no loss for the first 600ms, and then the deviation of the hand from the desired reach target penalized quadratically, together with the hand velocity.
For the double-reach task, each trial begins with the preparation phase right away, whose duration is sampled uniformly between 0.2 and 0.6s. The loss is similar to that described above, with one more term for the second target to be reached.
For the single-reach task, the input u!() is 3 3-dimensional, the first two channels relaying target information and the third conveying the go cue (50). For the double-reach task, the input is dimensional, as two more dimensions are needed for the second target.
Network structure The stochastic vanilla network dynamics considered here read:
dz = - + (6(()) @ (1 + oε)) ) + b Bu + F(0,, 02, TT)
(17)
dt
where (01, 02) and their derivatives denote the Markov state of the two arm joint angles. Here, W is matrix of recurrent weights, b is bias vector, B is the input weight matrix, and F is matrix of feedback weights. We set r = 20 ms, and use networks with 64 neurons. The nonlinearity d(.) is the ReLu function, and "firing rates" are corrupted by multiplicative noise (1 ae) where is an N-dimensional random vector of uniformly sampled elements between -0.5 and 0.5.
The network dynamics are integrated using the Euler method with a time step of 2 ms.
E.5 Details on Figure
Task structure The task is a one-shot three-class classification. We define three Gaussian distribu- tions of three-dimensional (3D) inputs: pi(u)= = Nww;; e,, 1) where is the 3 x 3 identity matrix and e, denotes its 7hh column. In each trial, these three input distributions are randomly assigned unique class label E 2.3). Each trial consists of two phases: the learning phase and the exploitation phase. During the learning phase, the network receives set of three components (the support set): the input is a randomly generated example, the label is the corresponding o encoded class label and teaching signal of value 1. The same set is presented for a fixed Tiem = 10 timesteps before the next set is presented for another Team timesteps where the input is generated from different class, until all three classes have been presented. During exploitation the network receives a set of two components (the query set): the input is a randomly generated example and a teaching signal of value
19

---

O (all entries in the label is set to zero). Similar to the learning phase, the input is an unseen example drawn from one class is presented for fixed Tesplon 10 timesteps, before another example drawn from a different class is presented for another exploir timesteps, until examples from all three classes have been presented. The goal of the neural network predict the class the last step of Tesplii for each of the 3 unseen examples during the exploitation phase (Figure 7
Network structure We use simple RNN with an inverted bottleneck for this task. The hidden state of the recurrent network evolves as:
- Az + Wo(Cz, + + + By: +
(18)
where A, B, C, e. are network parameters, á is the ReLU nonlinearity and yr the input and teaching signal concatenated together at the time step t. The predicted next step is obtained
y = Oz.
(19)
where the readout matrix. The number neurons is 100 (dimension of x) and the size of the inverted bottleneck is 400 (W is of size 400 x 10). We use the cross-entropy between the true label and the network output as the loss function, averaged over the c = 3 predictions that the network makes during each session.
E.6 Details on Figure
We conduct delayed addition task as Schmidt et al. 45] The input consists of two sequences both of length T: the first sequence contains entries drawn from uniform random distribution U[0, while the second sequence contains zeros except for one at two random timings: t1 < 10 and chosen such that every trial requires a long memory of at least
1 and T/2 time steps. At the last time step T, the network is expected to generate an output that the sum of the two numbers in the first sequence at t and f22
t2 < T/2. Constraints
We use vanilla RNN (same structure as that Appendix E.2) with 128 neurons to learn the task. We use the mean squared error between the target and output as the training criterion.
E.7 Details on Figure 10
We use SOFO to train five-layer MLP-mixer [52] on CIFAR-10 classification task. The network consists of a patchification layer (patch size 4 x 4) followed by a per-patch linear embedding to higher channel dimension of 128. This then followed by two mixer blocks, each block including one token-mixing MLP with dimension 64, and one channel-mixing MLP with dimension 128 (each MLP uses GELU nonlinearity). Finally, a fully-connected output layer performs classification. Other components such as skip-connections, dropout and layer norm the channels are also included. This adds up to about 10 parameters.
F Hyperparameter values
Here present the ranges of hyperparameter values tested for Adam and SOFO in Table and Table resppeciiyyyy all experiments carried out. The final values of hyperparameters used in the experiments for Adam and SOFO presented in Table and reesecerepeyttively. We finally present the range of hyperparameters tested and the final hyperparameters for FGD and first-order SOFO for the Lorenz task (Section 4.1p and for batched FORCE in the 3-bit flipflop task (Section 4.3 in Table 6
20

---

Table 2: Hyperparameters used for Adam. is to the learning rate.
batch size
1
Lorenz
256
2 10
Kalman
256
10 10m2
one-shot
256
10-3
single reach
256
10-3
double reach
256
10-3
flip-flop
32
10
delayed addition
256
10-4
MNIST
512
9 10
cifar-10
256
2
10
Table 3: Ranges of hyperparameters explored for Adam
batch size
?
Lorenz
256
10
10
Kalman
256
10
10
one-shot
256
10-5 10-3
single reach
256
10 10
double reach
256
10 10
flip-flop
32
256
10 - 10-2
delayed addition
256
10 10-2
MNIST cifar-10
512
10 10-2
256
10
10
Table Hyperparameters used for SOFO. is the number of tangents used, is the total number of network parameters, is the learning rate and is the relative damping applied on SOFO. Note that for experiments where range of K is used the corresponding entry in the table is left as empty.
batch size
K
K/L P ratio
1
Lorenz
256
128
4.6%
0.3
10 5
Kalman one-shot
256
256
2.5%
0.05
10
256
256
2.8%
0.1
10-3
single reach
256
256
5.4%
0.1
10-6
double reach
256
256
5.2%
0.1
10-6
flip-flop
32
256
2.5%
0.1
10-5
delayed addition
256
0.1
10-5
MNIST
512
0.8
10-7
cifar-10
256
0.5
10-7
Table Ranges of hyperparameters explored for SOFO
batch size
K
?
Lorenz
256
128
10-2
0.7
10 10-3
Kalman
256
256
0.5 1
10 10
one-shot
256
256
10-2
0.2
10-4-10-
10-2
single reach
256
256
10-3 -
10 - 10
double reach
256
256
10-3 1
10 10 10
flip-flop
32-256
256
10 10-2-1
10 -
delayed addition
256
10 0.1
10
MNIST
512
0.1 -
10 50-5-10-
10
cifar-10
256
0.1
10-6
21

---

Table 6: Range of the learning rate 1/ explored and all final hyperparameters used for FGD 4
first-order SOFO [28] for the Lorenz task (Section 4.1) and batched FORCE (Appendix B for the 3-bit flip-flop task (Section 4.3) Range of refers the range of the parameter explored.
batch size
K
7]
range of 1)
FGD
256
10
10
first-order SOFO
256
128
10-2
10-6
10-4
batched FORCE
256
0.1 - 100
22

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction accurately reflect the paper's contribution which is tackling the multiple challenges of training RNNs over long horizons, and the scope is focusing on serving the computational neuroscience community.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The suitable tasks for the proposed optimizer those involving relatively few trainable parameters but memory-hungry. The computational efficiency is discussed in Section 4.5 and also tested in Creffig:motorC and Figure 5B. A formal discuss of limitation can be found in the Discussion.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
23

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification:
We do not have any theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the suppiemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yes]
Justification:
The proposed algorithm is clearly presented Algorithm | in detail, matching exactly with the releasing code. For each simulated task, the used network structure, the hyperparameters, and the way generate datasets all included in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: Data and code are publicly accessible upon publication, including data genera- tion, and training scripts to reproduce all experimental results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm cc/
publit//giide//Cdeetttbmiss wnt.ssamiisioooooooooooooooooiie for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
//nips /nipsscc/pput cttp:/scccpubblcccuiiiiie guides /ttpp///bb///.....o/////////aaaa
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: All training details of the tasks are included in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The 2-sigma error bars of losses are shown for 3-5 runs with different random seeds as shown Figure 2C and Figure 4B.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
25

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Computer resources in type of compute workers, memory, time of execution are investigated in Section 4.5 (Figure 6). Time of execution as in wallclock time is also compared for motor reaching tasks.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: htt.g//rgr/////oooooo.o/.//occha
Answer: [Yes]
Justification: The paper conforms with the NeurIPS code ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper proposed generic algorithm for optimizing RNNs particular for
the computational neuroscience community and there is negative societal impacts.
Guidelines:
26

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g.. gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper proposed optimization algorithm and poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification:
The paper does not use any existing assets.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
27

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
batasets.papersstt
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out the asset's s creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
The dataset and code will be well documented with details about training and license.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details of the dataset/code/mooel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
28

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
29


Reviews:
[{"review_id": "kuakAi7h4P", "reviewer": "Reviewer_535W", "summary": "In this work, the authors develop a method to train recurrent neural networks with biological plausibility constraints, which are required for neuroscience applications, and where standard BPTT doesn't tend to work well. They use forward-mode gradient calculation (RTRL) that makes the memory requirement independent of length of input sequence. To mitigate the high compute cost of RTRL, they combine it with a second order optimizer that is parallelizable on a GPU and makes the overall training wall-clock time comparable to BPTT. The second order optimizer uses a randomized method to calculate the Generalized Gauss-Newton matrix. They demonstrate their method in a number of neuroscience inspired tasks.", "questions": "- It's not clear why a stochastic RNN was chosen for the motor task.\n- In line 90, the authors say \"In RNN applications, i indexes both time bins and RNN output dimensions\" which doesn't sound quite right.", "limitations": "The authors discuss the limitations clearly.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The overall problem domain is well defined and well motivated. Neuroscience does have constraints in terms of network architectures that can be used, and so it makes sense to improve plain backprop even if it comes at an additional computational cost (which in this particular approach is not an issue).\n\nTheir use of forward mode to handle commonly occurring long input sequences makes sense. The choice of GGN updates with a randomized algorithm is also reasonable.\n\nOverall their paper is well written and structured. The the experiments are well chosen.", "weaknesses": "I would have liked to see a more detailed ablation study to understand the contribution of the various components of the algorithm (e.g. choice of K, size of network)\nAnd perhaps a LSTM/GRU baseline just to understand the gap between vanilla RNNs and architectures used in ML (although this is not critically relevant for this study).\n\nI would also have liked a bit more discussion into why vanilla RNNs are the right architecture for neuroscience, since at best they would be crude simplifications of the complex dynamics of biological neurons. And I would expect its rate-based representations don't provide a good model for brain areas that tend to use more spike time based representations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your thoughtful review; we are glad you found the paper well written and structured. In the main rebuttal above, we have addressed the concerns common to all reviews, including a more systematic study of the effect of $K$ (number of random parameter tangents used to sketch the GGN update) on SOFO's performance relative to Adam. Given the time constraints of this one-week rebuttal period, we have performed this study in the context of a new task (delayed addition) that was suggested to us by Rev. Z9Qd, but we will progressively be adding similar experiments for some of the other tasks we used in the original submission.\n\nWe haven't had the time to do comparisons with LSTMs/GRUs, but we can definitely get started on this after the rebuttal deadline if you think this would add value. We expect that for LSTMs/GRUs, SOFO will have a reduced advantage over Adam, and that those networks will perform overall better than their vanilla counterparts.\n\nRegarding the choice of vanilla RNNs as models for neuroscience: as you rightly point out, neuroscience isn't yet at a point where there is a consensus on _the right architecture_ to describe the dynamics of neuronal networks. In this context, many studies have used vanilla RNNs as an Occam-razor-like compromise: simple integration of input firing rates appears to be what neurons do to a decent first order approximation, and until we have more clarity on whether and how neurons might be implementing more sophisticated mechanisms such as forget/input/output gates, sticking to simpler models appears sensible to many in the field. Thus, we see SOFO as relevant to neuroscience primarily because of the pervasiveness of vanilla RNNs in the field, not necessarily because they are the most accurate models. Indeed, we agree that these simpler models are crude, and it will be important to develop methods for training more biologically realistic (e.g. spiking) variants as they become more mainstream. We will add a paragraph of discussion along these lines, thanks for prompting us to do so.\n\n> not clear why a stochastic RNN was chosen for the motor task\n\nWe wanted to have at least one example of a stochastic RNN to cover more diversity of dynamics, as stochastic vanilla RNNs are also common in the literature, for example to study across-trial variability and its relationship with computation (e.g. Echeveste et al., 2020).\n\n> In line 90, the authors say \"In RNN applications, i indexes both time bins and RNN output dimensions\" which doesn't sound quite right.\n\nWhat we meant was that SOFO optimizes general loss functions that are summed over both RNN output dimensions (e.g. target behaviour of the RNN in a certain number $D$ of output channels) and over time. We could have used two distinct indices, say $j$ and $t$, and have a triple sum in Eq. (1), but since these two dimensions are completely exchangeable from the perspective of the SOFO algorithm, we thought we might as well collapse them into a single index ($i$, which then runs from 1 to DT). Another reason we simplified the notation this way is because SOFO should also be applicable to feedforward (non-recurrent) networks for which there is simply no time index -- the way Eq 1 is written, it still applies to such cases. We will clarify this in the text.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GApJC7c2aV", "reviewer": "Reviewer_mSGm", "summary": "In this work, the authors applied a second-order optimization with an approximated Hessian to train vanilla recurrent neural networks (RNNs). The authors have shown that the weight update using a random subspace approximation of Hessian can be implemented onto a GPU in a memory-efficient manner by using the Jacobian-vector product. They applied the algorithm to Lorenz attractor, adaptive Kalman filter, one-shot learning, and motor reaching tasks, and demonstrated that the proposed algorithm achieves better performance than the Adam optimizer while achieving comparable wall-clock time.", "questions": "Evaluations in the manuscript were mostly conducted in networks with a small number of neurons (N=64) and a small time-bin (2ms). Do you expect the proposed approach to also benefit larger networks with coarser discretization of time?", "limitations": "Related to the question above, I believe some comments on the hyperparameter choices would strengthen the clarity of the paper.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The manuscript is written clearly and technically sound. The authors implemented the proposed algorithm in several tasks including biologically-motivated ones. The presented results clearly demonstrate the advantage of the proposed algorithm over the Adam optimizer in their experimental setting.", "weaknesses": "The contribution to both machine learning and neuroscience are rather weak. The Gauss-Newton method with random subspace estimation was studied previously (Gower et al., NeurIPS 2019; Cartis et al., arXiv, 2022), and efficient GPU implementation of forward gradient-based optimization methods was also developed before (e.g., Baydin et al., arXiv, 2022). While its application to RNN is novel to the best of my knowledge, the authors did not provide any analytical results, for instance, on how sketching of Hessian estimation and noise in RNN dynamics corrupt the weight updates. The implementation did not take advantage of the simplicity of RNN architecture either. \n\nRegarding contribution to neuroscience, while the authors focus on vanilla RNNs due to the biological implausibility of gated recurrent units and LSTMs, their proposed learning algorithm also lacks biological plausibility. Although they mention a “striking resemblance with those recorded in the motor cortex,” this resemblance is not quantitatively evaluated. Additionally, it remains unclear if the proposed algorithm converges to a solution more similar to those implemented in the motor cortex compared to previous optimization algorithms.\n\nWhile the authors motivate this work with ‘reverse engineering’ approach for studying neural circuits, its potential contribution to neuroscience is not explained clearly. It would be helpful if the authors could explain what kind of biological questions they aim to address with this algorithm.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review; we are glad you found the paper clearly written and technically sound and that you found our experiments convincing. In the main rebuttal above, we have addressed the concerns common to all reviews, including your comment on the effect of network size. Here we would like to address your other concerns.\n\nRegarding your novelty concerns: to the best of our knowledge, we are the first to really take advantage of GPU parallelism to implement random-subspace Gauss-Newton in an efficient manner (where previous works on second-order sketching had used (i) sequentially evaluated finite-differences approximations of directional derivatives (Cartis et al., 2022), or (ii) memory-expensive reverse-mode differentiation (Gower et al.)). We find that parallelising Jacobian-vector products is really key to outperforming Adam (or other first-order optimizers) on _wall-clock time_ in all our experiments. We achieve this through a custom implementation of _batched_ forward-mode differentiation, which was in itself a challenge as it is not natively available in common GPU frameworks (e.g. pytorch). We will of course release this library on github together with the paper.\n\nOur work also improves substantially on Baydin et al.: their work on Forward-Gradient Descent (FGD) only used a single tangent vector to estimate the gradient, and therefore only required standard (non-batched) forward-mode AD. What we do instead is compute $K$ directional derivatives in parallel, leading to reduced stochasticity in our gradient and GGN estimates. As our experiments show, even in the easiest task of Figure 1B (Lorenz attractor), FGD fails; in fact, it even fails when extended to the large $K$ setting, and SOFO rescues this by incorporating the GGN matrix. In summary, while previous literature had indeed laid the groundwork for our study, SOFO uniquely expands and combines these previous strands to yield an efficient new method for training RNNs.\n\nRegarding analytical results: we agree that this was not the focus of our work; for this, we rely on SOFO inheriting the theoretical guarantees previously derived for sketching algorithms in general, not just for RNNs (e.g. Cartis et al., ICML 2020), and focus on showing empirical performance.  \n\nRegarding neural activity in the network trained on the motor task in Section 4.4 bearing a “striking resemblance with those recorded in the motor cortex”, we agree that this statement relies on mere visual inspection/comparison of firing rate time series (e.g. see Churchland et al., 2012 for examples of monkey M1 PSTHs) and was not made quantitative. We believe this is perhaps better left to a future study, as doing such comparisons carefully takes significant time and effort, and is not the focus of our paper. The focus here is on establishing the methodology that will enable such studies in future. We are in fact excited about this and have begun to analyse SOFO-trained RNN models of motor cortex; given our preliminary analyses, all we feel comfortable saying at this stage is that our trained RNNs capture at least two salient features of motor cortex dynamics prior to and during reaching: (i) movement-related activity exhibits consistent state-space rotations revealed by jPCA (Churchland et al., 2012) and (ii) is largely orthogonal to preparatory activity (Elsayed et al., 2016).\n\nRegarding the type of biological inquiries that are enabled by the training and subsequent dissection of RNNs, we would like to refer you to the great review by Omri Barak on “Recurrent neural networks as versatile tools for neuroscience research” (Cur. Op. Neurobiol., 2017). This type of approaches have already been influential (e.g. Mante, Sussillo et al, Nature 2013; Sussillo et al., Nat. Neurosci. 2015) and have helped understand of how brains “compute through dynamics”. Since these pioneering contributions, the use of trained RNNs has become increasingly common in neuroscience, see e.g. Driscoll et al., Nat Neurosci 2024 for the more recent developments. We will add these references to our paper. We believe that by significantly accelerating the training of RNNs, SOFO can enable faster progress in this space. In particular, there are a number of computations which brains perform, such as (active) learning or contextual adaptation, that rely on integrating inputs over long behavioural time scales. To the best of our knowledge, few attempts have been made to train RNNs to chart the space of network solutions to these long-horizon problems, and progress here will rely on memory-efficient training algorithms such as SOFO which are capable of handling long horizons. \n\n> their proposed learning algorithm also lacks biological plausibility\n\nJust to clarify: our algorithm is not meant as a model of biological learning; it is meant as an optimization tool to find plausible network solutions to various problems which brains solve.\n\n> Do you expect the proposed approach to also benefit larger networks with coarser discretization of time?\n\nFor larger networks, see our overall rebuttal. Regarding time discretisation: our one-shot classification task of Figure 3 is an example where we did not have an explicit \"dt\" term in the dynamics (equivalently, the dynamics did not incorporate an explicit leak term), i.e. it is an example of coarse discretization. More generally though, we believe that how coarsely one can afford to discretize time does not depend on the optimizer being used, but on the task itself: if, say, a task has trials that last for a second of biological time, and where the millisecond-scale details of the network output influence the task loss, then at least 1000 time bins must be used; in other words, we would expect all optimizers to be equally affected by a poor choice of time discretization that is not suitable for the task at hand.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1bWPq8mJY2", "reviewer": "Reviewer_Z9Qd", "summary": "The manuscript introduces a second order optimization algorithm (called SOFO) for training recurrent neural networks on neuroscience tasks. SOFO does not do backpropagation, instead uses a batched forward-mode differentiation that is memory efficient for large (in time) computation graphs. Authors show that RNNs trained with SOFO outperforms those trained with Adam, a common first-order optimizer for training RNNs, on various neuroscience related tasks. In these tests, SOFO is more memory efficient and converges faster.", "questions": "Overall, I believe the work is barely below the acceptance criteria. Yet, my score can *significantly* increase if the authors address the relevant baselines with appropriate experiments (vs FORCE [1], say on a 3-bit flip flop task as in [2] and varying numbers of neurons), and perform additional experiments to fully flesh out the extent to which SOFO is superior (scaling experiments). Please see my major questions/concerns below:\n\n- I am sure this comment is no surprise to the authors, but the common issue with 2nd order optimizers is the inability to scale to large number of neurons. This, on the other hand, is very important for a computational neuroscience student trying to train their task. Could you please add experiments that show the scaling of SOFO w.r.t. Adam as the number of neurons increase? For example, can you show what happens once we get to the levels of ~10,000 neurons? Can SOFO train well without any memory issues?\n\n- Relatedly, in this work, the authors are mainly using networks with only O(10) neurons. This information can only be found in the appendix. It would be helpful to have this information in relevant figure captions.\n\n- There are several claims about the runtime complexity being comparable to a first order method, but I do not see any formal proof that includes the number of neurons as well? Could you provide the runtime complexity in terms of N and T? I see that there is an SVD computation, so would it be fair to say the complexity scales as O(K^3)? Then, how does K scale with the number of neurons N? It seems K scales linearly with N (from the appendix). Then, SOFO would not have the same complexity as gradient descent. Could the authors clarify?\n\n- I believe the authors are missing some key citations and baselines. I would imagine FORCE [1] would be a directly relevant approach for a fast training of RNNs on neuroscience tasks, which is also a 2nd order approach and can handle relatively large number of neurons (1000 neurons in [2]). I believe the benchmarks should include comparisons with respect to FORCE whenever applicable.\n\n- There are also some other relevant citations that may deserve some mentions, e.g. full-FORCE [3]. There are some more recent works, which the authors may want to include in their relevant work section. In general, computational neuroscience literature is no stranger to 2nd order optimizers, which I think can be emphasized a bit better in the relevant work section to place this work into the broader literature.\n\n- The problem formulation seems to suggest that the loss function is defined over all time steps. What if I had a loss function that penalized only the output after few time steps and had no penalty during the delay period? Would SOFO also work in this case with equal efficiency? Could you show an example task if yes, or state this if no? For example, you can consider a delayed addition task [4], where the output only at the end of the trial matters.\n\nMinor points:\n\n- In Eq. (1), $l_{mi}(y_{mi}(.))$ seems to have only one argument. But later, it is assumed to be strongly convex in its first argument? I guess the authors refer to $\\theta$, whereas $x(\\theta)$ depends on $\\theta$. This is also why the full function is not convex w.r.t. $\\theta$? can you explicitly state this? Overall, the notation needs to be cleaned here.\n\n- \"Whilst using large batches may enable faster learning with larger learning rates, it is often the case that increasing the batch size beyond a certain point no longer helps in this respect\" This sentence requires a citation.\n\n- Lines 124-125: SOFO is not yet introduced, so these lines are confusing for the reader.\n\n**Edit:** The revisions changed my mind and I believe this work is a significant contribution to the neuroscience literature. The results presented were beyond what I thought they would be. I support the acceptance of this work.\n\nCitations\n\n[1] Sussillo, D., & Abbott, L. F. (2009). Generating coherent patterns of activity from chaotic neural networks. Neuron, 63(4), 544-557.\n\n[2] Sussillo, D., & Barak, O. (2013). Opening the black box: low-dimensional dynamics in high-dimensional recurrent neural networks. Neural computation, 25(3), 626-649.\n\n[3] DePasquale, B., Cueva, C. J., Rajan, K., Escola, G. S., & Abbott, L. F. (2018). full-FORCE: A target-based method for training recurrent networks. PloS one, 13(2), e0191527.\n\n[4] Schmidt, D., Koppe, G., Monfared, Z., Beutelspacher, M., & Durstewitz, D. (2019). Identifying nonlinear dynamical systems with multiple time scales and long-range dependencies. arXiv preprint arXiv:1910.03471.", "limitations": "I believe there needs to be a more detailed discussion of SOFO's scaling with respect to number of neurons. I think even if SOFO does not scale well, the paper could easily meet the bar for publication as long as it is stated in the main text with a proper experiment.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- Being able to train RNNs on neuroscience tasks in a scalable and timely manner is an urgent need in the field. \n- The fact that authors achieved these results with a relatively cheap GPU is impressive, and a welcome addition for being able to train RNNs on cheap hardware.", "weaknesses": "- I believe an important baseline, FORCE [1], is missing. The authors should show whether SOFO is preferable to FORCE when training RNNs on neuroscience tasks.\n- I think there needs to be more concrete theoretical and empirical discussions of scaling with respect to the number of neurons.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for your thoughtful review; we have been greatly encouraged this last week by your expressed willingness to potentially increase your score. In the main rebuttal above, we have addressed the concerns common to all reviews, including the effect of network size. Here, we address your other concerns.\n\nFirst, we appreciate your suggestion of using FORCE as another baseline. We have done that in the context of a 3-bit flip-flop task (as you suggested), with results shown in Figure 1 of the one-page PDF. We expected FORCE to be a hard baseline to beat, as RLS is in some sense the optimal thing to do in the “reservoir” setting. However, whilst FORCE does make more rapid initial progress, SOFO eventually wins. To make the comparison fair to FORCE, we implemented a batched version of FORCE-RLS (which we couldn't find anywhere, from a quick scan of the FORCE literature), using the rank-M version of the Woodbury identity for online inverse cov. estimation. This greatly accelerates FORCE learning in our experiments. Note, however, that we have only had the time to do crude tuning of FORCE's damping parameter $\\alpha$ (we tried 1, 10 and 100, and 10 was the best).\n\nInterestingly, we find that SOFO actually works best when the reservoir is non-chaotic at initialization ($g=0.5$ instead of $1.5$)  -- whereas FORCE tends to fail in this setting (as is well known). What exactly that implies for neuroscience applications is unclear. Training from an initially chaotic network might result in residual chaos in corners of state space that are not relevant to the task -- and whether that is a good bias to have presumably depends on whether the brain is chaotic outside task manifolds, something on which we believe the jury is still out. At any rate, by allowing efficient training from non-chaotic initializations, SOFO stand as nicely complementary to FORCE in the RNN training toolbox.\n\nMoreover, SOFO is more flexible than FORCE: it is a general purpose 2nd-order optimizer that can be used to tweak any of the network's parameters, not only the output weights. The Full-FORCE paper nicely demonstrated that by restricting the locus of learning to those few task-relevant output weights, FORCE comes up with solutions that can be brittle. Full-FORCE solved this problem by considering more output units than strictly required by the task and constructing useful pseudo-targets for these additional outputs. This yields higher-rank weight modifications and more robust dynamics. Going further, SOFO can be used to train the _entire_ set of recurrent weights, biases, feedback and input weights. We tried this: SOFO was able to train a much smaller network of 128 units and ReLu activations (instead of the usual FORCE tanh) on the 3-bit flip flop task, to a final test MSE more than 15 times smaller than FORCE's with 1000 neurons. We find these comparisons to FORCE enlightening and will make sure to include them in the main text, not just in the appendix -- thanks again for the suggestion.\n\nRegarding runtime complexity: SOFO applied to an RNN of size $S$ over a time horizon $T$ with batch size $M$ has runtime complexity $\\mathcal{O}(KTMS^2)$, plus $\\mathcal{O}(K^3)$ for the SVD of the sketched GGN as you note (the latter is typically negligible relative to the former, with K being just a few hundred). In contrast, standard backpropagation through time has complexity $\\mathcal{O}(TMS^2)$. Therefore, SOFO incurs an extra factor of $K$, and since K is a hyperparameter that the user is free (a priori) to set independently of $S$, this still makes SOFO technically $\\mathcal{O}(S^2)$. Perhaps the question is (as you rightly point out) whether maintaining a set level of performance implicitly ties $K$ to $S$. In theory, $K$ ought to be on the same order as the effective rank of the GGN matrix (which is hard to estimate a priori). From experimenting with SOFO, we find that performance mainly depends on the ratio of $K$ to the number of parameters to be optimized, and that will be application dependent (e.g. we may only optimize low-rank modifications of the weights, or all recurrent weights, etc). Moreover, all our experiments show that SOFO works well even when this ratio is as low as ~1%. It is also worth noting that SOFO is readily parallelizable over the $K$ tangents: we find that whether we set $K=1$, $10$ or $100$, the wallclock time per iteration does not suffer much, certainly not linearly in $K$. What fundamentally limits $K$ for us is GPU memory. We will make sure to discuss these issues more thoroughly in the paper.\n\nTo provide more clarity on the role of $K$, we have systematically varied $K$ in the delayed-addition task that was suggested to us by Rev. 3Uai (Figure 2 of the one-page PDF). As expected, SOFO gets better with increasing $K$. \n\nYou also raised concerns about SOFO's ability to train much larger networks (e.g. size 1000 as for FORCE). We trained networks of size up to $S=8000$ in the context of the 3-bit flip-flop task, but could have gone higher given that $S=8000$ only used half of our GPU memory. However, $S=8000$ worked well here because, just like FORCE, we only optimized the readout weights on this task (not the full recurrent weight matrix). Optimizing the full recurrent matrix would have the same runtime complexity, but a larger memory overhead given that we would need to store $K$ (~100) parameter tangents of size $S^2$. This is not an unsurmountable obstacle, as these tangents are random and can in principle be evaluated lazily (e.g. recomputed on demand based on a random seed). Sadly this would require more software engineering than we can pull off in this short rebuttal period.\n\nFinally, SOFO is indeed readily applicable to settings where a loss function is not defined for all time steps (would FORCE work in such a setting?) -- this is discussed on L.91, and the one-shot classification example of Figure 3 is an example of this.\n\nWe will address all your minor comments in the revision.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "suDW6sDpSb", "reviewer": "Reviewer_3Uai", "summary": "The present work addresses the difficulty to optimize RNNs that are used in the realm of neurosciences to model biological data, which by design may not benefit from modern RNN inductive biases (e.g. gates), nor be trainable by the Adam optimizer and may inherently hit memory limits due to handling long temporal sequences. To this end, the paper proposes a forward-only second-order optimizer which computes and invert a Gaussian-Newton curvature matrix in a random parameter subspace spanned by $K$ randomly sampled directions in the parameter space. The resulting algorithm amounts, up to a reparametrization of the loss, to exact Gauss-Newton optimization within this parameter subspace. Algorithmically, the procedure amounts to simultaneously computing activations and directional derivatives, a.k.a Jacobian-vector products (JVPs), throughout the forward computational path, for each of the K random directions. As such, the algorithm doesn't require storing any activations and can be parallelized across the K random directions on different GPUs. The algorithm is largely superior to Adam on a variety of neuroscience-oriented tasks in terms of training speed, resulting performance and memory footprint.\n\nMore precisely:\n- Section 2 formulates the optimization problem, introduces standard generalized Gauss-Newton optimization and forward-mode automatic differentiation and briefly surveys subspace optimization methods.\n- Section 3 introduces SOFO (Second Order Forward Only optimizer) reading as follows: at each training iteration, sample K random directions $u_k$ in the parameter space, and compute directional derivatives for the model's Jacobian $Y_k := \\partial_\\theta y \\cdot u_k$ and cost function $C_k := \\partial_\\theta c^\\top \\cdot u_k$ for each of these. Build the GGN matrix within the random subspace with entries $G_k := Y_k^\\top \\cdot \\partial_y^2 L \\cdot Y_k $, compute GN update $\\Delta\\tilde{\\theta}$ in the random subspace coordinates, and project it back to the original parameter space $\\Delta \\theta := \\sum_{k=1}^K \\Delta\\tilde{\\theta}_k \\theta_k$.\n- Section 4 presents results on four different tasks:\n  i) an RNN learns the trajectory of a 3-dimensional Lorentz attractor with sparse error signal (Section 4.1,). SOFO significantly outperforms Adam, its first-order counterpart (i.e. without gradient preconditioning) using $K$ directions, and a single direction (\"FGD\").\n  ii) Adaptive Kalman filtering where observations are produced from ground-truth latents (Linear Dynamical Systems), and an RNN is made to learn to recover these latents from observations. The ground-truth parameters of the LDS may change sporadically, yielding successive such LDS's, making the task more difficult for the RNN to learn. Still here, SOFO outperforms Adam and two baselines where the ground-truth LDS's parameters are known.\n iii)  SOFO is also evaluated on one-shot classification where an RNN is presented with samples from a support set with one example per class (each of these classes being a 3D Gaussian distribution) and has to classify three other samples from a query set. Here again, SOFO is superior to Adam.\niv) Finally, SOFO is evaluated on motor tasks where a target has to be memorized (\"idle\" and \"preparation\" phase) and reached (\"movement\" phase), either a single one (\"single reach task\") or two of these and in a sequential fashion (\"double-reach\" task). The RNN learns to output the appropriate torques to move the arm towards the targets. It is shown that while Adam struggles on the single reach task and totally fails on the double-reach task, SOTA successfully learns these two tasks in terms of resulting performance and training speed, although SOFO runtime exceeds that of Adam. Arm trajectories are depicted to illustrate and compare SOFO and Adam easily on the two tasks, and it is shown that neurons specialize in one of the three phases.\n- Finally, it is shown that while the memory footprint of Adam scales linearly with the sequence length of the task at hand, the memory cost of SOFO is (as expected) constant.\n- It is also worth mentioning that SOFO is tested on larger models on MNIST where it performs on par with Adam (Appendix, Fig. 7).", "questions": "- For memory profiling: I think that a forward-only equivalent of BPTT, namely Recurrent Real Time Learning (RTRL), should instead be considered here in combination with Adam. Just to clarify: using RTRL in place of BPTT will not yield any difference in performance in all the tasks considered as these are equivalent, so I am not questioning the relevance of the results. However, using RTRL instead of BPTT has a different memory cost. Generally, RTRL is discarded because of the cubic memory cost of storing $\\partial_\\theta h$ across time steps. However, because the dimension at use in all your experiments are pretty small, I don't think the memory cost of RTRL would be that huge. **Could you please carry out the memory profiling using RTRL**? \n- Could you please provide architectural details of the MNIST experiments?\n- Did you try evaluating SOFO on CIFAR-10? If you could evaluate SOFO on a 6/7 layers-deep architecture on CIFAR-10, that would be a great cherry on the cake.\n- L.1 \"a common source of anxiety\", L.116 \"embarrassingly parallelizable\", L.165 \"embarrasingly parallel problem\": why be so embarrassed and anxious? I don't think that this language adds much value to the already great presented work.\n- **Could you consider one of these more complex temporal tasks to check your method against**: sCIFAR, IMDB, ListOps? Taking these from \"Online learning of long range dependencies\" [Zucchet et al, 2023] where RTRL is used as an online learning method which is also constant in memory.", "limitations": "See my questions above. I would be very happy to have extra experiments by the end of the discussion period.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- All the claims made about SOFO hold: it significantly outperforms Adam on all tasks, it has a constant memory footprint because forward-only, it amounts to exact Gauss-Newton optimization in the random parameter subspace.\n- The paper is extremely neat and clear: the structure and language of the paper, the mathematical notations, the SOFO pseudo-algorithm, the presentation of the results.\n- The figure quality is high.\n- 4 different experimental setups are investigated. Beyond showing the performance of SOFO, I think it was not trivial to find a set of tasks where Adam fails and which is also of relevance to neuroscience.", "weaknesses": "- The models at use are very small, and at most two-layer deep in space (i.e. each hidden states at a given timestep can be inferred with at most two matrix multiplications).\n- I think that the proposed method should be properly benchmarked against RTRL (which is also an **online**/**forward-only** technique, constant in memory) and more complex temporal tasks (see questions below).\n- I'm not sure that the memory profiling analysis is the most relevant here. See my question below. \n- Architectural details about the MNIST experiment in the Appendix are not given.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your thoughtful review (and for the extensive summary of our paper!). In the main rebuttal above, we have addressed the concerns common to all reviews, including your comments on small network sizes. Here, we wish to address your more specific concerns.\n\nThank you for bringing up RTRL, which also has $\\mathcal{O}(T^0)$ memory cost. Having thought about it, RTRL is mathematically equivalent to a particular limit of SOFO: the limit where (i) SOFO doesn't use curvature information (i.e. dropping the inverse matrix term in Eq. 8) and (ii) the set of tangent vectors used at each iteration ($\\Theta$ in Eqs. 6 and 7) is a full basis equal to the identity matrix. In this limit, batched forward-mode differentiation (efficiently) implements the usual RTRL recursion to propagate exact, entire Jacobians of network activity w.r.t. the parameters, in (constant memory) forward mode. We can therefore predict that, for a model with $P$ parameters, the memory cost of RTRL will be $P/K$ times that of SOFO if SOFO uses $K$ random tangents. Given that the ratio $K/P$ used in our experiments is around $0.01$, this would make RTRL ~100 times more memory-hungry than SOFO. This means that unless we compensate by drastically reducing the (data) batch size, we wouldn't be able to run RTRL on our tasks with our architectures. To confirm these predictions, we implemented RTRL in our codebase on the Kalman example of Figure 2, and will be adding these new data points to the memory profiling of Figure 6. To summarize our findings: RTRL forces us to reduce the batch size to less than 10, and even then it nearly fills our GPU memory (12GB). In contrast, SOFO can easily accommodate a batch size of 256, at which point it only fills 1GB of memory. Thanks very much for reminding us of RTRL, it is in fact rather nice that we are able to frame RTRL as a particular limit of SOFO; we will add a paragraph of discussion on this.\n\nRegarding running more complex tasks: we haven't had a chance to run any of the tasks you suggested from Zucchet et al., 2023. Instead, we followed Rev. Z9Qd's suggestion of adding a delayed addition task, with results in Figure 2 of the one-page rebuttal PDF. In this task, SOFO again outperforms Adam.   \n\nWe also worked to provide the “great cherry on the cake”, by running SOFO on a 5-layer MLP-Mixer (Tolstikhin et al., 2021). This network consists of a patchification layer (patch size $(4,4)$) followed by a per-patch linear embedding to a higher channel dimension of $128$. This is then followed by two mixer blocks, each block including one token-mixing MLP with dimension $64$, and one channel-mixing MLP with dimension $128$ (each MLP uses a GELU nonlinearity). Finally, a fully-connected output layer performs classification. Other components such as skip-connections, dropout and layer norm on the channels are also included. This adds up to about $P \\approx 10^5$ parameters. We compared SOFO with Adam under increasing $K/P$ ratio (with $K$ the number of tangents), from $0.27$% to $1.09$%, with results shown in Figure 3 of the one-page rebuttal PDF. As expected, SOFO gets better as this ratio increases; for $K/P = 1.09$%, SOFO approaches the performance of Adam -- so it is a cherry on a cake, just not a very sweet one (?). We sort of expected this, as this network is a fairly easy optimization target for Adam; we suspect that if we did away with the skip connections and layer-norm tricks, Adam (but not SOFO) would struggle more. We are currently exploring this and will be adding a comment later on, as results become available. Note that we have tried our best to do hyperparameter tuning for all algorithms given the short turnaround time of this rebuttal.\n\n> Could you please provide architectural details of the MNIST experiments?\n\nThese can be found in L.838 in the appendix (2-layer fully-connected MLP, hidden size 100).\n\n> why be so embarrassed and anxious?\n\n(“embarassingly parallel” is a standard term in the field of parallel computing, but we are happy to replace it with the equivalent term \"perfectly parallel\" -- simply meaning that a parallel implementation does not require any clever strategy, the problem being inherently parallel already with no sharing of computation between instances (between the different JVPs, in our case). In any case, we agree that we could well do away with the \"anxiety\" vocabulary :) -- we actually had afterthoughts after submission, thinking that this phrasing was not particular thoughtful to grad students who actually struggle with serious life anxiety, as many do.)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for the time and effort they put into providing thoughtful reviews. We have addressed most of their concerns in reviewer-specific rebuttals, and here is an executive summary:

- We have added a comparison to FORCE learning, on a new task (3-bit flip flop) -- see Figure 1 of the one-page PDF; SOFO performs better than FORCE in this task.  We encourage all reviewers to take a quick look at our rebuttal to Rev. Z9Qd (who suggested this addition), where we discuss in what ways SOFO is not just quantitatively better than FORCE in this case, but also very complementary: it allows training networks initialized in a non-chaotic state (where FORCE relies on chaos), and is capable of training much smaller networks on this challenging task, using more realistic nonlinearities such as ReLU (where FORCE is largely restricted to tanh or other symmetric activation functions due to its inability to optimize single-neuron biases).
- In the context of this 3-bit flip-flop task, we have trained networks of size up to 8000, i.e. two orders of magnitude larger than the networks we trained in the paper.
- We have added results for a delayed addition task suggested by Rev. Z9Qd (Figure 2 of the one-page PDF). There, SOFO again outperforms Adam. In this task, we have also systematically varied the number of tangents $K$ (size of the GGN / gradient sketch) to provide more clarity as to the role of this hyper-parameter. With larger $K$, both SOFO's estimates of the gradient and GGN matrix get more accurate, leading to better performance (see also Appendix A.6). Empirically, we show that for RNN applications, SOFO works well when the ratio of $K$ to the number of parameters $P$ is as low as $1$%, which is one of the reasons why SOFO is so pratical in this setting. When using larger networks (as we did in the flip-flop task), GPU memory forces us to use a smaller $K/P$ ratio, and SOFO does suffer from this: training the network of size $S=8000$ requires more iterations than for $S=1000$; we will more clearly lay out this limitation of SOFO in our revision.
- We have also added a comparison between Adam and SOFO on an MLP-Mixer architecture (~$10^5$ parameters) trained on CIFAR10 (as suggested by Rev. 3Uai) -- see Figure 3 of the one-page PDF. This is a fairly easy target for Adam due to the presence of skip connections and layer-norm, and with the limits on $K$ imposed by GPU memory, the best SOFO does here is to nearly match Adam's performance. Note that this shallow network setup is not where we expect SOFO to be the most useful; its primary strength is in fixing the $\mathcal{O}(T)$ memory complexity of backpropagation through time in RNNs, making it $\mathcal{O}(T^0)$, enabling training in settings where backprop with either large networks or large data batches is not feasible for a given long time horizon.  
- Prompted by Rev. 3Uai, we have studied the relationship between SOFO and Real-Time Recurrent Learning (RTRL). RTRL turns out to be mathematically equivalent to a particular limit of SOFO: that of (i) not using curvature information and (ii) using a full, identity basis of parameter tangents (as opposed to a small number of random tangents). This is a lot more memory intensive, as we demonstrate by running memory profiling for RTRL; we find that we can only afford to run RTRL provided that we also dramatically reduce the data minibatch size.

We hope that the reviewers will find, as we do, that these additions strengthen our paper.

Meta-review: This paper suggests using second-order optimization, essentially a type of quasi-Newton method, for “vanilla” RNNs as used in neuroscience. While the method itself is not that new, the authors introduce an efficient GPU-based parallelization that makes runtimes comparable to those achieved with standard gradient-based techniques, but with largely superior performance. This is indeed an important methodological contribution which shifts the focus from improving architectures to improving training algorithms, following a recent trend in ML. While initially there were some concerns about novelty (see above), biological plausibility, or size of the models used (well addressed in the rebuttal), overall this manuscript received a fairly enthusiastic response from most referees, with all four referees ultimately on the accept side, which I support.
Decision: Accept (spotlight)