摘要

循环神经网络（RNNs）超越基于transformer的语言模型的关键优势在于它们关于序列长度的线性计算复杂度，使得它们在处理长序列推理中更加快速。然而，大多公开可用的RNNs（如，Mamba和RWKV）在少于10K标记的序列上训练，它们在更长上下文中的有效性目前仍极为不满意。在本文中，我们研究了RNN无法处理长上下文的原因，并提出了关键的缓解措施。我们研究了将最先进的RNN应用于长上下文时的两个实际问题：（1）无法推断长于训练长度的输入；（2）内存容量的上限。为了解决第一个问题，我们首先研究了状态崩溃（SC），这是一种在训练过程中没有遇到的序列长度上导致严重性能下降的现象。通过对照实验，我们将其归因于训练长度对循环状态的过度拟合。对于第二个问题，我们在长文档上训练了一系列Mamba-2模型，以实证估计语言建模和密钥检索中的循环状态容量。然后，提出了三种SC缓解方法来提高Mamba-2的长度泛化能力，使模型在没有SC的情况下处理超过1M个令牌。我们还发现，密钥检索中的循环状态容量与状态大小呈指数关系，我们在256K上下文长度上实证训练了一个具有接近完美密钥检索精度的Mamba-2 370M。这表明基于RNN的长上下文建模前景广阔。论文地址