金磊 发自 凹非寺
量子位 | 公众号 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
