Scaling Recurrent Neural Networks to a Billion Parameters with Zero-Order Optimization

• URL: http://arxiv.org/abs/2505.17852v1
• Date: Fri, 23 May 2025 13:04:06 GMT
• Title: Scaling Recurrent Neural Networks to a Billion Parameters with Zero-Order Optimization
• Authors: Francois Chaubard, Mykel Kochenderfer,
• Abstract summary: RNNs scale constant in FLOPs and GPU memory with increasing context length.<n>Transformers scale linearly in FLOPs and, at best, linearly in memory during generation.<n>Training large RNNs on long contexts remains impractical because standard optimization methods depend on Backpropagation Through Time (BPTT)
• Score: 0.0
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: During inference, Recurrent Neural Networks (RNNs) scale constant in both FLOPs and GPU memory with increasing context length, as they compress all prior tokens into a fixed-size memory. In contrast, transformers scale linearly in FLOPs and, at best, linearly in memory during generation, since they must attend to all previous tokens explicitly. Despite this inference-time advantage, training large RNNs on long contexts remains impractical because standard optimization methods depend on Backpropagation Through Time (BPTT). BPTT requires retention of all intermediate activations during the forward pass, causing memory usage to scale linearly with both context length and model size. In this paper, we show that Zero-Order Optimization (ZOO) methods such as Random-vector Gradient Estimation (RGE) can successfully replace BPTT to train RNNs with convergence rates that match, or exceed BPTT by up to 19 fold, while using orders of magnitude less memory and cost, as the model remains in inference mode throughout training. We further demonstrate that Central-Difference RGE (CD-RGE) corresponds to optimizing a smoothed surrogate loss, inherently regularizing training and improving generalization. Our method matches or outperforms BPTT across three settings: (1) overfitting, (2) transduction, and (3) language modeling. Across all tasks, with sufficient perturbations, our models generalize as well as or better than those trained with BPTT, often in fewer steps. Despite the need for more forward passes per step, we can surpass BPTT wall-clock time per step using recent advancements such as FlashRNN and distributed inference.

Related papers

• Convergence Analysis of Real-time Recurrent Learning (RTRL) for a class of Recurrent Neural Networks [0.0]
  Real-time recurrent learning (RTRL) is an online optimization algorithm.<n>It follows the true gradient loss on the data sequence as the number of sequence time steps.<n>One potential application area for RTRL is the analysis of financial data.
  arXiv  Detail & Related papers  (2025-01-14T11:46:36Z)
• Optimal Gradient Checkpointing for Sparse and Recurrent Architectures using Off-Chip Memory [0.8321953606016751]
  We introduce memory-efficient gradient checkpointing strategies tailored for the general class of sparse RNNs and Spiking Neural Networks.<n>We find that Double Checkpointing emerges as the most effective method, optimizing the use of local memory resources while minimizing recomputation overhead.
  arXiv  Detail & Related papers  (2024-12-16T14:23:31Z)
• UIO-LLMs: Unbiased Incremental Optimization for Long-Context LLMs [111.12010207132204]
  UIO-LLMs is an incremental optimization approach for memory-enhanced transformers under long-context settings.<n>We refine the training process using the Truncated Backpropagation Through Time (TBPTT) algorithm.<n>UIO-LLMs successfully handle long context, such as extending the context window of Llama2-7b-chat from 4K to 100K tokens with minimal 2% additional parameters.
  arXiv  Detail & Related papers  (2024-06-26T08:44:36Z)
• Winner-Take-All Column Row Sampling for Memory Efficient Adaptation of Language Model [89.8764435351222]
  We propose a new family of unbiased estimators called WTA-CRS, for matrix production with reduced variance. Our work provides both theoretical and experimental evidence that, in the context of tuning transformers, our proposed estimators exhibit lower variance compared to existing ones.
  arXiv  Detail & Related papers  (2023-05-24T15:52:08Z)
• RWKV: Reinventing RNNs for the Transformer Era [54.716108899349614]
  We propose a novel model architecture that combines the efficient parallelizable training of transformers with the efficient inference of RNNs. We scale our models as large as 14 billion parameters, by far the largest dense RNN ever trained, and find RWKV performs on par with similarly sized Transformers.
  arXiv  Detail & Related papers  (2023-05-22T13:57:41Z)
• Towards Memory- and Time-Efficient Backpropagation for Training Spiking Neural Networks [70.75043144299168]
  Spiking Neural Networks (SNNs) are promising energy-efficient models for neuromorphic computing. We propose the Spatial Learning Through Time (SLTT) method that can achieve high performance while greatly improving training efficiency. Our method achieves state-of-the-art accuracy on ImageNet, while the memory cost and training time are reduced by more than 70% and 50%, respectively, compared with BPTT.
  arXiv  Detail & Related papers  (2023-02-28T05:01:01Z)
• Online Training Through Time for Spiking Neural Networks [66.7744060103562]
  Spiking neural networks (SNNs) are promising brain-inspired energy-efficient models. Recent progress in training methods has enabled successful deep SNNs on large-scale tasks with low latency. We propose online training through time (OTTT) for SNNs, which is derived from BPTT to enable forward-in-time learning.
  arXiv  Detail & Related papers  (2022-10-09T07:47:56Z)
• GLEAM: Greedy Learning for Large-Scale Accelerated MRI Reconstruction [50.248694764703714]
  Unrolled neural networks have recently achieved state-of-the-art accelerated MRI reconstruction. These networks unroll iterative optimization algorithms by alternating between physics-based consistency and neural-network based regularization. We propose Greedy LEarning for Accelerated MRI reconstruction, an efficient training strategy for high-dimensional imaging settings.
  arXiv  Detail & Related papers  (2022-07-18T06:01:29Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.