深度學習三巨頭之一的Yoshua Bengio，剛剛發布了一篇有趣的新論文——
RNN 就是所需的全部嗎？
Were RNNs All We Needed?
不僅論文的名字有意思，其結論更是精彩。
研究表明，精簡十幾年前的 RNN 們，性能上可以與最近序列模型（如 Transformer 等）相媲美！
具體而言，Bengio 等人重新審視了諸如 LSTM（1997）和 GRU（2014）這些傳統的 RNN，認爲這些模型的缺點是由于需要時間反向傳播 （BPTT） 而導緻速度較慢。
所以他們直接大刀闊斧地移除了 LSTM 和 GRU 中的隐藏狀态依賴，讓它們不再需要 BPTT，從而可以高效地并行訓練。
而精簡改良版的 RNN 們，名字分别叫做minLSTM和minGRU。
它們和傳統 RNN 相比，不僅訓練時所需的參數量大幅減少，并且完全可并行化。
嗯，是頗有一種大道至簡的感覺了。
那麽 Bengio 等人具體又是如何實現的？我們繼續往下看。
精簡版 RNN
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 的最終形式爲：
Were RNNs All We Needed?
在精簡了 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
