RNN回归!图灵奖得主Bengio新作大道至简与Transformer一较高下
【导读】近日,深度学习三巨头之一的Yoshua Bengio,带领团队推出了全新的RNN架构,以大道至简的思想与Transformer一较高下。
在Transformer统治的AI时代之下,
散落在世界各地的「RNN神教」信徒,一直相信并期待着RNN回归的那天:
毕竟,凭借强大的顺序和上下文感知能力,RNN曾在各种任务中表现惊艳。
直到后来遭遇了反向训练的瓶颈,因Scaling Law而跌落神坛。
然而,人们并没有忘记RNN。
RWKV、Mamba、xLSTM等RNN衍生模型接连出现,欲挑战Transformer之霸主地位。
就在近日,又有重量级人物下场——
深度学习三巨头之一的Yoshua Bengio,带领团队推出了全新的RNN架构,以大道至简的思想与Transformer一较高下。
论文地址:https://arxiv.org/pdf/2410.01201v1
研究人员对传统的两种RNN架构LSTM和GRU,进行了大刀阔斧的改造,从中诞生了两个新模型:minLSTM和minGRU。
这俩极简主义的版本到底怎么样?咱们先看疗效。
首先是RNN最大的问题:训练速度。
上图展示了几种模型在T4 GPU上训练花费的时间,以及新模型带来的加速比。横轴为输入数据的序列长度,批量大小为64。
可以看到,相比于原版的LSTM和GRU,minLSTM、minGRU和Mamba的运行时间不会随序列长度而增加(后3个模型的线在左图中重叠了)。
当序列长度为4096时,新架构相对于传统版本达到了1300多倍的加速比!
相当于原版GRU需要3年才能做完的事情,minGRU一天就搞定了。
那么对线Transformer的战绩如何?
在本文测试的语言建模任务中,minGRU和minLSTM分别在600步左右达到最佳性能点。
相比之下,Transformer需要比minGRU多花大概2000步,训练速度慢了约2.5倍。
对此,YC上的网友表示:「我非常喜欢这个新架构的简单性」。
毕竟,俗话说的好,「最好的PR是那些删除代码的PR」。
模型架构
首先,这是传统的RNN架构:
LSTM在RNN的每个cell中加入了比较复杂的门控:
三个门控(input gate、output gate、forget gate)和输入的分量,都通过线性投影和非线性激活函数来得出,并且依赖于上一个时刻的隐藏状态ht-1。
这些值再经过线性和非线性计算,得到本时刻的输出ct和隐藏状态ht。
GRU在LSTM的基础上做了一些简化:
少了显式计算ct,用于门控的项也缩减到2个,相应的参数量和计算量也减少了。
那么我们就从相对简单的GRU入手,开始改造。
改造的目的是使RNN能够应用并行扫描(Parallel Scan)算法,解决自身训练困难的问题。
简单来说,就是将网络中的计算改造成vt = at ⊙ vt−1 + bt的形式。
minGRU
第一步,公式中含有对之前隐藏状态ht-1的依赖,没办法用并行扫描,所以把ht-1直接删掉。
ht-1没了,负责调控ht-1的rt也没用了,删掉。
第二步,双曲正切函数(tanh)负责限制隐藏状态的范围,并减轻因sigmoid(σ)而导致的梯度消失。
但是现在ht-1和rt都没了,tanh也失去了存在的意义,删掉。
那么最终,minGRU就是下面这三个公式:
相比于原版,参数量和计算量再次减少,最重要的是能够使用并行扫描来显著加快训练速度。
minLSTM
经过上面的叙述,minLSTM的由来就很好理解了。
首先还是去除隐藏状态的依赖:
接着是拿掉相关的tanh:
最后,为了保证LSTM输出的尺度与时间无关,以及hidden state在缩放上与时间无关,还需要删掉output gate。
output gate没了,ct也就没必要单独存在了,删掉;剩下的两个门控通过归一化来调配hidden state进入的比例。
——emmm......好像变成GRU了,算了不管了。
最终改造好的minLSTM是下面这个样子:
Were RNNs All We Needed?
别急,先打内战证明价值。
除了传统的RNN(LSTM和GRU),这里特别关注与Mamba的比较。
首先是训练上的提升:
实验在批次大小64的情况下改变序列长度,测量了模型执行前向传递、计算损失和向后传递计算梯度的总运行时间以及内存占用。
在运行时间方面,minLSTM、minGRU与Mamba实现了类似的效率。
序列长度为512时的运行时间(超过100次的平均值),分别为 2.97、2.72和2.71毫秒;序列长度为4096时,运行时间分别为3.41、3.25和3.15。
相比之下,LSTM和GRU的运行时间随序列长度线性增加。所以序列长度为512时,minGRU和minLSTM的训练加速了175倍和235倍;序列长度为4096时,加速比达到了1324和1361。
内存方面,利用并行扫描算法时会创建更大的计算图,所以minGRU、minLSTM和Mamba ,比传统RNN需要更多的内存(大概多出88%)。
——但这并不重要,因为对于RNN来说,训练时间才是瓶颈。
去除隐藏状态的效果
minLSTM和minGRU的训练效率是通过降低它们的门控对先前隐藏状态的依赖来实现的。
尽管单层minLSTM或minGRU的门控只与输入有关,而与时间无关,但是在深度学习中,模型是通过堆叠模块来构建的。
从第二层开始,minLSTM和minGRU的门也将与时间相关,从而对更复杂的函数进行建模。
下表比较了不同层数的模型在选择性复制任务上的性能。我们可以看到时间依赖性的影响:将层数增加会大大提高模型的性能。
训练稳定性
层数的另一个影响是稳定性,随着层数的增加,精度的方差减小。
此外,尽管minLSTM和minGRU都解决了选择性复制任务,但我们可以看到minGRU在经验上是一种比minLSTM更稳定的方法(更高的一致性和更低的方差)。
minLSTM丢弃旧信息并添加新信息,使用两组参数(forget gate 和input gate)控制比率。在训练期间,两组参数会向不同的方向进行调整,使得比率更难控制和优化。相比之下,minGRU的丢弃和添加信息由一组参数控制,更容易优化。
选择性复制
选择性复制任务的输入元素相对于其输出是随机间隔的,为了解决这项任务,模型需要执行内容感知推理,记住相关token并过滤掉不相关的token。
上表将minLSTM和minGRU与可以并行训练的知名RNN模型进行了比较(S4,H3,Hyena和Mamba(S6)),基线结果引自Mamba论文。
在所有这些基线中,只有Mamba的S6,以及本文的minGRU和minLSTM能够解决此任务,体现了LSTM和GRU的内容感知门控机制。
强化学习
下面开始对战Transformer。
考虑D4RL基准中的MuJoCo运动任务,包括三个环境:HalfCheetah、Hopper和Walker。
对于每个环境,模型在三个数据质量不同的数据集上进行训练:Medium(M)、Medium-Replay(M-R)和Medium-Expert(M-E)。
上表将minLSTM和minGRU与各种决策模型进行了比较,包括原始的Decision Transformer(DT)、Decision S4 (DS4) 、Decision Mamba和Aaren。
由结果可知,minLSTM和minGRU的性能优于Decision S4,与Decision Transformer、Aaren和Mamba相媲美(Decision S4的递归转换不是输入感知的,这会影响它的性能)。就平均分数而言,minLSTM和minGRU的表现优于除Decision Mamba之外的所有基线。
语言建模
最后考虑语言建模任务,使用nanoGPT框架在莎士比亚的作品上训练字符级GPT。
上图绘制了具有交叉熵损失的学习曲线,可以发现minGRU、 minLSTM、 Mamba和Transformers分别实现了1.548、1.555、1.575和1.547的可比测试损耗。
Mamba的表现略差于其他模型,但训练速度更快(400步),minGRU和minLSTM分别花费575步和625步。而Transformer直接比minGRU多了2000 步,慢了大概2.5倍。