[論文レビュー] Hungry Hungry Hippos: Towards Language Modeling with State Space Models

Daniel Y. Fu, Tri Dao+4 more|arXiv (Cornell University)|Dec 28, 2022

Topic Modeling被引用数 116

ひとこと要約

この論文は、言語モデル用の表現力ギャップを埋める新しい状態空間モデル層 H3 と、長いシーケンス上でのSSMの効率的な訓練と推論を可能にするハードウェア認識FFTベースの加速 FlashConv を導入します。

ABSTRACT

State space models (SSMs) have demonstrated state-of-the-art sequence modeling performance in some modalities, but underperform attention in language modeling. Moreover, despite scaling nearly linearly in sequence length instead of quadratically, SSMs are still slower than Transformers due to poor hardware utilization. In this paper, we make progress on understanding the expressivity gap between SSMs and attention in language modeling, and on reducing the hardware barrier between SSMs and attention. First, we use synthetic language modeling tasks to understand the gap between SSMs and attention. We find that existing SSMs struggle with two capabilities: recalling earlier tokens in the sequence and comparing tokens across the sequence. To understand the impact on language modeling, we propose a new SSM layer, H3, that is explicitly designed for these abilities. H3 matches attention on the synthetic languages and comes within 0.4 PPL of Transformers on OpenWebText. Furthermore, a hybrid 125M-parameter H3-attention model that retains two attention layers surprisingly outperforms Transformers on OpenWebText by 1.0 PPL. Next, to improve the efficiency of training SSMs on modern hardware, we propose FlashConv. FlashConv uses a fused block FFT algorithm to improve efficiency on sequences up to 8K, and introduces a novel state passing algorithm that exploits the recurrent properties of SSMs to scale to longer sequences. FlashConv yields 2$ imes$ speedup on the long-range arena benchmark and allows hybrid language models to generate text 2.4$ imes$ faster than Transformers. Using FlashConv, we scale hybrid H3-attention language models up to 2.7B parameters on the Pile and find promising initial results, achieving lower perplexity than Transformers and outperforming Transformers in zero- and few-shot learning on a majority of tasks in the SuperGLUE benchmark.

研究の動機と目的

• 合成タスクを用いて、言語モデリングにおける状態空間モデル（SSMs）とアテンションの表現力ギャップを評価する。
• 過去のトークンを記憶し、シーケンス全体でトークンを比較できるSSM層（H3）を設計する。
• 標準ベンチマークでTransformersの性能に匹敵するかそれを上回るよう、SSMsのハードウェア効率を改善する。
• FlashConvを用いた大規模SSMベースモデルのスケーラブルで高速な訓練と生成を実証する。

提案手法

• シフト行列と対角行列を備えた二つのSSM層と、入力投影の乗法を導入し、トークンの記憶と跨トークンの比較を可能にするH3を導入する。
• H3が合成言語タスクでアテンションと一致し、OpenWebTextでTransformersとの差をほぼ縮める（0.4パープレキシティ以内）ことを示す。
• 2つのアテンション層を備えたハイブリッドH3–アテンションモデルを提案し、OpenWebTextでTransformersを1.0パープレキシティ上回る。
• ブロックFFTと状態伝播アルゴリズムを備えた統合FFTベース畳み込みのFlashConvを開発し、長いシーケンスとより長いモデルへSSMをスケールさせる。
• FlashConvは長いシーケンスで2×などのスピードアップを生み、テキスト生成をTransformersより2.4×高速にする。
• Pile上で最大2.7BパラメータまでハイブリッドH3–アテンションモデルをスケールさせ、パープレキシティとSuperGLUEの性能を報告する。
• 同程度のサイズでHybrid H3がTransformerより最大2.4×高速な推論スループットを示す比較を提供する。

実験結果

リサーチクエスチョン

• RQ1 SSMsは合成タスクにおいてアテンションと同等の表現力を示せるか。
• RQ2 どのようなアーキテクチャ変更（例：シフト/対角SSMと乗法的相互作用）がSSMsに過去のトークンを記憶させ、シーケンス全体で比較を可能にするか。
• RQ3 FlashConvのようなハードウェア認識アルゴリズムは現代のアクセラレータでSSMとアテンションのランタイムギャップを縮小できるか。
• RQ4 ハイブリッドH3–アテンションモデルは標準ベンチマークでTransformerと比較してパープレキシティとゼロ/フェewショット性能で競争力を持つか。
• RQ5 FlashConvはSSMベースのモデルを十億単位のパラメータへスケールさせ、訓練と推論の速度向上を現実的に達成できるか。

主な発見

• H3は合成言語タスクでアテンションと一致し、OpenWebTextでTransformersとの差をほぼ縮める（0.4パープレキシティ以内）。
• ハイブリッドH3–アテンションモデル（2つのアテンション層）はOpenWebTextでTransformersを1.0パープレキシティ上回る。
• Pile上で125M–2.7Bのサイズで訓練されたハイブリッドH3–アテンションモデルは、パープレキシティでTransformerのベースラインを上回り、強力なSuperGLUEゼロ/数-shot性能を示す。
• FlashConvは長いシーケンスで最大2×のスピードアップを実現し、SSMベースの生成をTransformersより2.4×高速にする。
• Pile上で2.7BパラメータへスケールしたSSMベースモデルは、いくつかの設定でTransformersより低いパープレキシティを達成し、ゼロ/数-shotの結果で競争力を示す。
• FlashConvはLong Range ArenaでTransformersベースラインおよび他のSSM手法と比較して最先端のスピードを達成する。