[論文レビュー] xLSTM: Extended Long Short-Term Memory

Maximilian Beck, Korbinian Pöppel+7 more|arXiv (Cornell University)|May 7, 2024

Topic Modeling被引用数 86

ひとこと要約

xLSTMは指数ゲーティングと2つのメモリ拡張変種（スカラー記憶を持つsLSTMと行列記憶を持つmLSTM）を導入し、LSTMの制限を克服することで、TransformerおよびState Space Modelsと比較して言語モデリングの性能とスケーリングに有利な成果を達成する。

ABSTRACT

In the 1990s, the constant error carousel and gating were introduced as the central ideas of the Long Short-Term Memory (LSTM). Since then, LSTMs have stood the test of time and contributed to numerous deep learning success stories, in particular they constituted the first Large Language Models (LLMs). However, the advent of the Transformer technology with parallelizable self-attention at its core marked the dawn of a new era, outpacing LSTMs at scale. We now raise a simple question: How far do we get in language modeling when scaling LSTMs to billions of parameters, leveraging the latest techniques from modern LLMs, but mitigating known limitations of LSTMs? Firstly, we introduce exponential gating with appropriate normalization and stabilization techniques. Secondly, we modify the LSTM memory structure, obtaining: (i) sLSTM with a scalar memory, a scalar update, and new memory mixing, (ii) mLSTM that is fully parallelizable with a matrix memory and a covariance update rule. Integrating these LSTM extensions into residual block backbones yields xLSTM blocks that are then residually stacked into xLSTM architectures. Exponential gating and modified memory structures boost xLSTM capabilities to perform favorably when compared to state-of-the-art Transformers and State Space Models, both in performance and scaling.

研究の動機と目的

• 大規模な言語モデリングにおいて、従来のLSTMの制限（ストレージの再設計、メモリ容量、並列性の欠如）を克服する必要性を動機付ける。
• 指数ゲーティングを備えた2つのメモリ拡張LSTM変種（sLSTMとmLSTM）を提案する。
• これらの変種を残差アーキテクチャに統合したxLSTMブロックを導入し、スケーラブルなモデリングを実現する。
• 合成タスクと大規模言語モデリング実験を通じて、特定の設定でxLSTMがTransformer/SSMのベースラインと競合する、またはそれを上回ることを示す。

提案手法

• 正規化と安定化を伴う指数ゲーティングを導入し、再設計可能なストレージ決定を可能にする。
• スカラー記憶を持つsLSTMを、スカラー更新とセル/ヘッド間の記憶の混合とともに開発する。
• 行列記憶と共分散更新ルールを備えたmLSTMを開発し、高容量ストレージと並列取得を可能にする。
• sLSTMまたはmLSTMを残差ブロックのバックボーン内に埋め込むことでxLSTMブロックを作成する（変種に応じてアッププロジェクションの前後を選択）。
• プレLayerNormバックボーンを用いてxLSTMブロックをxLSTMアーキテクチャに積み上げ、スケーラブルなモデリングを実現する。
• CUDA最適化実装を提供し、Transformerと比較したメモリ・計算特性について論じる。

実験結果

リサーチクエスチョン

• RQ1 指数ゲーティングと新しいメモリ構造は、スケール時に基本的なLSTMの制限（ストレージの再設計、メモリ容量、および逐次ボトルネック）を解消できるか？
• RQ2 sLSTMとmLSTMは、数十億パラメータへスケールさせた言語モデリングタスクで、TransformerおよびState Space Modelsとどのように比較されるか？
• RQ3 残差スタッキングを用いたxLSTMアーキテクチャは、モデルサイズとデータ規模を跨いで、競争力のあるパープレキシティと下流タスク性能を提供するか？

主な発見

• 指数ゲーティングと行列/スカラー記憶設計により、検証パープレキシティでxLSTMが複数のベースラインモデルを上回る。
• 合成タスクおよび長距離タスクで、xLSTMは従来のLSTMおよび一部のTransformer/SSM系と比べて状態追跡とメモリ容量の向上を示す。
• xLSTM[1:0]およびxLSTM[7:1は、SlimPajama実験でモデルサイズを問わず強い検証パープレキシティを達成し、好ましいスケーリング挙動を示す。
• 300Bトークンのトレーニングでは、xLSTM変種はシーケンス長外挙、長い文脈パープレキシティ、下流タスク性能でRWKV、Llama、Mambaと比較して競争力を維持する。
• アブレーション研究は、指数ゲーティングと行列記憶が性能向上の主要因であることを示し、学習可能で入力依存のゲートが追加の利点を提供する。
• xLSTMアーキテクチャは、線形に近いメモリ考慮と並列可能な成分を備え、Transformer中心の領域に対するスケーラブルな代替手段となり得る。