用零阶优化训练百亿参数长序列RNN,内存占用降为千分之一。
Scaling Recurrent Neural Networks to a Billion Parameters with Zero-Order Optimization
- 用随机向量梯度估计替代反向传播,训练时无需存储中间激活值。
- 在语言建模等任务上,收敛速度比传统方法快19倍,且泛化能力更强。
- 适合训练超大规模长序列模型,尤其对显存受限场景友好。
推理时,循环神经网络(RNN)的计算量和显存占用随上下文长度增长保持恒定,因其将所有历史标记压缩为固定大小的隐状态。而变压器模型在生成时计算量和显存均线性增长,必须显式关注所有先前标记。尽管推理优势显著,但传统优化方法依赖反向传播通过时间(BPTT),需在前向传播中保存所有中间激活值,导致内存消耗随上下文长度和模型规模线性增长,难以训练大规模长上下文RNN。本文证明,零阶优化(ZOO)方法如随机向量梯度估计(RGE)可成功替代BPTT,在训练过程中始终以推理模式运行,实现与或优于BPTT的收敛速度(最高快19倍),同时内存和成本降低数个数量级。我们进一步表明,中心差分RGE(CD-RGE)对应于优化一个平滑的代理损失,具有内在正则化效果,提升泛化能力。在三种任务设置下:过拟合、转换和语言建模,该方法性能均匹配或超越BPTT,使用足够扰动时,模型泛化能力相当甚至更优,且常在更少步数内完成。尽管每步需更多前向传递,但借助FlashRNN和分布式推理等技术,仍可在实际时间上超越BPTT。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。