使用零阶优化将循环神经网络扩展至十亿参数规模

在推理期间，循环神经网络 (RNN) 的 FLOPs 和 GPU 内存占用随上下文长度的增大而恒定，因为它们将所有先前标记压缩到固定大小的内存中。相比之下，transformer 在生成期间的 FLOPs 最多线性增长，内存占用也最多线性增长，因为它们必须显式地注意所有先前标记。尽管存在这种推理优势，但在长上下文上训练大型 RNN 仍不可行，因为标准优化方法依赖通过时间反向传播 (BPTT)。BPTT 需要在前向传递期间保留所有中间激活，从而导致内存占用随上下文长度和模型大小的线性增长。本文表明，零阶优化 (ZOO) 方法，如随机向量梯度估计 (RGE)，可以成功地取代 BPTT 来训练 RNN，其收敛率可匹配甚至比 BPTT 高出最多 19 倍，而且只需极少的内存和计算成本，因为模型在训练期间始终处于推理模式。我们进一步展示，中心差分 RGE (CD-RGE) 对应于优化平滑代理损失，内在地对训练进行正则化，从而提高泛化能力。我们的方法在三个设置下与 BPTT 相匹配或优于：(1) 过拟合，(2) 序列转换，和 (3) 语言建模。在所有任务中，只要扰动足够大，我们的模型在泛化方面表现出与或优于使用 BPTT 训练的模型，往往所需步数更少。尽管每一步需要更多的前向传递，但我们可以通过近期进展如 FlashRNN 和分布式推理来超越 BPTT 的每步墙钟时间。