LLM2D
填充曼巴:基于循环神经网络的长上下文建模中的状态崩溃和状态容量
Stuffed Mamba: State Collapse and State Capacity of RNN-Based Long-Context Modeling
作者: Yingfa Chen, Xinrong Zhang, Shengding Hu, Xu Han, Zhiyuan Liu, Maosong Sun
发布日期: 10/10/2024
arXiv ID: oai:arXiv.org:2410.07145v1

摘要

循环神经网络(RNN)相较于基于Transformer的语言模型的一个重要优势在于其关于序列长度的线性计算复杂度,这使得它们在推理过程中处理长序列的速度快得多。然而,大多数公开可用的RNN(例如,Mamba和RWKV)是在少于10K个token的序列上训练的,它们在更长上下文中的有效性迄今为止仍然令人不满意。本文研究了RNN无法处理长上下文的原因,并提出了一些关键的缓解方法。我们考察了将最先进的RNN应用于长上下文时的两个实际问题:(1)无法推断到比训练长度更长的输入,以及(2)内存容量的上限。为了解决第一个问题,我们首先研究了*状态崩溃*(SC),这是一种导致序列长度在训练期间未遇到的性能严重下降的现象。通过对照实验,我们将此归因于由于循环状态对于训练长度参数过多而导致的过拟合。对于第二个问题,我们在长文档上训练了一系列Mamba-2模型,以经验估计语言建模和密钥检索中的循环状态容量。然后,提出了三种SC缓解方法来提高Mamba-2的长度泛化能力,允许模型处理超过100万个token而不会出现SC。我们还发现,密钥检索中的循环状态容量随状态大小呈指数级缩放,并且我们通过实验证明了一个具有近乎完美密钥检索精度的3.7亿参数的Mamba-2模型,其上下文长度为256K。这表明基于RNN的长上下文建模具有光明的前景。