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一較高下。

研究人員對傳統的兩種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搞出來了，能打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