在Transformer統(tǒng)治的AI時(shí)代之下，

散落在世界各地的「RNN神教」信徒，一直相信并期待著RNN回歸的那天：

圖片

畢竟，憑借強(qiáng)大的順序和上下文感知能力，RNN曾在各種任務(wù)中表現(xiàn)驚艷。

直到后來(lái)遭遇了反向訓(xùn)練的瓶頸，因Scaling Law而跌落神壇。

然而，人們并沒(méi)有忘記RNN。

圖片

RWKV、Mamba、xLSTM等RNN衍生模型接連出現(xiàn)，欲挑戰(zhàn)Transformer之霸主地位。

就在近日，又有重量級(jí)人物下場(chǎng)——

深度學(xué)習(xí)三巨頭之一的Yoshua Bengio，帶領(lǐng)團(tuán)隊(duì)推出了全新的RNN架構(gòu)，以大道至簡(jiǎn)的思想與Transformer一較高下。

圖片

論文地址：https://arxiv.org/pdf/2410.01201v1

研究人員對(duì)傳統(tǒng)的兩種RNN架構(gòu)LSTM和GRU，進(jìn)行了大刀闊斧的改造，從中誕生了兩個(gè)新模型：minLSTM和minGRU。

這倆極簡(jiǎn)主義的版本到底怎么樣？咱們先看療效。

首先是RNN最大的問(wèn)題：訓(xùn)練速度。

圖片

上圖展示了幾種模型在T4 GPU上訓(xùn)練花費(fèi)的時(shí)間，以及新模型帶來(lái)的加速比。橫軸為輸入數(shù)據(jù)的序列長(zhǎng)度，批量大小為64。

可以看到，相比于原版的LSTM和GRU，minLSTM、minGRU和Mamba的運(yùn)行時(shí)間不會(huì)隨序列長(zhǎng)度而增加（后3個(gè)模型的線在左圖中重疊了）。

當(dāng)序列長(zhǎng)度為4096時(shí)，新架構(gòu)相對(duì)于傳統(tǒng)版本達(dá)到了1300多倍的加速比！

相當(dāng)于原版GRU需要3年才能做完的事情，minGRU一天就搞定了。

那么對(duì)線Transformer的戰(zhàn)績(jī)?nèi)绾危?/span>

圖片

在本文測(cè)試的語(yǔ)言建模任務(wù)中，minGRU和minLSTM分別在600步左右達(dá)到最佳性能點(diǎn)。

相比之下，Transformer需要比minGRU多花大概2000步，訓(xùn)練速度慢了約2.5倍。

對(duì)此，YC上的網(wǎng)友表示：「我非常喜歡這個(gè)新架構(gòu)的簡(jiǎn)單性」。

圖片

畢竟，俗話說(shuō)的好，「最好的PR是那些刪除代碼的PR」。

模型架構(gòu)

下面來(lái)感受一下極簡(jiǎn)模型的誕生過(guò)程。

首先，這是傳統(tǒng)的RNN架構(gòu)：

圖片

LSTM在RNN的每個(gè)cell中加入了比較復(fù)雜的門控：

圖片

三個(gè)門控（input gate、output gate、forget gate）和輸入的分量，都通過(guò)線性投影和非線性激活函數(shù)來(lái)得出，并且依賴于上一個(gè)時(shí)刻的隱藏狀態(tài)ht-1。

圖片

這些值再經(jīng)過(guò)線性和非線性計(jì)算，得到本時(shí)刻的輸出ct和隱藏狀態(tài)ht。

GRU在LSTM的基礎(chǔ)上做了一些簡(jiǎn)化：

圖片

少了顯式計(jì)算ct，用于門控的項(xiàng)也縮減到2個(gè)，相應(yīng)的參數(shù)量和計(jì)算量也減少了。

圖片

那么我們就從相對(duì)簡(jiǎn)單的GRU入手，開(kāi)始改造。

改造的目的是使RNN能夠應(yīng)用并行掃描（Parallel Scan）算法，解決自身訓(xùn)練困難的問(wèn)題。

