譯者 | 李睿

審校 | 重樓

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

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

Transformer的局限性和RNN的復(fù)興

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

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

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

Transformer的局限性重新引起了人們對循環(huán)模型的興趣。在過去的一年，引入了新的循環(huán)架構(gòu)，例如S4和Mamba，這些架構(gòu)有望在實(shí)現(xiàn)可比性能的同時解決Transformers的可擴(kuò)展性問題。這些模型使用“并行前綴掃描”等算法，通過并行化輸入序列上的計算來加快訓(xùn)練速度。

重新審視LSTM和GRU

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

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

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

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

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

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

研究人員寫道：“因此，在minGRU需要一天的時間來完成固定次數(shù)的訓(xùn)練的情況下，傳統(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的性能與最先進(jìn)的循環(huán)序列模型Mamba進(jìn)行了比較。他們測量了訓(xùn)練時間、內(nèi)存使用情況和多項任務(wù)的表現(xiàn)，包括選擇性復(fù)制、強(qiáng)化學(xué)習(xí)（RL）和語言建模。

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

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

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

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

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

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

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

盡管如此，這項研究的結(jié)果意義重大，因?yàn)檠芯勘砻鳎S著新信息的出現(xiàn)，重新審視原有的想法是值得的。

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

原文標(biāo)題：Minimized RNNs offer a fast and efficient alternative to Transformers，作者：Ben Dickson