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

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

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

圖片

xLSTM 中的核心增強功能

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

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

關于 xLSTM 更詳細的說明，請參考對應的論文：https://arxiv.org/pdf/2405.04517

LSTM 基礎知識

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

原始 LSTM 的計算涉及幾個關鍵組件：輸入門、遺忘門、輸出門和單元狀態。

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

圖片

xLSTM 解釋

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

1.sLSTM

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

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

圖片

2.mLSTM

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

在 mLSTM 中，每個狀態不再是單個向量，而是一個矩陣。這使得它能夠在單個時間步內捕獲更復雜的數據關系和模式。mLSTM 特別適合處理大規模數據集或需要識別高度復雜數據模式的任務。

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

圖片

效率與性能分析

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

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

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

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

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

圖片

下面是 xLSTM 的開源實現：https://github.com/muditbhargava66/PyxLSTM