タイパン：選択的注意を備えた効率的かつ表現豊かな状態空間言語モデル

要旨

自然言語処理（NLP）において、効率的な長いコンテキストの言語モデリングは依然として重要な課題です。Transformerは言語タスクで優れていますが、トレーニング中の二次計算量と推論時の線形スケーリングメモリコストにより、長いシーケンスに苦労しています。最近のState Space Models（SSMs）としてMambaなどがあり、定数のメモリ使用量を提供しますが、コンテキスト内の広範な検索が必要なタスクでは性能が低いです。私たちは、Mamba-2とSelective Attention Layers（SALs）を組み合わせた新しいハイブリッドアーキテクチャ、Taipanを提案します。これらのSALsは、長距離の相互作用が必要なトークンを特定し、重要でない特徴を削除し、その後アテンションモジュールを使用して表現を拡張します。このアプローチは、Mambaの効率性とTransformerのようなメモリ集約型タスクでのパフォーマンスをバランスさせます。アテンション予算を制約することで、Taipanは計算効率を保ちつつ、最大100万トークンまでのコンテキスト長で正確な予測を拡張します。私たちの実験は、さまざまなスケールやタスクでTaipanの優れたパフォーマンスを示し、効率的な長いコンテキストの言語モデリングに向けた有望な解決策を提供しています。

English

Efficient long-context language modeling remains a significant challenge in Natural Language Processing (NLP). While Transformers dominate language tasks, they struggle with long sequences due to quadratic computational complexity in training and linearly scaling memory costs during inference. Recent State Space Models (SSMs) such as Mamba offer alternatives with constant memory usage, but they underperform in tasks requiring extensive in-context retrieval. We introduce Taipan, a novel hybrid architecture that combines Mamba-2 with Selective Attention Layers (SALs). These SALs identify tokens requiring long-range interactions, remove less important features, and then augment their representations using the attention module. This approach balances Mamba's efficiency with Transformer-like performance in memory-intensive tasks. By constraining the attention budget, Taipan extends accurate predictions to context lengths of up to 1 million tokens while preserving computational efficiency. Our experiments demonstrate Taipan's superior performance across various scales and tasks, offering a promising solution for efficient long-context language modeling.

タイパン：選択的注意を備えた効率的かつ表現豊かな状態空間言語モデル

Taipan: Efficient and Expressive State Space Language Models with Selective Attention

要旨

Summary

Support

Support