20 世紀(jì) 90 年代，長(zhǎng)短時(shí)記憶（LSTM）方法引入了恒定誤差選擇輪盤(pán)和門(mén)控的核心思想。三十多年來(lái)，LSTM 經(jīng)受住了時(shí)間的考驗(yàn)，并為眾多深度學(xué)習(xí)的成功案例做出了貢獻(xiàn)。然而，以可并行自注意力為核心 Transformer 橫空出世之后，LSTM 自身所存在的局限性使其風(fēng)光不再。

當(dāng)人們都以為 Transformer 在語(yǔ)言模型領(lǐng)域穩(wěn)坐江山的時(shí)候，LSTM 又殺回來(lái)了 —— 這次，是以 xLSTM 的身份。

5 月 8 日，LSTM 提出者和奠基者 Sepp Hochreiter 在 arXiv 上傳了 xLSTM 的預(yù)印本論文。

論文的所屬機(jī)構(gòu)中還出現(xiàn)了一家叫做「NXAI」的公司，Sepp Hochreiter 表示：「借助 xLSTM，我們縮小了與現(xiàn)有最先進(jìn) LLM 的差距。借助 NXAI，我們已開(kāi)始構(gòu)建自己的歐洲 LLM。」

• 論文標(biāo)題：xLSTM: Extended Long Short-Term Memory
• 論文鏈接：https://arxiv.org/pdf/2405.04517

具體來(lái)說(shuō)，xLSTM 從三個(gè)層面解決了 LSTM 以往所存在的局限性：

(i) 無(wú)法修改存儲(chǔ)決策。

可以通過(guò)「最近鄰搜索」（Nearest Neighbor Search）問(wèn)題來(lái)舉例說(shuō)明這一局限性：在給定參考向量的情況下，必須按順序掃描序列，尋找最相似的向量，以便在序列末端提供其附加值。圖 2 左側(cè)顯示了這項(xiàng)任務(wù)的均方誤差。當(dāng)發(fā)現(xiàn)更相似的向量時(shí)，LSTM 難以修改存儲(chǔ)的值，而新的 xLSTM 通過(guò)指數(shù)門(mén)控彌補(bǔ)了這一限制。

(ii) 存儲(chǔ)容量有限，即信息必須壓縮成標(biāo)量單元狀態(tài)。

圖 2 右側(cè)給出了 Wikitext103 上不同 token 頻率的 token 預(yù)測(cè)困惑度。由于 LSTM 的存儲(chǔ)容量有限，它在不常見(jiàn) token 上的表現(xiàn)較差。xLSTM 通過(guò)矩陣內(nèi)存解決了這一問(wèn)題。

(iii) 由于內(nèi)存混合而缺乏可并行性，需要進(jìn)行順序處理。例如，從一個(gè)時(shí)間步到下一個(gè)時(shí)間步的隱藏狀態(tài)之間的隱藏 - 隱藏連接。

與此同時(shí)，Sepp Hochreiter 和團(tuán)隊(duì)在這篇新論文中回答了一個(gè)關(guān)鍵問(wèn)題：如果克服這些局限性并將 LSTM 擴(kuò)展到當(dāng)前大語(yǔ)言模型的規(guī)模時(shí)，能實(shí)現(xiàn)怎樣的性能？

將 LSTM 擴(kuò)展到數(shù)十億參數(shù)

為了克服 LSTM 的局限性，xLSTM 對(duì)等式（1）中的 LSTM 理念進(jìn)行了兩項(xiàng)主要修改。

在原來(lái)的 LSTM 中，恒定誤差選擇輪盤(pán)是由單元輸入 z_t 對(duì)單元狀態(tài) c_(t-1)（綠色）進(jìn)行的加法更新，并由 sigmoid 門(mén)（藍(lán)色）進(jìn)行調(diào)節(jié)。輸入門(mén) i_t 和遺忘門(mén) f_t 控制這一更新，而輸出門(mén) o_t 控制存儲(chǔ)單元的輸出，即隱藏狀態(tài) h_t。存儲(chǔ)單元的狀態(tài)被 ψ 歸一化或壓縮，然后輸出門(mén)控得到隱藏狀態(tài)。

xLSTM 的修改包括指數(shù)門(mén)控和新穎的內(nèi)存結(jié)構(gòu)，因此豐富了 LSTM 家族的兩個(gè)成員：

(i) sLSTM（第 2.2 節(jié)），具有標(biāo)量?jī)?nèi)存、標(biāo)量更新和內(nèi)存混合功能；

(ii) mLSTM（第 2.3 節(jié)），具有矩陣內(nèi)存和協(xié)方差（外積）更新規(guī)則，完全可并行處理。