簡(jiǎn)單來(lái)說(shuō)，就是將網(wǎng)絡(luò)中的計(jì)算改造成vt = at ⊙ vt?1 + bt的形式。

minGRU

第一步，公式中含有對(duì)之前隱藏狀態(tài)ht-1的依賴，沒(méi)辦法用并行掃描，所以把ht-1直接刪掉。

圖片

ht-1沒(méi)了，負(fù)責(zé)調(diào)控ht-1的rt也沒(méi)用了，刪掉。

第二步，雙曲正切函數(shù)（tanh）負(fù)責(zé)限制隱藏狀態(tài)的范圍，并減輕因sigmoid（σ）而導(dǎo)致的梯度消失。

但是現(xiàn)在ht-1和rt都沒(méi)了，tanh也失去了存在的意義，刪掉。

圖片

那么最終，minGRU就是下面這三個(gè)公式：

圖片

相比于原版，參數(shù)量和計(jì)算量再次減少，最重要的是能夠使用并行掃描來(lái)顯著加快訓(xùn)練速度。

minLSTM

經(jīng)過(guò)上面的敘述，minLSTM的由來(lái)就很好理解了。

首先還是去除隱藏狀態(tài)的依賴：

圖片

接著是拿掉相關(guān)的tanh：

圖片

最后，為了保證LSTM輸出的尺度與時(shí)間無(wú)關(guān)，以及hidden state在縮放上與時(shí)間無(wú)關(guān)，還需要?jiǎng)h掉output gate。

output gate沒(méi)了，ct也就沒(méi)必要單獨(dú)存在了，刪掉；剩下的兩個(gè)門控通過(guò)歸一化來(lái)調(diào)配hidden state進(jìn)入的比例。

圖片

——emmm......好像變成GRU了，算了不管了。

最終改造好的minLSTM是下面這個(gè)樣子：

圖片

Were RNNs All We Needed?

全新的RNN搞出來(lái)了，能打Transformer嗎？

別急，先打內(nèi)戰(zhàn)證明價(jià)值。

除了傳統(tǒng)的RNN（LSTM和GRU），這里特別關(guān)注與Mamba的比較。

首先是訓(xùn)練上的提升：

圖片

實(shí)驗(yàn)在批次大小64的情況下改變序列長(zhǎng)度，測(cè)量了模型執(zhí)行前向傳遞、計(jì)算損失和向后傳遞計(jì)算梯度的總運(yùn)行時(shí)間以及內(nèi)存占用。

在運(yùn)行時(shí)間方面，minLSTM、minGRU與Mamba實(shí)現(xiàn)了類似的效率。

序列長(zhǎng)度為512時(shí)的運(yùn)行時(shí)間（超過(guò)100次的平均值），分別為 2.97、2.72和2.71毫秒；序列長(zhǎng)度為4096時(shí)，運(yùn)行時(shí)間分別為3.41、3.25和3.15。

相比之下，LSTM和GRU的運(yùn)行時(shí)間隨序列長(zhǎng)度線性增加。所以序列長(zhǎng)度為512時(shí)，minGRU和minLSTM的訓(xùn)練加速了175倍和235倍；序列長(zhǎng)度為4096時(shí)，加速比達(dá)到了1324和1361。

內(nèi)存方面，利用并行掃描算法時(shí)會(huì)創(chuàng)建更大的計(jì)算圖，所以minGRU、minLSTM和Mamba ，比傳統(tǒng)RNN需要更多的內(nèi)存（大概多出88%）。

——但這并不重要，因?yàn)閷?duì)于RNN來(lái)說(shuō)，訓(xùn)練時(shí)間才是瓶頸。

去除隱藏狀態(tài)的效果

minLSTM和minGRU的訓(xùn)練效率是通過(guò)降低它們的門控對(duì)先前隱藏狀態(tài)的依賴來(lái)實(shí)現(xiàn)的。

盡管單層minLSTM或minGRU的門控只與輸入有關(guān)，而與時(shí)間無(wú)關(guān)，但是在深度學(xué)習(xí)中，模型是通過(guò)堆疊模塊來(lái)構(gòu)建的。

