Random Long-Context Access for Mamba via Hardware-aligned Hierarchical Sparse Attention

要約

変圧器よりも再発性ニューラルネットワーク（RNN）の重要な利点は、線形計算と空間の複雑さにより、長いシーケンスに対するより速いトレーニングと推論を可能にすることです。
ただし、RNNは根本的に歴史的コンテキストにランダムにアクセスすることができず、注意メカニズムを単に統合するだけで効率の利点を損なう可能性があります。
この制限を克服するために、\ textbf {h} ierarchical \ textbf {s} parse \ textbf {a} ttention（hsa）を提案します。
HSAは、入力をチャンクに分割し、最高$ k $チャンクを選択し、情報を階層的に集約します。
コアイノベーションは、各チャンク内の細粒トークンレベルの情報に基づいて、トークンとチャンクまでの関連性を学習することにあります。
このアプローチは、ドメイン内とドメイン外のコンテキストの長さの両方で、チャンク選択の精度を高めます。
HSAを効率的にするために、ハードウェアに並べられたカーネル設計をさらに紹介します。
HSAとMAMBAを組み合わせることで、RAMBAを導入します。これは、4Kの長さのコンテキストのみでトレーニング前にトレーニングを行うにもかかわらず、6400万のコンテキストにわたってPassKey検索で完全な精度を達成し、ほぼ一定のメモリフットプリントを備えたさまざまなダウンストリームタスクの大幅な改善を実現します。
これらの結果は、ランバのロングコンテキストモデリングにおける大きな可能性を示しています。

要約(オリジナル)

A key advantage of Recurrent Neural Networks (RNNs) over Transformers is their linear computational and space complexity enables faster training and inference for long sequences. However, RNNs are fundamentally unable to randomly access historical context, and simply integrating attention mechanisms may undermine their efficiency advantages. To overcome this limitation, we propose \textbf{H}ierarchical \textbf{S}parse \textbf{A}ttention (HSA), a novel attention mechanism that enhances RNNs with long-range random access flexibility while preserving their merits in efficiency and length generalization. HSA divides inputs into chunks, selecting the top-$k$ chunks and hierarchically aggregates information. The core innovation lies in learning token-to-chunk relevance based on fine-grained token-level information inside each chunk. This approach enhances the precision of chunk selection across both in-domain and out-of-domain context lengths. To make HSA efficient, we further introduce a hardware-aligned kernel design. By combining HSA with Mamba, we introduce RAMba, which achieves perfect accuracy in passkey retrieval across 64 million contexts despite pre-training on only 4K-length contexts, and significant improvements on various downstream tasks, with nearly constant memory footprint. These results show RAMba’s huge potential in long-context modeling.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google