今天給大家分享一個超強的算法模型，xLSTM。

xLSTM（Extended Long Short-Term Memory）是對傳統(tǒng) LSTM（Long Short-Term Memory）模型的擴展和改進，旨在提升其在處理時間序列數(shù)據(jù)和序列預(yù)測任務(wù)中的性能。

傳統(tǒng)的 LSTM 模型通過引入遺忘門、輸入門和輸出門，解決了標準 RNN（Recurrent Neural Network）在長序列數(shù)據(jù)中存在的梯度消失和梯度爆炸問題。然而，LSTM 仍然存在一些局限性，尤其是在處理非常長的序列或需要更高維度特征提取時。xLSTM 的提出是為了克服這些限制。

圖片

xLSTM 中的核心增強功能

具體改進包括如下幾個方面。

1. 指數(shù)門控
   xLSTM 論文中引入的指數(shù)門控機制是對 LSTM 中使用的傳統(tǒng) S 形門控的重大改進。通過對輸入門和遺忘門采用指數(shù)激活，xLSTM 增強了模型在處理新信息時有效修改和更新其記憶的能力。
   在傳統(tǒng)的 LSTM 中，S 型門控函數(shù)限制了模型對記憶單元狀態(tài)進行重大改變的能力，尤其是當函數(shù)值接近 0 或 1 時。這種限制阻礙了 LSTM 快速適應(yīng)新數(shù)據(jù)的能力，并可能導(dǎo)致內(nèi)存更新不理想。
   xLSTM 通過將 S 型激活替換為指數(shù)激活來解決此問題。指數(shù)門控允許記憶細胞狀態(tài)發(fā)生更明顯的變化，使模型能夠快速整合新信息并相應(yīng)地調(diào)整其記憶。規(guī)范化器狀態(tài)有助于穩(wěn)定指數(shù)門控并保持輸入門和遺忘門之間的平衡。
2. 先進的內(nèi)存結(jié)構(gòu)
   xLSTM 采用了 sLSTM 和 mLSTM 等先進的內(nèi)存設(shè)計，它們分別使用了新的混合技術(shù)和基于矩陣的內(nèi)存。這些結(jié)構(gòu)允許更好的并行處理能力和更高效的數(shù)據(jù)存儲和檢索，這對于處理大規(guī)模數(shù)據(jù)至關(guān)重要。
3. 殘差連接
   通過在 LSTM 單元內(nèi)集成殘差塊，xLSTM 可以維持甚至增強整個網(wǎng)絡(luò)的梯度流，這對于有效訓(xùn)練更深層的模型至關(guān)重要。
4. 可并行架構(gòu)
   xLSTM 論文中最重要的進步之一是引入了可并行化的架構(gòu)，這解決了傳統(tǒng) LSTM 的一個主要限制。
   在傳統(tǒng)的 LSTM 中，標記的處理是按順序進行的，每個標記一次處理一個，這限制了模型利用并行性的能力，并導(dǎo)致訓(xùn)練和推理時間變慢。
   xLSTM 架構(gòu)引入了 mLSTM（矩陣內(nèi)存 LSTM）和 sLSTM（標量 LSTM）塊的靈活組合，從而實現(xiàn)了對 token 的并行處理。mLSTM 塊旨在同時對整個 token 序列進行操作，從而實現(xiàn)與 Transformer 模型實現(xiàn)的并行性類似的高效并行計算。
   mLSTM 模塊利用了前面討論過的矩陣內(nèi)存機制，使其能夠并行捕獲和處理所有 token 中豐富的高維信息。這種并行處理能力顯著加快了訓(xùn)練和推理過程，使得 xLSTM 比傳統(tǒng) LSTM 具有更高的計算效率。
   另一方面，sLSTM 塊保留了傳統(tǒng) LSTM 的順序處理特性，使模型能夠捕獲對特定任務(wù)可能很重要的某些順序依賴關(guān)系。
   在 xLSTM 架構(gòu)中，可以靈活地以不同的比例組合 mLSTM 和 sLSTM 塊，從而在并行性和順序建模之間取得平衡，從而能夠適應(yīng)各種語言建模任務(wù)。

關(guān)于 xLSTM 更詳細的說明，請參考對應(yīng)的論文：https://arxiv.org/pdf/2405.04517

LSTM 基礎(chǔ)知識

為了解釋 xLSTM，我們首先簡單回顧一下 LSTM。