sLSTM 和 mLSTM 都通過(guò)指數(shù)門(mén)控增強(qiáng)了 LSTM。為了實(shí)現(xiàn)并行化，mLSTM 放棄了內(nèi)存混合，即隱藏 - 隱藏遞歸連接。mLSTM 和 sLSTM 都可以擴(kuò)展到多個(gè)存儲(chǔ)單元，其中 sLSTM 具有跨單元內(nèi)存混合的特點(diǎn)。此外，sLSTM 可以有多個(gè)頭，但不存在跨頭的內(nèi)存混合，而只存在每個(gè)頭內(nèi)單元間的內(nèi)存混合。通過(guò)引入 sLSTM 頭和指數(shù)門(mén)控，研究者建立了一種新的內(nèi)存混合方式。對(duì)于 mLSTM 而言，多頭和多單元是等價(jià)的。

將這些新的 LSTM 變體集成到殘差塊模塊中，就得到了 xLSTM 塊。將這些 xLSTM 塊剩余堆疊到架構(gòu)中，就形成了 xLSTM 架構(gòu)。xLSTM 架構(gòu)及其組件見(jiàn)圖 1。

xLSTM 塊應(yīng)在高維空間中對(duì)過(guò)去進(jìn)行非線性總結(jié)，以便更好地分離不同的歷史或上下文。分離歷史是正確預(yù)測(cè)下一個(gè)序列元素（如下一個(gè) token）的先決條件。研究者在此采用了 Cover 定理，該定理指出，在高維空間中，非線性嵌入模式比在原始空間中更有可能被線性分離。

他們考慮了兩種殘差塊結(jié)構(gòu)：(i) post up-projection 的殘差塊（如 Transformer），它非線性地概括了原始空間中的歷史，然后線性地映射到高維空間，應(yīng)用非線性激活函數(shù)，再線性地映射回原始空間（圖 3 左側(cè)和圖 1 第三欄，更詳細(xì)的版本見(jiàn)圖 9）。(ii) pre up-projection 的殘差塊（如狀態(tài)空間模型），它線性地映射到高維空間，在高維空間中非線性地總結(jié)歷史，然后線性地映射回原始空間。對(duì)于包含 sLSTM 的 xLSTM 塊，研究者主要使用了 post up-projection 塊。對(duì)于包含 mLSTM 的 xLSTM 塊，使用 pre up-projection 塊，因?yàn)樵诟呔S空間中內(nèi)存容量會(huì)變大。

實(shí)驗(yàn)

隨后，研究者對(duì) xLSTM 進(jìn)行了實(shí)驗(yàn)評(píng)估，并將其與現(xiàn)有的語(yǔ)言建模方法進(jìn)行了比較。

第 4.1 節(jié)討論了 xLSTM 在合成任務(wù)中的具體能力。首先，研究者測(cè)試了 xLSTM 的新指數(shù)門(mén)控與內(nèi)存混合在形式化語(yǔ)言上的有效性。然后，他們?cè)u(píng)估了 xLSTM 的新矩陣內(nèi)存在多次查詢聯(lián)想記憶任務(wù)（Multi-Query Associative Recall，MQAR）中的有效性。最后，研究者評(píng)估了 xLSTM 在 Long Range Arena（LRA）中處理長(zhǎng)序列的性能。

第 4.2 節(jié)比較了當(dāng)前各種語(yǔ)言建模方法的驗(yàn)證集復(fù)雜度，包括在同一數(shù)據(jù)集上對(duì) xLSTM 進(jìn)行消融研究，然后對(duì)不同方法的縮放行為進(jìn)行評(píng)估。

研究者在自回歸語(yǔ)言建模設(shè)置中使用 SlimPajama 的 15B token 訓(xùn)練了 xLSTM、Transformers、狀態(tài)空間模型（SSM）等模型。表 1 中的結(jié)果顯示，xLSTM 在驗(yàn)證復(fù)雜度方面優(yōu)于所有現(xiàn)有方法。

圖 6 顯示了該實(shí)驗(yàn)的擴(kuò)展結(jié)果，表明 xLSTM 對(duì)于更大規(guī)模的模型也有良好的表現(xiàn)。

消融研究則表明，性能改進(jìn)源于指數(shù)門(mén)控和矩陣內(nèi)存。

第 4.3 節(jié)進(jìn)行了更深入的語(yǔ)言建模實(shí)驗(yàn)。

研究者增加了訓(xùn)練數(shù)據(jù)量，對(duì)來(lái)自 SlimPajama 的 300B 個(gè) token 進(jìn)行了訓(xùn)練，并比較了 xLSTM、RWKV-4、Llama 和 Mamba。他們訓(xùn)練了不同大小的模型（125M、350M、760M 和 1.3B），進(jìn)行了深入的評(píng)估。首先，評(píng)估這些方法在推斷較長(zhǎng)語(yǔ)境時(shí)的表現(xiàn)；其次，通過(guò)驗(yàn)證易混度和下游任務(wù)的表現(xiàn)來(lái)測(cè)試這些方法；此外，在 PALOMA 語(yǔ)言基準(zhǔn)數(shù)據(jù)集的 571 個(gè)文本域上評(píng)估了這些方法；最后，評(píng)估了不同方法的擴(kuò)展行為，但使用的訓(xùn)練數(shù)據(jù)多了 20 倍。

可以看出，xLSTM 在性能和擴(kuò)展性上都更勝一籌。

更多研究細(xì)節(jié)，可參考原論文。