Bengio精簡了傳統RNN，性能可與Transformer媲美

深度學習三巨頭之一的Yoshua Bengio，剛剛發佈了一篇有趣的新論文——

不僅論文的名字有意思，其結論更是精彩。

研究表明，精簡十幾年前的RNN們，性能上可以與最近序列模型（如Transformer等）相媲美！

具體而言，Bengio等人重新審視了諸如LSTM（1997）和GRU（2014）這些傳統的RNN，認爲這些模型的缺點是由於需要時間反向傳播 （BPTT） 而導致速度較慢。

所以他們直接大刀闊斧地移除了LSTM和GRU中的隱藏狀態依賴，讓它們不再需要BPTT，從而可以高效地並行訓練。

而精簡改良版的RNN們，名字分別叫做minLSTM和minGRU。

它們和傳統RNN相比，不僅訓練時所需的參數量大幅減少，並且完全可並行化。

嗯，是頗有一種大道至簡的感覺了。

那麼Bengio等人具體又是如何實現的？我們繼續往下看。

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的最終形式爲：

在精簡了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