[論文レビュー] MambaByte: Token-free Selective State Space Model

Junxiong Wang, Tushaar Gangavarapu+2 more|arXiv (Cornell University)|Jan 24, 2024

Cellular Automata and Applications被引用数 7

ひとこと要約

MambaByte は選択的状態空間モデルに基づくトークン不要のバイトレベル言語モデルで、サブワードTransformersと競合する性能を発揮し、生成はより速く、スケーラブルです。

ABSTRACT

Token-free language models learn directly from raw bytes and remove the inductive bias of subword tokenization. Operating on bytes, however, results in significantly longer sequences. In this setting, standard autoregressive Transformers scale poorly as the effective memory required grows with sequence length. The recent development of the Mamba state space model (SSM) offers an appealing alternative approach with a fixed-sized memory state and efficient decoding. We propose MambaByte, a token-free adaptation of the Mamba SSM trained autoregressively on byte sequences. In terms of modeling, we show MambaByte to be competitive with, and even to outperform, state-of-the-art subword Transformers on language modeling tasks while maintaining the benefits of token-free language models, such as robustness to noise. In terms of efficiency, we develop an adaptation of speculative decoding with tokenized drafting and byte-level verification. This results in a $2.6 imes$ inference speedup to the standard MambaByte implementation, showing similar decoding efficiency as the subword Mamba. These findings establish the viability of SSMs in enabling token-free language modeling.

研究の動機と目的

• トークン化バイアスとロバスト性の問題を緩和するため、トークンフリーの言語モデリングを動機づける。
• パッチを当てずに長いシーケンスへ効率的にスケールするバイトレベルモデルを開発する。
• Mambaの選択的状態空間アーキテクチャを活用し、線形時間でのシーケンス処理を可能にする。
• Transformers、MegaByte、SSMベースラインと比較した計算量・パラメータ数を揃えた性能を評価する。
• バイト単位の自己回帰設定における生成効率の利点を示す。

提案手法

• Mambaの選択的状態空間モデルを、直接バイト列上で動作するよう適用する。
• 入力に依存するB、C、Deltaを、連続時間SSMダイナミクスの離散化として、入力の関数として用いる。
• 訓練時に線形再帰を効率的に計算するために並列スキャンを利用する。
• 言語モデリングのためにSSM層をゲート付きネットワークアーキテクチャに組み込む。
• 長文データセット（PG19、Stories、Books、ArXiv、Code）で、計算量・パラメータ数を揃えた実験を実施する。
• バイトレベルのTransformer、SSM、MegaByte、サブワードベースラインと比較する。

Figure 1 : Benchmarking byte-level models with a fixed parameter budget. Language modeling results on PG19 ( $8,192$ consecutive bytes), comparing the standard Transformer (Vaswani et al., 2017 ; Su et al., 2021 ) , MegaByte Transformer (Yu et al., 2023 ) , gated diagonalized S4 (Mehta et al., 2023

実験結果

リサーチクエスチョン

• RQ1 トークンフリーなバイトレベルモデルは、長文言語モデリングにおいてトークン化されたサブワード変換器と同等かそれ以上を達成できるか？
• RQ2 選択的状態空間モデルはバイト列に対して計算量と生成速度の利点を提供するか？
• RQ3 既存のバイトレベルおよびサブワードアーキテクチャと計算量・パラメータ数を揃えた比較で、MambaByteはどのように性能を示すか？

主な発見

バイトレベルモデル	コンテキスト	訓練バイト数	テスト BPB	PG19	Stories	Books	ArXiv	Code
Transformer-320M	1024	80 B	1.057	1.057	1.064	1.097	0.816	0.575
PerceiverAR-248M	8192	80 B	1.104	1.104	1.070	1.104	0.791	0.546
MegaByte-758M+262M (patch: 8)	8192	80 B	1.000	1.000	0.978	1.007	0.678	0.411
MambaByte-353M	8192	30 B*	0.930	0.930	0.908	0.966	0.663	0.396

• MambaByte は同じ FLOPs/バイト予算の下で、複数の長文データセットで MegaByte を一貫して上回る。
• 小型のMambaByte（353M）は平均的に BPB でバイトレベルの Transformer および PerceiverAR を上回る。
• 972M パラメータのとき、 substantial bytes で訓練した場合、PG19 で最先端サブワードモデルと競合するまたは上回る困惑度を達成。
• MambaByteは再帰的な状態空間ベースの設計によりMegaByteより生成が速い。
• 計算量を揃えた設定で、MambaByte はより少ない計算量とデータで Transformer に近い損失を達成する。
• トークンフリーのバイトモデルは、長文文脈タスクにおいてトークナイザ依存モデルの実用的な代替となり得る。

Figure 2 : Illustration of the Mamba SSM. (a) The discrete-time input $x[k]$ , along with input-selective $\Delta[k]$ . (b) The continuous-time signal $x(t)$ . (c) Mathematically, the SSM transforms the continuous-time $x(t)$ through an $n$ -dimensional hidden state (here, $n=4$ ) using parameters $