Second-order forward-mode optimization of recurrent neural networks for neuroscience

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.

Thank 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.

Regarding 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.

We 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.

Could you please provide architectural details of the MNIST experiments?

These can be found in L.838 in the appendix (2-layer fully-connected MLP, hidden size 100).

why be so embarrassed and anxious?

(“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.)

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.

First, 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).

Interestingly, 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.

Moreover, 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.

Regarding 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.

To 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$.

You 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.

Finally, 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.

We will address all your minor comments in the revision.

Thank you very much for constructively engaging with our work!

We will of course add the details of batched FORCE in an Appendix.

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.

Regarding 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.

Our 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.

Regarding 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.

Regarding 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).

Regarding 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.

their proposed learning algorithm also lacks biological plausibility

Just 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.

Do you expect the proposed approach to also benefit larger networks with coarser discretization of time?

For 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.

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.

We 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.

Regarding 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.

not clear why a stochastic RNN was chosen for the motor task

We 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).

In line 90, the authors say "In RNN applications, i indexes both time bins and RNN output dimensions" which doesn't sound quite right.

What 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.