Bengio精简了传统RNN，性能可与Transformer媲美

金磊 发自 凹非寺
量子位 | 公众号 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这项研究如何？欢迎在评论区留言讨论。

论文地址：
https://arxiv.org/abs/2410.01201