金磊发自凹非寺
量子位公众号 QbitAI
深度学习三巨头之一的Yoshua Bengio，刚刚发布了一篇有趣的新论文——RNN 就是所需的全部吗？
Were RNNs All We Needed?
不仅论文的名字有意思，其结论更是精彩。
研究表明，精简十几年前的 RNN 们，性能上可以与最近序列模型（如 Transformer 等）相媲美！
具体而言，Bengio 等人重新审视了诸如 LSTM（1997）和 GRU（2014）这些传统的 RNN，认为这些模型的缺点是由于需要时间反向传播 （BPTT） 而导致速度较慢。
所以他们直接大刀阔斧地移除了 LSTM 和 GRU 中的隐藏状态依赖，让它们不再需要 BPTT，从而可以高效地并行训练。
而精简改良版的 RNN 们，名字分别叫做minLSTM和minGRU。
它们和传统 RNN 相比，不仅训练时所需的参数量大幅减少，并且完全可并行化。
嗯，是颇有一种大道至简的感觉了。
那么 Bengio 等人具体又是如何实现的？我们继续往下看。
精简版 RNN
Transformer 和它的变体们可以说是近几年大热的架构，但与此同时缺点也是较为明显，那便是在处理长序列时的计算复杂度问题。
具体来说，Transformer 模型在序列长度上的计算复杂度是二次方的，这使得它在处理长序列时资源的消耗就比较高。
因此就需要能够在训练时有效地处理长序列，同时在推理时保持高效性能的替代方案——简化版的 RNN。
这个过程的关键便是隐藏状态依赖，让它们不再需要 BPTT，让效率直接飙升。
minGRU
首先我们来看下 Bengio 团队对 GRU 的处理，即minGRU，总共分为 2 步。
第一步，去除之前隐藏状态的依赖。
在传统的 GRU 模型中，更新门 zt 和候选隐藏状态h~t的计算依赖于前一时刻的隐藏状态 ht-1。这导致模型在训练时无法实现并行处理，因为每个时间步的计算都依赖于前一个时间步的结果。
为了解决这个问题，minGRU 对 GRU 进行了修改，使更新门和候选隐藏状态的计算仅依赖于当前时刻的输入 xt，而不依赖于 ht-1：
通过这种方式，minGRU 的每一时刻的计算可以独立于其他时刻并行执行。
第二步，去除候选状态的范围限制。
在第一步中，候选隐藏状态h~t仍然使用双曲正切函数（tanh）来限制其值的范围在 [−1,1][−1,1] 之间。虽然这有助于模型的稳定性，但它并不是并行化所必需的。
minGRU 进一步简化模型，去除了对h~t的范围限制，将其替换为一个无需激活函数的线性变换：
这样，候选隐藏状态的计算变得更加简单，并且没有任何范围限制。
在这种结构下，minGRU 不仅减少了模型参数，而且可以利用并行扫描算法在训练时实现并行化，从而显著提高了处理长序列的速度。
此外，minGRU 的输出尺度在时间上是独立的，这有助于优化过程中的数值稳定性。整体变化如下：
minLSTM
接下来，我们再来看下 Bengio 团队对 LSTM 的处理，即minLSTM，共分为三步。
第一步，去除之前隐藏状态的依赖。
在传统的 LSTM 模型中，遗忘门 ft、输入门 it 和候选细胞状态c~t的计算依赖于前一时刻的隐藏状态 ht-1。
这导致模型在训练时无法实现并行处理，因为每个时间步的计算都依赖于前一个时间步的结果。
为了解决这个问题，minLSTM 对 LSTM 进行了修改，使遗忘门、输入门和候选细胞状态的计算仅依赖于当前时刻的输入 xt，而不依赖于 ht-1：
通过这种方式，minLSTM 的每一时刻的计算可以独立于其他时刻并行执行。
第二步，去除候选状态的范围限制。
在第一步中，候选细胞状态c~t仍然使用双曲正切函数（tanh）来限制其值的范围在 [−1,1][−1,1] 之间。虽然这有助于模型的稳定性，但它并不是并行化所必需的。
minLSTM 进一步简化模型，去除了对c~t的范围限制，将其替换为一个无需激活函数的线性变换：
这样，候选细胞状态的计算变得更加简单，并且没有任何范围限制。
第三步，确保输出在时间上是独立的。
在许多序列建模设置中（例如文本生成），优化目标/输出在时间上是独立的。
为了确保 LSTM 的输出在时间上是独立的，minLSTM 对遗忘门和输入门进行了归一化，确保它们的和为1，并且细胞状态的尺度在时间上是独立的：
通过这种方式，minLSTM 确保了其输出在时间上是独立的，这有助于优化过程中的数值稳定性。
minLSTM 的最终形式为：
Were RNNs All We Needed?
在精简了 RNN 们之后，Bengio 团队也展示了实验结果。
例如下图显示了 minGRU、minLSTM 和 Mamba 模型在训练效率方面的比较，具体包括训练运行时间、加速比和内存占用。
这些指标是在 T4 GPU 上，以 64 的批次大小进行测量的：
以及在下图中，还展示了在 Shakespeare 语言建模任务中，不同模型的学习曲线。
这个任务使用字符级生成对抗训练，目的是评估模型在文本生成任务中的表现，简化 RNN 模型在处理语言建模任务时具有较好的有效性和高效率（特别是在需要快速训练和部署的应用场景中）：
总而言之，Bengio 团队认为，经过简化的 RNN 可能仍然是处理长序列任务的理想选择，尤其是在资源有限的场景下，因此也提出了问题 “Were RNNs All We Needed?”
华人一作
在这项研究中，作者除了 Bengio 之外，还有一点值得关注，那便是一作是一位华人，Leo Feng。
从公开的个人网站来看，Leo Feng 师从 Bengio，目前是蒙特利尔大学的博士生，目前正在 Borealis AI 进行研究实习。
Leo Feng 的研究范围包括元学习和高效模型的设计，其本科毕业于牛津大学。
那么你觉得精简版 RNN 这项研究如何？欢迎在评论区留言讨论。
论文地址：
