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

在推理过程中，循环神经网络（RNNs）随着上下文长度的增加，在浮点运算（FLOPs）和GPU内存使用上保持常数规模，因为它们将所有先前的标记压缩到固定大小的记忆中。相比之下，Transformer在生成时的浮点运算呈线性增长，而内存使用即使在最佳情况下也呈线性增长，因为它们必须显式关注所有之前的标记。尽管RNN在推理时具有这一优势，但在长上下文中训练大型RNN仍然不切实际，因为标准优化方法依赖于通过时间反向传播（BPTT）。BPTT要求在前向传递过程中保留所有中间激活值，这导致内存使用随着上下文长度和模型规模呈线性增长。 在本文中，我们展示了零阶优化（ZOO）方法，例如随机向量梯度估计（RGE），可以成功替代BPTT来训练RNN，其收敛速度与BPTT相当，甚至最高可超过BPTT达19倍，同时使用的内存和成本要低几个数量级，因为在整个训练过程中模型始终保持在推理模式。我们进一步证明，中心差分随机梯度估计（CD-RGE）对应于优化一个平滑的代理损失函数，这本质上对训练进行了正则化，并提高了泛化能力。我们的方法在三种设置下与BPTT表现相当或更优：（1）过拟合、（2）转导、以及（3）语言建模。在所有任务中，只要提供足够的扰动，我们的模型泛化能力与通过BPTT训练的模型相当甚至更好，且通常在更少的步骤内完成。尽管每一步需要更多的前向传递，但通过最近的技术进步如FlashRNN和分布式推理，我们可以在每步的实际运行时间内超越BPTT。

该论文试图解决训练大型RNN在长上下文场景下的内存和计算成本问题。传统方法依赖于BPTT，其内存使用随上下文长度和模型大小线性增长，这使得训练变得不切实际。这是一个长期存在的问题，但本文提出了一种全新的优化方法来应对这一挑战。

论文的关键思路是用零阶优化（ZOO）方法，如随机向量梯度估计（RGE）及其改进版CD-RGE，替代传统的BPTT进行RNN训练。这种方法使模型在整个训练过程中保持推理模式，显著降低了内存需求，并且在某些情况下收敛速度比BPTT快19倍。此外，CD-RGE通过优化平滑的代理损失函数，提供了内在的正则化效果，从而改善了泛化性能。

论文展示了三种任务（过拟合、转导和语言建模）上的实验结果，证明了所提出方法在匹配或超越BPTT的同时，大幅减少了训练成本。实验设计涵盖了不同规模的数据集，并验证了模型在充分扰动条件下的泛化能力。尽管每步需要更多前向传播，但结合FlashRNN和分布式推理技术，可以实现比BPTT更快的每步训练时间。代码尚未明确提及是否开源，但研究为高效训练RNN开辟了新方向，值得进一步探索更复杂任务上的表现。

最近的相关研究包括：(1) 使用截断BPTT减少内存消耗的工作；(2) 针对Transformer架构提出的稀疏注意力机制以降低计算复杂度的研究；(3) 探索基于进化算法或其他启发式方法训练神经网络的尝试。例如，《Optimizing Recurrent Models with Truncated Backpropagation》讨论了BPTT的变体，而《Sparse Transformers》则专注于减少Transformer的计算负担。此外，《Evolutionary Strategies as a Scalable Alternative to Reinforcement Learning》提出了进化策略作为另一种无梯度优化方法。