DeciMamba: Exploring the Length Extrapolation Potential of Mamba

要約

長距離シーケンスの処理は、入力長の 2 次複雑さのため、Transformer にとって大きな課題となります。
有望な代替案は Mamba です。Mamba は高いパフォーマンスを示し、必要な計算リソースを大幅に削減しながら Transformer レベルの機能を実現します。
この論文では、Mamba の長さ一般化機能について調査しますが、この機能は比較的制限されていることがわかりました。
一連の視覚化と分析を通じて、この制限は、トレーニング中に使用されるシーケンスの長さによって決まる制限された有効受容野に起因することがわかりました。
この制約に対処するために、Mamba 用に特別に設計されたコンテキスト拡張メソッドである DeciMamba を導入します。
このメカニズムは、S6 レイヤー内に埋め込まれた非表示のフィルタリング メカニズムの上に構築されており、追加のトレーニングがなくても、トレーニングされたモデルが適切に外挿できるようになります。
現実世界の長距離 NLP タスクに関する実証実験では、DeciMamba がトレーニング中に見られたものよりも 25 倍長いコンテキスト長を推定でき、追加の計算リソースを利用することなくそれを行うことができます。
コードとモデルを公開します。

要約(オリジナル)

Long-range sequence processing poses a significant challenge for Transformers due to their quadratic complexity in input length. A promising alternative is Mamba, which demonstrates high performance and achieves Transformer-level capabilities while requiring substantially fewer computational resources. In this paper we explore the length-generalization capabilities of Mamba, which we find to be relatively limited. Through a series of visualizations and analyses we identify that the limitations arise from a restricted effective receptive field, dictated by the sequence length used during training. To address this constraint, we introduce DeciMamba, a context-extension method specifically designed for Mamba. This mechanism, built on top of a hidden filtering mechanism embedded within the S6 layer, enables the trained model to extrapolate well even without additional training. Empirical experiments over real-world long-range NLP tasks show that DeciMamba can extrapolate to context lengths that are 25x times longer than the ones seen during training, and does so without utilizing additional computational resources. We will release our code and models.

arxiv情報

著者 Assaf Ben-Kish,Itamar Zimerman,Shady Abu-Hussein,Nadav Cohen,Amir Globerson,Lior Wolf,Raja Giryes
発行日 2024-06-20 17:40:18+00:00
arxivサイト arxiv_id(pdf)

提供元, 利用サービス

arxiv.jp, Google

カテゴリー: cs.AI, cs.LG パーマリンク