本文介紹了長短期記憶網(wǎng)絡（LSTM）和門控循環(huán)單元（GRU）這兩種流行的RNN變體的最小化版本——minLSTM和minGRU，這兩個模型在訓練期間可以完全并行，并且使用的參數(shù)顯著減少，使它們成為Tansformer的快速和有效的替代方案。

Transformer如今已經(jīng)成為大型語言模型（LLM）和其他序列處理應用程序的主要架構。然而，它們固有的二次方計算復雜性成為了將Transformer擴展至超長序列時的巨大障礙，顯著增加了成本。這引發(fā)了人們對具有線性復雜性和恒定內(nèi)存需求的架構的興趣。

Mila公司和Borealis AI的研究人員重新審視了遞歸神經(jīng)網(wǎng)絡（RNN）作為Transformer架構的潛在替代方案。他們在發(fā)布的一篇論文中，介紹了長短期記憶網(wǎng)絡（LSTM）和門控循環(huán)單元（GRU）這兩種流行的RNN變體的最小化版本——minLSTM和minGRU，這兩個模型在訓練期間可以完全并行，并且使用的參數(shù)顯著減少，使它們成為Transformer的快速和有效的替代方案。

Transformer的局限性和RNN的復興

每當Transformer模型接收的輸入序列長度翻倍時，就需要四倍的內(nèi)存和計算量。這種二次方計算復雜性使得Transformer在處理長序列來說成本高昂，在資源受限的環(huán)境中尤為棘手。

另一方面，RNN按順序處理輸入數(shù)據(jù)，并且相對于序列長度具有線性計算復雜度。它們在推理過程中還需要恒定的內(nèi)存，這使得它們適合于非常長的序列。然而，傳統(tǒng)的RNN存在梯度消失和梯度爆炸的問題。當用于更新神經(jīng)網(wǎng)絡權值的梯度過小或過大，阻礙有效學習時，就會出現(xiàn)梯度消失和梯度爆炸。從而嚴重影響學習效果。這一局限性制約了RNN學習長期依賴關系的能力。為了解決這一問題，LSTM和GRU引入了創(chuàng)新的門控機制，以精確調(diào)控信息在網(wǎng)絡中的傳遞。

盡管傳統(tǒng)的LSTM和GRU具有優(yōu)勢，但它們有一個關鍵的限制：只能按順序計算。這意味著它們必須在訓練期間使用時間反向傳播（BPTT），這是一個緩慢的過程，極大地限制了它們擴展到長場景的能力。

重新審視LSTM和GRU

受到最近提出的序列模型之間算法相似性的啟發(fā)，研究人員重新審視了LSTM和GRU。他們發(fā)現(xiàn)，通過從門控機制中消除對先前隱藏狀態(tài)的依賴，這些模型可以使用并行掃描算法進行高效訓練。

傳統(tǒng)的LSTM和GRU有多個門來控制網(wǎng)絡中的信息流。這些門依賴于之前的隱藏狀態(tài)來確定保留或丟棄多少當前輸入和之前的內(nèi)存。這創(chuàng)建了一個順序依賴關系，要求模型一次處理一個令牌。

圖1 minGRU通過刪除組件和計算簡化了GRU

研究人員發(fā)現(xiàn)，他們可以消除對先前隱藏狀態(tài)的依賴，同時保持計算中的時間一致性。這使得能夠通過并行掃描算法訓練模型。他們通過刪除一些不必要的計算進一步簡化了架構，生成了最小化的LSTM（minLSTM）和最小化的GRU（minGRU），這兩種模型不僅使用的參數(shù)大幅減少，而且訓練速度也得到了顯著提升。

minGRU和minLSTM通過實現(xiàn)并行計算解決了傳統(tǒng)RNN的訓練瓶頸。在T4 GPU上，與傳統(tǒng)的序列長度為512個令牌相比，這一變化使得minGRU的速度提高了175倍，minLSTM的速度提高了235倍。隨著序列的延長，這種改善變得更加顯著。對于長度為4096的序列，minGRU和minLSTM的速度比傳統(tǒng)版本快1300倍以上。

圖2 minLSTM通過刪除不必要的組件和計算來簡化LSTM

研究人員寫道：“因此，在minGRU需要一天的時間來完成固定次數(shù)的訓練的情況下，傳統(tǒng)的GRU可能需要3年以上的時間?！?/p>

與傳統(tǒng)GRU相比，minGRU將所需參數(shù)的數(shù)量減少了87%，與傳統(tǒng)LSTM相比，minLSTM將所需參數(shù)的數(shù)量減少了85%。

最小化RNN與SOTA循環(huán)模型

研究人員將minLSTM和minGRU的性能與最先進的循環(huán)序列模型Mamba進行了比較。他們測量了訓練時間、內(nèi)存使用情況和多項任務的表現(xiàn)，包括選擇性復制、強化學習（RL）和語言建模。

在運行時間方面，minLSTM和minGRU取得了與Mamba相似的結果。雖然它們比傳統(tǒng)的RNN使用更多的內(nèi)存，但它們的內(nèi)存效率仍然比Mamba高。

圖3 訓練最小化 RNN 模型（橙色和藍色實線）比傳統(tǒng) RNN（虛線）更快，并且比 Mamba使用更少的內(nèi)存

在選擇性復制（需要內(nèi)容感知推理和記憶的任務）方面，minLSTM和minGRU的表現(xiàn)與Mamba相當。

在D4RL基準的RL實驗中，minLSTM和minGRU的性能優(yōu)于除Decision Mamba之外的所有基線，其中差異很小。

在語言建模任務中，minLSTM和minGRU在訓練過程中達到峰值性能的速度略慢于Mamba，但最終收斂的損失更低。值得注意的是，它們比Transformer的效率高得多，Transformer需要2.5倍的時間才能達到最佳性能。

圖4 在語言建模任務中，最小化的RNN（橙色和藍色實線）的損失低于Mamba，收斂速度至少是Transformers的2.5倍

與其他研究Transformer替代方案的類似工作一樣，最小化RNN研究的局限性之一是實驗的規(guī)模。這些架構是否能夠在非常大的模型和場景窗口下提供類似的結果還有待觀察。

盡管如此，這項研究的結果意義重大，因為研究表明，隨著新信息的出現(xiàn)，重新審視原有的想法是值得的。

研究人員在文章中寫道：“鑒于這些簡化的RNN在實際應用中所展現(xiàn)出的卓越性能，并且它們與眾多最新提出的循環(huán)序列方法存在著本質(zhì)上的相似之處，我們不禁反思并提出疑問，‘我們是否僅僅依靠RNN就足以滿足需求？’。”

原文標題：??Minimized RNNs offer a fast and efficient alternative to Transformers??，作者：Ben Dickson