重新思考循环神经网络中的双线性状态转换：隐藏单元不仅是记忆，更是计算的主角

在2025年5月，高通AI研究团队的M.Reza Ebrahimi和Roland Memisevic发布了一篇深入探讨循环神经网络(RNN)结构的研究论文。这篇发表在arXiv预印本平台(arXiv:2505.21749v1)上的论文，重新审视了神经网络中一个被长期忽视的观点：隐藏单元不仅仅是被动的记忆存储器，它们还能作为网络计算过程中的积极参与者。

一、研究背景：隐藏单元的双重身份

当我们谈论循环神经网络时，大多数人会将其隐藏单元视为"记忆存储器"，就像电脑中的RAM一样，主要功能是记住过去看到的信息。这种观点导致了大量研究集中在改进这些"记忆单元"的信息保留能力上，比如通过各种"门控机制"来控制信息的流入流出。

然而，高通研究团队提出了一个不同的视角：隐藏单元不应该仅仅被视为被动的记忆仓库，而应该被看作是网络进行实际计算的积极参与者。想象一下，这就像是将计算机中的RAM转变为一个小型处理器，不仅存储数据，还能对数据进行处理和变换。

这种观点并不完全是新发明，而是对一种被称为"双线性操作"（bi-linear operations）的旧概念的重新审视。所谓双线性操作，指的是隐藏单元与输入嵌入之间的乘法交互作用。简单来说，就是让隐藏状态的更新不仅受到当前输入的影响，还受到前一时刻隐藏状态的影响，而且这种影响是相乘的关系，而非简单的相加。

二、研究的核心：状态跟踪任务与双线性转换

为什么研究团队如此关注隐藏单元的这种双线性更新机制呢？答案在于一类特殊的问题：状态跟踪任务。

想象你正在玩一个电子游戏，游戏中的角色会根据你的输入（比如按下向左、向右、跳跃等按钮）改变自己的状态。游戏角色的下一个状态不仅取决于你当前的输入，还取决于它目前所处的状态。例如，如果角色正在空中，那么按"下蹲"键可能没有效果；但如果角色在地面上，同样的按键却会让角色蹲下。这就是一个典型的状态跟踪任务。

在形式化的语言中，状态跟踪可以被定义为正确表示一个状态机在响应给定输入序列时所经历的任意长度的状态序列。这等同于在乔姆斯基形式语言层级结构中对有限自动机（FA）或正则语言的建模。

虽然状态跟踪看起来是神经网络应该能轻松学习的简单任务，但令人惊讶的是，许多模型在从数据中学习这类任务时表现不佳。这包括流行的Transformer模型，它们无法在比训练数据更长的序列上执行状态跟踪，即使是经过大规模预训练的基于Transformer的语言模型也是如此。

高通研究团队表明，双线性状态转换为表示状态跟踪任务中隐藏状态的演变提供了一种自然的归纳偏置。这种归纳偏置使得模型能够以一种更自然的方式学习状态转换，从而在处理状态跟踪任务时表现出色。

三、双线性RNN的数学原理与层级结构

那么，什么是双线性状态转换呢？让我们用一个简单的例子来解释：

想象你正在学习弹钢琴。每个音符（输入）对你的弹奏状态（隐藏状态）的影响不仅取决于音符本身，还取决于你当前的弹奏状态。比如，如果你正在弹奏一个和弦，那么下一个音符可能需要你改变手指的位置；而如果你正在弹奏单音，下一个音符的影响可能完全不同。这种"输入"和"当前状态"之间的相互作用就是双线性关系的一个比喻。

在数学上，双线性RNN的隐藏状态更新可以表示为：

h_t_i = (h_{t-1})^T W_i x_t = ∑_{jk} W_{ijk} x_t_k h_{t-1}_j

这里，h_t_i是时间t时隐藏状态向量的第i个元素，W_{ijk}是一个三维参数张量的元素，x_t是时间t的输入向量。这个公式本质上是说，当前时刻的隐藏状态是前一时刻隐藏状态与当前输入的乘法交互作用。

研究团队发现，双线性状态更新形成了一个自然的层级结构，对应于日益复杂的状态跟踪任务：

1. 最简单的形式是实数对角线双线性RNN，它可以学习奇偶校验等简单任务。 2. 复数对角线双线性RNN可以学习模块化加法等阿贝尔群结构任务。 3. 因子化/块对角线双线性RNN可以处理更复杂的任务。 4. 完全无约束的双线性RNN可以模拟任意状态机。

有趣的是，目前流行的线性循环网络，如Mamba，位于这个层级结构的最低复杂度中心，因此它们的状态跟踪能力是有限的。

四、研究发现：双线性RNN的惊人能力

研究团队通过理论分析和实验证明了几个关键发现：

