不是RNN的锅！清华团队深入分析长上下文建模中的状态崩溃，Mamba作者点赞

在自然语言处理（NLP）领域，长上下文建模一直是个挑战。最近，来自清华大学的团队发表了一篇论文，深入分析了基于RNN（循环神经网络）的长上下文模型中的状态崩溃现象，并提出了有效的缓解方法。这篇论文引起了广泛关注，甚至得到了Mamba模型作者的点赞。

RNN相比于基于Transformer的语言模型，具有线性的计算复杂度，这使得它们在处理长序列时更加高效。然而，大多数公开可用的RNN模型（如Mamba和RWKV）在训练时使用的序列长度小于10K，导致它们在更长的上下文中表现不佳。

论文首先研究了RNN在处理长上下文时所面临的状态崩溃（SC）现象。SC是指模型在遇到超过训练长度的输入时，性能急剧下降。通过控制实验，研究团队发现SC是由于RNN的状态在训练长度上过参数化，导致模型在处理更长序列时无法正确泛化。

为了解决SC问题，研究团队提出了三种训练无关的缓解方法和一种基于连续训练的方法。

1. 减少记忆与增加遗忘：通过增加状态衰减或减少插入信息的量，使模型在处理长序列时能够更好地遗忘旧信息，避免状态溢出。

2. 状态归一化：在每个时间步对状态进行归一化，确保状态的范数始终低于某个阈值，从而避免状态的剧烈变化。

3. 滑动窗口机制：利用状态可以表示为先前插入信息的加权和这一特性，模拟滑动窗口机制，使模型能够在不重新处理整个窗口的情况下，处理长序列。

4. 训练更长的序列：通过在超过状态容量的序列长度上进行训练，鼓励模型学习如何平滑地遗忘最早的信息，从而提高模型在长上下文中的泛化能力。

研究团队在Mamba-2模型上进行了广泛的实验，以验证这些缓解方法的有效性。实验结果表明，这些方法能够显著提高模型在长上下文中的性能，使模型能够处理超过1M的tokens而不会崩溃。

此外，研究团队还分析了状态容量与状态大小之间的关系，并发现状态容量是状态大小的线性函数。他们还发现，在passkey检索任务中，模型的准确性是状态大小的指数函数，这表明RNN在处理长上下文时具有巨大的潜力。

这篇论文的发表，为解决RNN在长上下文建模中的状态崩溃问题提供了新的思路和方法。它不仅得到了Mamba模型作者的认可，也为其他研究者提供了宝贵的参考。

然而，这篇论文也存在一些局限性。首先，它主要研究了Mamba-2模型，而没有对其他RNN模型进行广泛的实验。其次，它的连续训练方法相对昂贵，可能不适合所有应用场景。最后，它使用的passkey检索任务相对简单，可能无法完全反映模型在真实世界中的长上下文处理能力。

尽管存在一些局限性，这篇论文为未来的研究提供了许多可能的方向。例如，可以进一步研究如何将这些缓解方法应用于其他RNN模型，或者探索更高效的连续训练方法。此外，还可以研究如何将这些方法与更复杂的任务相结合，以评估模型在真实世界中的长上下文处理能力。

论文地址：https://arxiv.org/pdf/2410.07145v1