RNN回归!Bengio新作大道至简与Transformer一较高下

新智元报道

:alan

【新智元导读】 近日,深度学习三巨头之一的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的形式。

第一步,公式中含有对之前隐藏状态ht-1的依赖,没办法用并行扫描,所以把ht-1直接删掉。

ht-1没了,负责调控ht-1的rt也没用了,删掉。

第二步,双曲正切函数(tanh)负责限制隐藏状态的范围,并减轻因sigmoid(σ)而导致的梯度消失。

但是现在ht-1和rt都没了,tanh也失去了存在的意义,删掉。

那么最终,minGRU就是下面这三个公式:

相比于原版,参数量和计算量再次减少,最重要的是能够使用并行扫描来显著加快训练速度。

经过上面的叙述,minLSTM的由来就很好理解了。

首先还是去除隐藏状态的依赖:

接着是拿掉相关的tanh:

最后,为了保证LSTM输出的尺度与时间无关,以及hidden state在缩放上与时间无关,还需要删掉output gate。

output gate没了,ct也就没必要单独存在了,删掉;剩下的两个门控通过归一化来调配hidden state进入的比例。

——emmm......好像变成GRU了,算了不管了。

最终改造好的minLSTM是下面这个样子:

Were RNNs All We Needed?

全新的RNN搞出来了,能打Transformer吗?

别急,先打内战证明价值。

除了传统的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倍。

参考资料:

https://arxiv.org/pdf/2410.01201v1


1982年,美国加州理工学院物理学家John hopfield 发明了一种单层反馈神经网络 Hopfield network,用来解决组合优化问题。 这是最早的RNN的雏形。 86年,michael I. Jordan 定义了recurrent的概念,提出 Jordan network。 1990年, 美国认知科学家Jeffrey L. Elman 对jordan network进行了简化,并采用BP算法进行训练,便有了如今最简单的包含单个自连接节点的RNN 模型。 但此时RNN由于梯度消失(gradient vanishing)及梯度爆炸(gradient exploding)的问题,训练非常困难,应用非常受限。 直到1997年,人工智能研究所的主任Jurgen Schmidhuber 提出长短期记忆(LSTM),LSTM使用门告散控单元及记忆机制大大缓解了早期RNN训练的问题。 同样在1997年,Mike Schuster 提出双向RNN模型(Bidirectional RNN)。 这两种模型大大改进了早期RNN结构,拓宽了RNN的应用范围,为后续序列建模的发展奠定了基础。 此时RNN虽然在一些序列建模任务上取得了不错的效果,但由于计算资源消耗大,后续几年一直没有太大的进展。 2010年,Tomas Mikolov对bengio提出的feedforward Neural network language model (NNLM) 进行了改进,提出了颤友陆基于RNN的语言模型(RNN LM),并将其用在语音识别任务中,大幅提升了识别精度。 再此基础上Tomas Mikolov于2013年提出了大名鼎鼎的word2vec,与NNLM及RNNLM不同,word2vec的目标不再专注于建模语言模型,而是如何利用语言模型学习每个单词的语义化向量(distributed representation),其中distributed representation概念最早要来源于Hinton 1986年的工作。 Word2vec引发了深度学习在自然语言处理领域的浪潮,除此之外还启发了knowledge representation,network representation等新的领域。 另一方面,2014年,Bengio团队与google几乎同时提出了seq2seq架构,将RNN用于机器翻译。 没过多久,Bengio团队又提出attention机制,对seq2seq架构进行改进。 自此机器翻译全面进入到神经机器翻译(NMT)的时代,NMT不仅过程简单,而且效果要远超统计机器翻译的效果。 目前主流的机器翻译系统几乎都采用了神经机器翻译的技术。 除此之外,attention机制也被广泛用于基于深度学习的各种任务中。 近两年,相关领域仍有一些突破性进展,2017年,facebook人工智能实验室提出基于卷积神经网络的seq2seq架构,将rnn替换为带有门控单元的cnn,提升效果的同时大幅加快了模型训练速度,此后不久,google提出transformer架构,使用self-attention代替原有的RNN及CNN,更进一步降低了模型复杂度茄顷。 在词表示学习方面,Allen人工智能研究所2018年提出上下文相关的表示学习方法ELMo,利用双向LSTM语言模型对不同语境下的单词学习不同的向量表示,在6个nlp任务上取得了提升。 OpenAI团队在此基础上提出预训练模型GPT,把LSTM替换为transformer来训练语言模型,在应用到具体任务时,与之前学习词向量当作特征的方式不同,GPT直接在预训练得到的语言模型最后一层接上softmax作为任务输出层,然后再对模型进行微调,在多项任务上GPT取得了更好的效果。 不久之后,Google提出BERT模型,将GPT中的单向语言模型拓展为双向语言模型(Masked Language Model),并在预训练中引入了sentence prediction任务。 BERT模型在11个任务中取得了最好的效果,是深度学习在nlp领域又一个里程碑式的工作。

首发 多功能支架版预售 519 小米 显示器 元 Redmi
佳源国际 资本退潮与行业洗牌 大唐地产集体退市 大发地产