首先，研究表明，对于隐藏状态来说，成为前一时刻隐藏状态的纯线性函数（而非仿射函数）总是足够的，在某些情况下甚至是必要的。这意味着，移除任何加性项（如偏置项）使隐藏状态具有尺度不变性，这反过来允许在训练和/或推理过程中对隐藏状态进行归一化，而不会牺牲线性递归的特性。

一个特别有趣的发现是，具有随机冻结权重和训练线性读出层的纯线性（非仿射）RNN可以从仅仅两个训练样本中学习奇偶校验任务，并达到完美的准确率。这就像一个初学者钢琴家只需要听两个例子就能完美地弹奏一首复杂的曲子！

此外，研究团队还发现，具有对角线转换矩阵的线性RNN是一种特殊情况，仅限于学习具有交换结构的状态跟踪任务。即使对于复值对角线转换矩阵也是如此。因此，具有2×2块对角线转换矩阵的线性RNN无法学习一般状态机（这是一个否定性结果）。

五、实验验证：双线性模型在复杂任务中的表现

为了验证理论分析，研究团队在三种状态跟踪任务上测试了各种模型：模块化加法、随机状态机和模块化算术。

在模块化加法任务中，模型处理一系列随机整数，并预测它们在模m下的和。对于随机状态机任务，模型必须模拟一个随机生成的有限状态机，其中输入字母表Σ和状态集Q都与Zm相同。在模块化算术任务中，模型需要处理交替出现的整数和算术运算符（如+、×、-），并计算这些操作依次应用的结果。

实验结果令人印象深刻：双线性模型在所有任务上的表现普遍最佳。随着块大小增加，双线性块对角线变体的性能也有所提高。值得注意的是，实数对角线模型（块大小为1的双线性块对角线模型）只能学习奇偶校验（即m=2的模块化加法）；但将块大小增加到2就能学习更大m值的模块化加法。

非线性循环模型，如LSTM和简单RNN，在这些状态跟踪任务上也表现良好。研究团队推测，隐藏状态和输入之间的乘法交互可能来自于门控机制和非线性激活函数。

而Mamba模型虽然可以学习小状态大小m的任务，但很大程度上无法推广到更长的序列。Transformer模型在长度泛化方面的失败也是文献中广为人知的观察结果。

六、数据效率与乘法vs加法交互的影响

关于双线性模型的一个重要问题是数据效率。由于参数数量随输入嵌入维度和隐藏维度的乘积增长，它们的参数数量可能非常大。

为了了解双线性模型的数据效率，研究团队在固定训练集大小的情况下训练和评估了各种模型。结果表明，尽管参数数量庞大，但这些模型的数据效率并不比LSTM差。这一点即使对于完全双线性模型也是如此。

另一个有趣的实验是关于乘法与加法交互的影响。研究团队发现，对于完全双线性模型，无论在哪种任务上，加性贡献都不会改变性能；但对于旋转模型，加性项会降低模块化加法任务上的性能。

七、学习奇偶校验的随机网络：一个令人惊讶的发现

最后一个实验展示了一个特别令人惊讶的发现：具有冻结随机权重和可训练线性读出层的纯双线性模型可以解决奇偶校验任务，即使循环参数在训练期间被冻结（只训练读出层）。

这就像给一个人一个随机设置的钢琴，却发现他们只需调整最后一个音符就能弹奏出完美的曲子！这与Mamba等模型形成鲜明对比，后者的状态转换是对角线且为正值，因此即使在学习期间调整循环参数也无法学习奇偶校验。

八、研究意义与未来展望

这项研究的意义在于重新思考了循环神经网络中隐藏单元的角色，从被动的记忆存储转变为计算的积极参与者。研究表明，双线性模型在学习状态跟踪任务方面非常有效，这与许多线性RNN如Mamba、LRU等形成对比。

不过，需要注意的是，双线性模型的参数数量会随隐藏状态数量的增加而大致立方增长。未来研究的一个重要问题是，是否有办法在保持强大的状态跟踪能力的同时减少参数数量。一个密切相关的问题是，这种减少是否可能适得其反（或者庞大的参数数量甚至可能有益）在大规模多任务场景中，如语言建模。

研究还建议区分有无交换结构的状态跟踪任务，后者可由更多类型的模型学习，特别是更小的模型。这引发了另一个未来研究方向：在哪些真实世界场景中，交换状态跟踪任务足够或不足够。

总的来说，这项研究重新审视了循环神经网络中的双线性状态转换，展示了它们在状态跟踪任务中的强大能力，并为未来的研究提供了有价值的见解。就像钢琴手不仅需要记住音符，还需要根据当前状态和输入来决定下一个动作一样，双线性RNN的隐藏单元也不仅仅是记忆存储器，而是计算过程中的积极参与者。