LongMamba: Enhancing Mamba’s Long Context Capabilities via Training-Free Receptive Field Enlargement

要約

状態空間モデル（SSM）は、言語モデリングのためのトランスモデルの効率的な代替品として浮上しており、コンテキストの長さが増加するにつれて線形計算の複雑さと一定のメモリ使用量を提供します。
しかし、長いコンテキストの処理における効率にもかかわらず、最近の研究では、MAMBAモデルなどのSSMが一般に長いコンテキスト理解タスクのトランスと比較してパフォーマンスが低いことが示されています。
この重要な不足に対処し、効率的かつ正確な長いコンテキストの両方の理解を達成するために、Mambaモデルの長いコンテキスト機能を大幅に強化するトレーニングなしの手法であるLongmambaを提案します。
Longmambaは、マンバの隠されたチャネルを受容的なフィールドの長さに基づいてローカルおよびグローバルチャネルに分類できるという発見に基づいて構築されており、グローバルチャネルは主に長いコンテキスト機能を担当しています。
これらのグローバルチャネルは、入力コンテキストが長くなるにつれて重要なボトルネックになる可能性があります。
具体的には、入力の長さがトレーニングシーケンスの長さをほぼ上回る場合、グローバルチャネルは、適応的に受容フィールドを拡張することで制限を示し、マンバの長いコンテキストのパフォーマンスが低下します。
ロングマンバの重要なアイデアは、記憶に重要でないトークンの蓄積を防ぐことにより、これらのグローバルチャネルの隠れた状態記憶崩壊を軽減することです。
これは、最初にグローバルチャネルで重要なトークンを識別し、次にトークンフィルタリングを適用して、これらの重要なトークンのみを蓄積することによって達成されます。
Longmambaは、合成および現実世界の長いコンテキストシナリオ全体の広範なベンチマークを通じて、Mambaのロングコンテキストパフォーマンスの新しい基準を設定し、追加のトレーニングを必要とせずに運用範囲を大幅に拡張します。
私たちのコードは、https：//github.com/gatech-eic/longmambaで入手できます。

要約(オリジナル)

State space models (SSMs) have emerged as an efficient alternative to Transformer models for language modeling, offering linear computational complexity and constant memory usage as context length increases. However, despite their efficiency in handling long contexts, recent studies have shown that SSMs, such as Mamba models, generally underperform compared to Transformers in long-context understanding tasks. To address this significant shortfall and achieve both efficient and accurate long-context understanding, we propose LongMamba, a training-free technique that significantly enhances the long-context capabilities of Mamba models. LongMamba builds on our discovery that the hidden channels in Mamba can be categorized into local and global channels based on their receptive field lengths, with global channels primarily responsible for long-context capability. These global channels can become the key bottleneck as the input context lengthens. Specifically, when input lengths largely exceed the training sequence length, global channels exhibit limitations in adaptively extend their receptive fields, leading to Mamba’s poor long-context performance. The key idea of LongMamba is to mitigate the hidden state memory decay in these global channels by preventing the accumulation of unimportant tokens in their memory. This is achieved by first identifying critical tokens in the global channels and then applying token filtering to accumulate only those critical tokens. Through extensive benchmarking across synthetic and real-world long-context scenarios, LongMamba sets a new standard for Mamba’s long-context performance, significantly extending its operational range without requiring additional training. Our code is available at https://github.com/GATECH-EIC/LongMamba.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google