空腹なハングリーヒッポ：状態空間モデルを用いた言語モデルの構築へ

状態空間モデル（State Space Models, SSMs）は、特定のモダリティにおいて最先端の系列モデル性能を示しているが、言語モデルにおいては自己注意（attention）メカニズムに比べて性能が劣る。さらに、シーケンス長に対してほぼ線形にスケーリングする一方で、Transformerと比べてハードウェア利用効率が低いため、実際の処理速度は依然として遅い。本論文では、SSMsとattentionの間の表現力のギャップについての理解を深めるとともに、SSMsとattentionの間のハードウェア上の障壁を低減する点で進展を遂げた。まず、合成言語モデルタスクを用いてSSMsとattentionの性能差の原因を解明した。その結果、従来のSSMsは、シーケンス内の早期のトークンを記憶し直す能力、およびシーケンス内の異なるトークンを比較する能力に課題を抱えていることが明らかになった。これらの能力を明示的に強化するため、新しいSSMレイヤー「H3」を提案する。H3は合成言語においてattentionと同等の性能を達成し、OpenWebTextデータセットにおいてTransformerと0.4のPPL（Perplexity）差以内にまで迫った。さらに、2層のattention層を保持するハイブリッド125MパラメータのH3-attentionモデルは、驚くべきことにOpenWebTextにおいてTransformerよりも1.0 PPL優れた性能を示した。次に、現代のハードウェア上でSSMsの学習効率を向上させるために、「FlashConv」を提案する。FlashConvは、8Kまでのシーケンス長で効率を向上させる融合ブロックFFTアルゴリズムを採用し、SSMsの再帰的性質を活かした新規な状態伝搬アルゴリズムを導入することで、より長いシーケンスにもスケーラブルな処理を実現した。FlashConvは長距離タスクベンチマークで2倍の速度向上を達成し、ハイブリッド言語モデルがTransformerに比べてテキスト生成速度を2.4倍に加速した。FlashConvを用いて、Pileデータセット上でH3-attentionハイブリッドモデルを最大27億パラメータまでスケーリングしたところ、初期の結果として、Transformerより低いPPLを達成し、SuperGLUEベンチマークの多数のタスクにおいてゼロショットおよびフェイショット学習においてTransformerを上回る性能を示した。