Research map / Architectures and efficiency / RNNs and normalization
Second-order forward-mode optimization of recurrent neural networks for neuroscience
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes SOFO, a second-order optimizer that uses batched forward-mode AD to compute sketched Gauss-Newton updates in random subspaces, enabling constant-memory RNN training on neuroscience tasks.
Problem
Training vanilla RNNs as models of brain dynamics is hard: ill-conditioned loss surfaces, long time horizons exceed backpropagation memory, and GPU parallelism is underused; plausibility constraints rule out gated architectures.
Approach
SOFO reparameterizes the model in a random K-dimensional subspace each iteration and computes an exact sketched Generalized Gauss-Newton update using batched forward-mode JVPs, giving O(1) memory in time horizon and per-iteration cost close to first-order methods.
Key results
- On the Kalman filter task SOFO trains successfully within under 200 iterations while Adam converges to a suboptimal solution.
- SOFO's memory stays constant with time horizon T while Adam's grows linearly, exceeding a 12GB RTX 2080 Ti around T~30K steps.
- SOFO is roughly 1.5-3x slower than Adam per iteration in wallclock time, and on the double-reach task Adam fails to learn while SOFO succeeds.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use SOFO instead of Adam | Vanilla RNN · Lorenz attractor | training loss | Helped |
| Add second-order GGN preconditioning to the subspace update (ablation) | SOFO · Lorenz attractor | training loss | Helped |
| Use K=1 and omit second-order preconditioning (ablation) | SOFO · Lorenz attractor | training loss | Hurt |
| Use SOFO instead of Adam | Vanilla RNN · Adaptive Kalman filter | training loss: less than 200 training iterations | Helped |
| Use SOFO instead of FORCE | RNN · 3-bit flip-flop | evaluation loss | Mixed |
| Train all network weights in the smaller 128-unit network (ablation) | SOFO · 3-bit flip-flop | test mean squared error: more than 15 times smaller | Helped |
| Use SOFO instead of Adam | Stochastic vanilla RNN · Memory-guided single-reach | training loss | Helped |
| Use SOFO instead of Adam | Stochastic vanilla RNN · Memory-guided double-reach | training loss | Helped |
| Use SOFO instead of Adam | RNN · One-shot classification | classification accuracy: 90% ... in 150 training iterations, compared to 1500 iterations for Adam | Helped |
| Increase the tangent count K (ablation) | SOFO · Delayed addition (T=128) | training loss | Helped |
| Increase the number of neurons while keeping K fixed at 256 (ablation) | SOFO · 3-bit flip-flop | convergence speed: 1000 to 8000 | Hurt |
| Train an MLP-Mixer on CIFAR-10 with SOFO instead of Adam | MLP-Mixer · CIFAR-10 | evaluation accuracy | No clear effect |
Surprising findings
- Using a random subspace with K/P around 1% suffices for SOFO to outperform Adam on RNN tasks, despite optimizing only a small fraction of parameters per iteration.
- SOFO performs poorly relative to its advantage on the MLP-Mixer CIFAR-10 task, at best matching Adam, and its convergence worsens as network size grows with fixed K.
Limitations
- Performance deteriorates as the number of parameters grows for fixed K, since the GGN sketch gets noisier with smaller K/P ratio.
- Advantage over Adam is less evident on richer learning tasks such as CIFAR-10 with an MLP-Mixer, where it at best matches Adam.
- Runtime is linear in K times that of reverse-mode AD in principle, and hyperparameters were manually tuned.
Experimental setup
Datasets: Lorenz attractor (synthetic) linear dynamical system sessions (synthetic) 3-bit flip-flop (synthetic) single and double reaching task (simulated arm) delayed addition (synthetic) one-shot classification (synthetic) MNIST CIFAR-10
Benchmarks: Lorenz attractor sparse supervision adaptive Kalman filter 3-bit flip-flop memory-guided single and double reach delayed addition one-shot classification MNIST CIFAR-10
Metrics: training loss evaluation loss wallclock time per iteration GPU memory usage classification accuracy
Baselines: Adam FORCE Forward-gradient descent (FGD) first-order SOFO Kalman filter baseline
Builds on: Generalized Gauss-Newton forward-gradient descent FORCE real-time recurrent learning stochastic subspace descent
Closest papers by meaning
- Emergent mechanisms for long timescales depend on training curriculum and affect performance in memory tasks (ICLR 2024 · similarity 0.87)
- Recurrent neural networks: vanishing and exploding gradients are not the end of the story (NeurIPS 2024 · similarity 0.86)
- Recurrent neural network dynamical systems for biological vision (NeurIPS 2024 · similarity 0.86)
- FlashRNN: I/O-Aware Optimization of Traditional RNNs on modern hardware (ICLR 2025 · similarity 0.86)
- Exact, Tractable Gauss-Newton Optimization in Deep Reversible Architectures Reveal Poor Generalization (NeurIPS 2024 · similarity 0.86)
- Forward Learning with Top-Down Feedback: Empirical and Analytical Characterization (ICLR 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.