原始 LSTM 的計算涉及幾個關(guān)鍵組件：輸入門、遺忘門、輸出門和單元狀態(tài)。

以下是 LSTM 單元的標準計算步驟。

圖片

xLSTM 解釋

之所以被稱為 xLSTM，是因為它將原始 LSTM 擴展為 LSTM 的多種變體，例如 sLSTM 和 mLSTM，每種變體都針對特定的性能和功能進行了優(yōu)化，以處理各種復(fù)雜的序列數(shù)據(jù)問題。

1.sLSTM

sLSTM 在 LSTM 的基礎(chǔ)上添加了標量更新機制。該設(shè)計通過對內(nèi)部存儲單元進行細粒度控制來優(yōu)化門控機制，使其更適合處理具有細微時間變化的序列。

sLSTM 通常利用指數(shù)門控和歸一化技術(shù)來提高處理長序列數(shù)據(jù)的穩(wěn)定性和準確性。通過這種方法，sLSTM 能夠提供與更復(fù)雜模型相當?shù)男阅?，同時保持較低的計算復(fù)雜度，使其特別適合資源受限的環(huán)境或需要快速響應(yīng)的應(yīng)用程序。

圖片

2.mLSTM

mLSTM（矩陣 LSTM）通過將原始 LSTM 中的向量運算擴展為矩陣運算，顯著增強了模型的內(nèi)存容量和并行處理能力。

在 mLSTM 中，每個狀態(tài)不再是單個向量，而是一個矩陣。這使得它能夠在單個時間步內(nèi)捕獲更復(fù)雜的數(shù)據(jù)關(guān)系和模式。mLSTM 特別適合處理大規(guī)模數(shù)據(jù)集或需要識別高度復(fù)雜數(shù)據(jù)模式的任務(wù)。

將矩陣引入狀態(tài)表示使得 mLSTM 能夠更自然、更有效地處理多維數(shù)據(jù)。通過以矩陣形式處理數(shù)據(jù)，mLSTM 可以同時處理多個數(shù)據(jù)點，從而提高吞吐量和學(xué)習(xí)過程的速度。這種能力在圖像和視頻處理等領(lǐng)域尤其有價值，因為這些領(lǐng)域的數(shù)據(jù)本質(zhì)上以矩陣形式存在。

圖片

效率與性能分析

xLSTM 論文對所提架構(gòu)的效率和性能進行了全面分析，突出了其相對于基于 Transformer 的模型的優(yōu)勢。作者進行了一系列實驗和比較，以證明 xLSTM 卓越的計算效率和建模能力。

xLSTM 的一個關(guān)鍵效率優(yōu)勢在于其時間和內(nèi)存復(fù)雜度。傳統(tǒng)的基于 Transformer 的模型相對于序列長度 N 表現(xiàn)出 O(N2) 的二次方時間和內(nèi)存復(fù)雜度。這意味著隨著序列長度的增加，Transformer 的計算成本和內(nèi)存需求呈二次方增長，使其處理長序列的效率較低。

相比之下，xLSTM 實現(xiàn)了線性時間復(fù)雜度 O(N) 和恒定內(nèi)存復(fù)雜度 O(1)（相對于序列長度）。這比 Transformers 有了顯著的改進，因為它允許 xLSTM 更有效地處理更長的序列，而無需二次方增加計算成本和內(nèi)存使用量。線性時間復(fù)雜度可以縮短訓(xùn)練和推理時間，而恒定內(nèi)存復(fù)雜度確保即使對于長序列，內(nèi)存需求仍然是可控的。

為了驗證效率和性能聲明，作者通過在包含 150 億個 token 的大規(guī)模數(shù)據(jù)集上訓(xùn)練多個模型進行了比較評估。評估中包括的模型是基于 Transformer 的語言模型 (LLM)、RWKV 模型和 xLSTM 的不同變體。

評估結(jié)果為 xLSTM 的卓越性能提供了有力證據(jù)。特別是，由一個 mLSTM 塊和零個 sLSTM 塊組成的 xLSTM[1:0] 變體在所有測試模型中實現(xiàn)了最低的困惑度。困惑度是語言建模中廣泛使用的指標，用于衡量模型預(yù)測序列中下一個標記的能力。困惑度越低，語言建模性能越好。

圖片

下面是 xLSTM 的開源實現(xiàn)：https://github.com/muditbhargava66/PyxLSTM