ePoster

Second-order forward-mode optimization of RNNs for neuroscience

Youjing Yu, Rui Xia, Qingxi Ma, Mate Lengyel, Guillaume Hennequin
COSYNE 2025(2025)
Montreal, Canada

Conference

COSYNE 2025

Montreal, Canada

Resources

Authors & Affiliations

Youjing Yu, Rui Xia, Qingxi Ma, Mate Lengyel, Guillaume Hennequin

Abstract

Training recurrent neural networks (RNNs) to perform neuroscience tasks can be challenging. Unlike in machine learning where any architectural modification of an RNN (e.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 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. Moreover, 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, the learning of an adaptive Kalman filter algorithm trained over a long horizon and many more. In addition, SOFO can train flexible network models without requiring chaotic initializations, which makes it nicely complementary to FORCE learning as another second-order optimization method in the RNN training toolbox. 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.

Unique ID: cosyne-25/second-order-forward-mode-optimization-f0decf9b