原作者帶隊再次改造xLSTM，7B模型速度最快超Mamba 50%，權重程式碼全開源
机器之心發表於2025-03-19
近年來，大型語言模型（LLM）透過大量計算資源在推理階段取得了解決複雜問題的突破。推理速度已成為 LLM 架構的關鍵屬性，市場對高效快速的 LLM 需求不斷增長。其中，採用 Transformer 架構的模型雖然佔據了主流，但在輸入序列長度增加時，計算量會呈二次方增長。因此，自上個世紀 90 年代興起的 LSTM 捲土重來，它的提出者和奠基者 Sepp Hochreiter 在去年 5 月推出了 xLSTM，將 LSTM 擴充套件到數十億引數，成為 Transformer 的有力替代品，提供了與序列長度線性相關的計算擴充套件和穩定的記憶體佔用。然而，xLSTM 在擴充套件至更大引數規模時存在限制，推理速度和效率具體如何也沒做系統測評。近日，Sepp Hochreiter 等來自 NXAI、JKU 的研究者再次對 xLSTM 進行了最佳化，現在可以擴充套件到 70 億引數了。具體來講，xLSTM 7B 模型基於 DCLM 資料集，使用 128 塊 H100 GPU，在 8192 上下文長度下訓練了 2.3 萬億 token。研究者對原始 xLSTM 架構進行了改進，確保訓練效率和穩定性，同時保持任務效能。新架構依靠 mLSTM 單元和並行訓練模式，實現高效能的同時最大化速度。- 論文標題：xLSTM 7B: A Recurrent LLM for Fast and Efficient Inference
- 論文地址：https://arxiv.org/pdf/2503.13427
- 程式碼地址：https://github.com/NX-AI/xlstm
- Hugging Face 地址：https://huggingface.co/NX-AI/xLSTM-7b
透過修改模組架構，研究者最佳化了吞吐量，在低維空間執行 mLSTM 並新增前饋 MLP 層，同時去除了不必要的元件以提高 GPU 利用率。最佳化後的架構在保持相似效能的同時，將 token 吞吐量提高了 2 到 4 倍。研究者還最佳化了訓練穩定性，特別是 mLSTM 單元的門控機制，有效解決了梯度問題。在各類任務評估中，xLSTM 7B 與同規模 Transformer 和 Mamba 模型表現相當。透過架構最佳化，該模型在推理效率測試中實現了最高的預填充和生成吞吐量，同時保持最低的 GPU 記憶體佔用。論文作者之一 Günter Klambauer 表示，xLSTM 7B 成為了最快、最高效的 7B 語言模型！xLSTM 7B 架構的核心是 mLSTM 單元，它的迴圈和並行模式可以實現高效的訓練和推理。為了充分發揮該單元的潛力，研究者重新審視了相鄰塊結構的設計。與 Mamba 等其他線性 RNN 類似，以前的 xLSTM 架構將與通道卷積相結合的 mLSTM 單元置於線性上投影和下投影之間，這被稱為預上投影（pre up-projection ）塊。這些塊將序列混合和通道混合結合在一個塊中，因此均勻堆疊，而無需交錯位於前饋 MLP 層。儘管預上投影塊架構已展示出了對 1.4B 引數 xLSTM 的競爭性語言建模效能，但由於以下幾方面的原因，它在計算效率方面付出了很大代價：- 在預上投影塊中，mLSTM 在比模型嵌入維數高得多的維數上執行，這導致 mLSTM 操作的計算成本和 GPU 記憶體使用量大幅增加。
- 省略位置前饋 MLP 層會導致模型中高效線性層 FLOP 的比例下降。
- 以前的 xLSTM 架構使用幾個額外的元件，例如可學習的殘差連線、通道卷積以及用於計算查詢、鍵和值的小（塊對角化）投影層。如果沒有自定義核心融合，這些小操作會導致 GPU 上出現多個短核心呼叫，無法有效利用張量核心，從而大幅降低 GPU 利用率。
- 以前，輸入和遺忘門預啟用是透過連線的查詢、鍵和值投影計算出來的。而在大規模張量並行訓練設定中，這需要每個 mLSTM 塊進行額外的全歸約操作，從而增加總體通訊成本。
因此，為了將 xLSTM 擴充套件到更大的模型大小，研究者透過解決以上四個限制來最佳化 mLSTM 塊以實現最大效率。對於最佳化 mLSTM 塊，研究者首先在模型的嵌入維數而不是更高維數的空間中操作 mLSTM 單元，並在每個 mLSTM 層之後放置位置前饋 MLP 層。此修改增加了高度最佳化的線性層（即矩陣乘法）FLOP 的比例，並降低了 mLSTM 操作的計算成本。顯著減少的 GPU 記憶體使用量使得在訓練期間可以使用更大的批大小，從而提高了訓練效率。此外，研究者放棄了通道卷積和可學習的殘差連線等操作，並用密集線性層替換塊查詢、鍵和值投影。這再次增加了線性層 FLOP，並確保有效使用 mLSTM 層內的張量核。最後，確保每個 head 的門預啟用都是獨立計算的。 這些最佳化產生了下圖 1 和下圖 8 中改進後的 mLSTM 塊和 xLSTM 架構，其中在 xLSTM 7B 架構中堆疊了 32 個 mLSTM 塊。下表 4 為 xLSTM 7B 的超引數，包括模型引數（近 70 億）、詞表大小（50257）、塊數量（32）、模型維數（4096）以及 head 數（8）。研究者觀察到，本文最佳化在 1.4B 引數的模型訓練中實現了 3.5 倍的加速，但在驗證困惑度方面略有損失，可以透過增加幾個訓練步驟來緩解，詳見下表 2。研究者發現，先前在 7B 引數規模下的 xLSTM 架構在訓練初期階段常出現不穩定現象。具體而言，他們觀察到在較高學習率條件下訓練會導致梯度幅度和損失值劇烈波動。本文透過以下方法解決了這些穩定性問題：1. 使用 RMSNorm 的預歸一化（Pre-Norm with RMSNorm）下圖 9 中的實驗證實，預歸一化技術同樣適用於 xLSTM 架構的預歸一化層。因此，研究者在 xLSTM 架構中將 LayerNorm 替換為 RMSNorm（全稱為 Root Mean Square Normalization）。2. 門控軟上限限制（Gate Soft-Capping）為了降低潛在的大幅異常特徵和相關損失峰值，研究者對輸入門和遺忘門的預啟用值
應用了軟上限限制，使其值被限制在特定上限值 a 的 - a 與 a 之間。本文采用 a=15 對門控進行限制，所使用的函式為3. 負輸入門偏置初始化（Negative Input Gate Bias Initialization）研究者發現，在訓練初期，xLSTM 模型會出現較大的梯度範數峰值，這對模型的最終效能產生不利影響（詳見下圖 11）。將輸入門初始化為較大的負值（如 - 10）能有效緩解這些梯度範數峰值，從而提升模型效能。綜上所述，這些最佳化措施使 xLSTM 7B 的預訓練過程變得極為穩定，如下圖 2 所示。研究者首先在 7B 引數規模上，將 xLSTM 7B 與最先進的 Transformer 和迴圈神經網路（RNN）大語言模型進行了基準測試。結果總結在下表 1 中，顯示 xLSTM 7B 在 7B 規模模型中排名居中，其中一些表現更好的模型受益於更大規模的訓練資料集。研究者認為，如果使用更大且更精心策劃的訓練資料集，尤其是在早期訓練階段更加註重數學和程式碼資料，xLSTM 7B 可能會達到最強 7B 模型的效能水平。研究者將 xLSTM 與幾種基線模型進行了比較：作為 Transformer 基線的 Llama 2 7B（未進行長文字微調）和 Llama 3.1 8B（已進行長達 131K 詞元的長文字微調），作為狀態空間模型（State Space Model，SSM）基線的 CodestralMamba 和 FalconMamba，以及作為額外迴圈神經網路（Recurrent Neural Network，RNN）基線的 RWKV-5/6。下表 3 展示了 RULER 評估結果。對於 xLSTM 7B，預訓練中的長文字降溫（cooling）階段極大地提升了其長文字處理能力，使其效能與狀態空間模型相當，並且優於 RWKV-5/6。值得注意的是，長文字 xLSTM 7B 在 131K 上下文長度時實現了 20% 的平均準確率，儘管在降溫階段訓練時僅使用了最多 32K 的上下文長度。這一點尤為顯著，因為與具有不斷增長的 KV 快取（Key-Value cache）的 Transformer 不同，xLSTM 7B 必須在有限容量的固定大小記憶體中儲存整個序列的資訊（見表 3）。本研究主要關注本地單使用者推理場景，這在模型部署到邊緣裝置時較為常見。除非另有說明，研究在單個英偉達 H100 GPU 上對批大小為 1 的 xLSTM 7B 模型進行生成式推理基準測試，並將其與 Llama 2 和 Llama 3 模型進行了比較。如下圖 4 所示，由於注意力機制隨輸入上下文長度呈二次方增長，Transformer 模型在較長預填充長度下的文字生成速度顯著降低。研究表明，xLSTM 7B 的文字生成速度比 Mamba 快約 50%，這主要得益於其最佳化的塊設計。即使在預填充長度為 0 的情況下，xLSTM 7B 也比採用類似塊設計的基於 Llama 的 Transformer 模型更快。研究者測量了不同生成長度下的 token 生成時間和 GPU 記憶體使用情況（不包括預填充）。圖 5（左）展示了迴圈模型在計算時間上呈線性增長，與 Transformer 呈二次方增長的對比；圖 5（右）則顯示了迴圈模型記憶體佔用保持恆定，而 Transformer 的 KV 快取隨生成長度線性增長的對比。得益於最佳化的模組設計，mLSTM 在低維空間中執行，使得 xLSTM 7B 模型與 Mamba 模型相比具有顯著更低的記憶體佔用（如下圖 5 右側所示）和更短的生成時間（如圖 5 左側所示）。TTFT（Time To First Token）在語言模型作為使用者介面（可能在邊緣裝置上）的應用場景中，較短的響應時間至關重要。下圖 6 展示了不同模型在處理各種長度的預填充（prefill）內容後，生成 1 個或 100 個 token 所需的響應時間或延遲。在所有預填充長度條件下，xLSTM 7B 模型均表現出最快的響應速度。 
