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

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

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

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

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

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

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

重新審視LSTM和GRU

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

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

圖1 minGRU通過(guò)刪除組件和計(jì)算簡(jiǎn)化了GRU

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

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

圖2 minLSTM通過(guò)刪除不必要的組件和計(jì)算來(lái)簡(jiǎn)化LSTM

研究人員寫道：“因此，在minGRU需要一天的時(shí)間來(lái)完成固定次數(shù)的訓(xùn)練的情況下，傳統(tǒng)的GRU可能需要3年以上的時(shí)間。”

與傳統(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)行了比較。他們測(cè)量了訓(xùn)練時(shí)間、內(nèi)存使用情況和多項(xiàng)任務(wù)的表現(xiàn)，包括選擇性復(fù)制、強(qiáng)化學(xué)習(xí)（RL）和語(yǔ)言建模。

在運(yùn)行時(shí)間方面，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之外的所有基線，其中差異很小。

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

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

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

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

研究人員在文章中寫道：“鑒于這些簡(jiǎ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