從第二層開(kāi)始，minLSTM和minGRU的門也將與時(shí)間相關(guān)，從而對(duì)更復(fù)雜的函數(shù)進(jìn)行建模。

下表比較了不同層數(shù)的模型在選擇性復(fù)制任務(wù)上的性能。我們可以看到時(shí)間依賴性的影響：將層數(shù)增加會(huì)大大提高模型的性能。

圖片

訓(xùn)練穩(wěn)定性

層數(shù)的另一個(gè)影響是穩(wěn)定性，隨著層數(shù)的增加，精度的方差減小。

此外，盡管minLSTM和minGRU都解決了選擇性復(fù)制任務(wù)，但我們可以看到minGRU在經(jīng)驗(yàn)上是一種比minLSTM更穩(wěn)定的方法（更高的一致性和更低的方差）。

minLSTM丟棄舊信息并添加新信息，使用兩組參數(shù)（forget gate 和input gate）控制比率。在訓(xùn)練期間，兩組參數(shù)會(huì)向不同的方向進(jìn)行調(diào)整，使得比率更難控制和優(yōu)化。相比之下，minGRU的丟棄和添加信息由一組參數(shù)控制，更容易優(yōu)化。

選擇性復(fù)制

選擇性復(fù)制任務(wù)的輸入元素相對(duì)于其輸出是隨機(jī)間隔的，為了解決這項(xiàng)任務(wù)，模型需要執(zhí)行內(nèi)容感知推理，記住相關(guān)token并過(guò)濾掉不相關(guān)的token。

圖片

上表將minLSTM和minGRU與可以并行訓(xùn)練的知名RNN模型進(jìn)行了比較（S4，H3，Hyena和Mamba（S6）)，基線結(jié)果引自Mamba論文。

在所有這些基線中，只有Mamba的S6，以及本文的minGRU和minLSTM能夠解決此任務(wù)，體現(xiàn)了LSTM和GRU的內(nèi)容感知門控機(jī)制。

強(qiáng)化學(xué)習(xí)

下面開(kāi)始對(duì)戰(zhàn)Transformer。

考慮D4RL基準(zhǔn)中的MuJoCo運(yùn)動(dòng)任務(wù)，包括三個(gè)環(huán)境：HalfCheetah、Hopper和Walker。

對(duì)于每個(gè)環(huán)境，模型在三個(gè)數(shù)據(jù)質(zhì)量不同的數(shù)據(jù)集上進(jìn)行訓(xùn)練：Medium（M）、Medium-Replay（M-R）和Medium-Expert（M-E）。

圖片

上表將minLSTM和minGRU與各種決策模型進(jìn)行了比較，包括原始的Decision Transformer（DT）、Decision S4 （DS4） 、Decision Mamba和Aaren。

由結(jié)果可知，minLSTM和minGRU的性能優(yōu)于Decision S4，與Decision Transformer、Aaren和Mamba相媲美（Decision S4的遞歸轉(zhuǎn)換不是輸入感知的，這會(huì)影響它的性能）。就平均分?jǐn)?shù)而言，minLSTM和minGRU的表現(xiàn)優(yōu)于除Decision Mamba之外的所有基線。

語(yǔ)言建模

最后考慮語(yǔ)言建模任務(wù)，使用nanoGPT框架在莎士比亞的作品上訓(xùn)練字符級(jí)GPT。

圖片

上圖繪制了具有交叉熵?fù)p失的學(xué)習(xí)曲線，可以發(fā)現(xiàn)minGRU、 minLSTM、 Mamba和Transformers分別實(shí)現(xiàn)了1.548、1.555、1.575和1.547的可比測(cè)試損耗。

Mamba的表現(xiàn)略差于其他模型，但訓(xùn)練速度更快（400步），minGRU和minLSTM分別花費(fèi)575步和625步。而Transformer直接比minGRU多了2000 步，慢了大概2.5倍。

參考資料：https://arxiv.org/pdf/2410.01201v1