【新智元导读】RNN模型在长上下文中表现不佳？近日，来自清华的研究团队对此进行了深入的实验分析，结果表明：不是RNN的锅。
与Transformer相比，RNN模型的一大优势是应对长序列的能力。
比如Mamba，内部状态大小始终保持不变，计算随序列长度线性增长，吃得多，消化快。
理论虽如此，但实际情况却是，目前的这些RNN模型在长上下文中的有效性并不能令人满意。
为啥会这样？空有效率但实际上能力不行？
近日，来自清华的研究团队对此进行了深入的实验研究：
一是无法推断比训练长度更长的输入，原因是较短的训练数据导致了循环状态过拟合；
二是内存容量的上限，由于模型无法有效遗忘很久以前的信息，导致新的信息存不进来了。
——这俩问题明显不是RNN的锅。
而经过研究人员的对症下药，Mamba-2（370M）在256K上下文长度上达到了近乎完美的密钥检索精度。
所以结论就是，Mambayes！「RNN神教」前景一片光明！
对此，Mamba的作者AlbertGu点赞转发，并发表了相当详细的见解：
「这是一篇很棒的论文（名字也很棒）——关于状态空间模型（SSM）的状态容量和长上下文能力的巧妙实验。」
令人惊讶的是，对于每个状态大小M，当训练上下文长度达到或超过某个临界值K时，都会出现一个转折点，在这个点上SSM就能够稳健地实现长度泛化。这是因为当上下文长度小于K时，循环状态没有被充分利用，导致模型在训练期间会「过拟合」。但一旦通过足够长序列的训练使模型的状态容量得到充分利用，它就会自动获得泛化能力。值得注意的是，K与M竟然呈线性关系！——这表明每个token可能存在某种固有的信息含量（即存在一个值B，使得上下文中的每个token对应B字节的循环状态）。这个B值可能是由模型架构决定的？
最后，AlbertGu用一句话总结：要让你的Mamba吃得饱饱的，它就能发挥出最佳状态！
喂饱你的Mamba
先来复习一下基础知识。
本文以Mamba2作为主要研究对象，内部的计算表示为下图中的并行结构：
整体的输入输出遵循SSM（也即RNN）的形式：
而把上图中模块内部所有的计算写出来，就是下面这一坨公式：
之前提到的两个问题，核心在于模型的内部状态，也就是ht的表现。
之前有研究表明，当上下文长度超过其训练长度时，Mamba-1和RWKV-4的性能会严重下降。
顺着这个思路，研究人员在两个方向上进行了实验分析：状态崩溃（STATECOLLAPSE）和容量上限（STATECAPACITY）。
状态崩溃
上图展示了Mamba-2和RWKV-6在训练长度之外的语言建模损失。为了可控性和合成任意长度的提示，这个损失是在仅由「\n」字符组成的提示上计算的（称为「newlines」提示）。
结果表明，当上下文长度远大于其训练长度时，两个RNN的性能都会严重下降，最后就跟瞎猜差不多了。
语言建模可能无法反映下游能力，上图给出了Mamba-2（在8K上下文窗口上训练）在密钥检索任务上的评估结果。
我们可以发现，Mamba-2在8K上下文中具有近乎完美的检索准确性，但在序列长度超过16K后就没法看了，无论模型参数量大小。
从上面的公式来看，这种结果可能出人意料，因为内部状态ht的更新应该具有稳定的指数内存衰减，即对于最后k个token具有良好的检索准确性。
问题出在哪里？
作者对Mamba-2370M中每一层的递归状态进行了统计，发现当上下文长度超过训练长度时，一些头部的平均值和方差会急剧变化：
图5显示了模型第38层第2个头的状态，在t=20K时方差爆炸。从中可以发现这种方差爆炸在很大程度上可以归因于少数异常通道，其余大多数通道则相对稳定。
分析一下公式，与ht计算有关的t、Bt和xt：
如上图所示，虽然三者都是输入的函数，但xt相对稳定，而Bt比t更早发生爆炸，进一步探索还能发现生成t和Bt的卷积权重明显更大。
作者认为，产生SC的原因是，对于训练长度来说，状态容量过大，模型能够实现强大的语言建模性能，而无需学习如何忘记。
解决方案
为了缓解SC，使模型沿序列长度更好地泛化，作者提出了3种解决方案，总的思想是修改状态的update规则来避免其溢出。
Method1:ForgetMoreandRememberLess
通过增加状态衰减量（忘记更多）或减少输入信息的数量（记住更少）来减少SC，作者选择干预Bt和αt（分别控制输入强度和内存衰减强度）。
Method2:StateNormalization
在每次更新后对状态进行归一化，以确保状态的范数始终低于阈值：
PS：这种方式会将模型转换为非线性RNN，无法以与原始模型相同的方式并行化，预填充速度要慢得多。
Method3:SlidingWindowbyStateDifference
利用状态ht可以写为加权和的形式，来模拟滑动窗口机制，无需在每一步都从窗口的开头重新处理。
此方法适用于所有可以写成加权和的RNN，包括RWKV5和6、RetNet、GLA等。尽管会使生成的计算和内存成本翻倍，但仍然是一个可以接受的权衡，因为RNN的生成成本比Transformer低很多。
以上3个是不需要训练的方案，而基于SC是由状态参数过拟合引起的假设，我们也可以尝试使用超过状态容量的序列长度来训练模型。
容量上限
根据以上的讨论，当且仅当训练长度包含的信息少于状态容量时，才会发生SC，所以我们可以通过实验间接估计模型的状态容量。
研究人员训练了多个具有不同状态大小和训练长度的Mamba-2，并将SC未发生的最小训练长度视为状态容量。
实验数据选择RedPajama-V2，一个从CommonCrawl中提取的30Ttoken的开放数据集，进行去重以确保数据质量。
在评估过程中，对长度超过16Ktoken的文档进行抽样，如果不够长，则对其进行拼接。
研究人员试验了具有不同状态大小的模型配置，包括来自Mamba-2官方checkpoint的三个预训练模型，大小分别为130M、370M和780M，另外3个模型（36M、47M、85M）则从头开始训练。
实验结果
上图展示了在Mamba-2780M上无训练长度泛化方法的结果。我们可以看到，虽然LongMamba大大提高了模型的长度泛化性（3倍以上），但它在较短的序列上会导致明显更大的困惑度，并且仍然不可避免地表现出SC。
相比之下，本文的所有的方法都成功地抑制了SC，使模型能够泛化到超过64K个token。
三种方案中，状态归一化在较短序列上的性能大大低于其他方法，这可能是因为归一化折叠状态会改变heads之间的规范比率，破坏了学习机制。
上图显示了Mamba-2在语言建模和密钥检索方面的状态容量。两个图中最右边的数据点对应于Mamba-2370M。
左边的图可以拟合出一个线性关系，而右边的图则表明Mamba-2在密钥检索方面的容量与状态大小呈指数级关系。
这是因为上下文中的信息量不会随着其长度的增加而增加。换句话说，模型存储了恒定数量的信息，而状态的组合数量随着元素数量呈指数增长。