Mambaはこのように使用できます。バイトをトークンに分割せずに効率的に学習します。

言語モデルを定義するときは通常、基本的な単語分割方法を使用して文を単語、サブワード、または文字に分割します。その中で、サブワード分割は、トレーニング効率と語彙外の単語を処理する能力との間で自然なトレードオフを実現するため、最も人気のある選択肢となっています。しかし、いくつかの研究では、タイプミス、スペルや大文字小文字の差異、形態の変化に対する堅牢性の欠如など、サブワード分割の問題が指摘されています。

そのため、一部の研究者は異なるアプローチを採用し、バイトシーケンスを使用する方法を採用しました。これは、間に単語の分割を介さずに、生データから予測へのエンドツーエンドのマッピングです。サブワード モデルと比較すると、バイト レベルの言語モデルは、さまざまな書き方や形態のバリエーションにわたってより簡単に一般化できます。もちろん、テキストをバイトとしてモデル化すると、生成されるシーケンスは対応するサブワードよりもはるかに長くなります。この場合、効率性の向上はアーキテクチャの改善を通じて達成する必要があります。

自己回帰トランスフォーマーは言語モデリングで主流ですが、効率性の問題があります。計算コストは​​シーケンスの長さの 2 乗で増加するため、長い (バイト) シーケンスには適応しにくくなります。研究者たちは、たとえばトークンのグループが中間層内でマージされる長さを考慮したモデリング手法を開発することにより、Transformer の内部表現を圧縮して長いシーケンスの処理を容易にしました。最近、Yuら[2023]は、固定サイズのバイトフラグメントをサブワードの模擬圧縮形式として使用するMegaByte Transformerを提案した。したがって、MegaByte は計算コストを削減できます。ただし、これは最善のアプローチではない可能性があります。

新しい論文では、コーネル大学の研究者が、効率的でシンプルなバイトレベルの言語モデルである MambaByte を紹介しています。このモデルは、最近導入された Mamba アーキテクチャを直接採用したものです。 Mamba は、状態空間モデル (SSM) によって開拓されたアプローチに基づいて構築され、テキストなどの離散データに対してより効率的な選択メカニズムを導入し、効率的な GPU 実装を提供します。私たちの単純な観察は、Mamba を（変更せずに）使用すると、言語モデリングにおける主な計算ボトルネックが緩和され、パッチ適用が不要になり、利用可能な計算リソースが効率的に使用されるということです。

• 論文タイトル: MambaByte: トークンフリー選択状態空間モデル
• 論文リンク: https://arxiv.org/pdf/2401.13660.pdf

実験では、MambaByte を Transformers、SSM、MegaByte (パッチ適用) アーキテクチャと比較しました。これらはすべて、複数の長いテキスト データセットで固定パラメータと固定計算設定の下で比較されました。図 1 に主な調査結果をまとめます。

バイトレベルのトランスフォーマーと比較すると、MambaByte はより高速に優れたパフォーマンスを実現し、計算効率も大幅に向上します。著者らはまた、既存の最先端のサブワードモデルと比較したトークンフリー言語モデルの実現可能性についても検討しています。この点に関して、MambaByte はさまざまなサブワード ベースライン モデルと競合しますが、より長いシーケンスも処理できることがわかりました。私たちの結果は、MambaByte が既存のトークナイザー依存モデルの強力な代替手段であり、エンドツーエンドの学習を促進するために使用できることを示唆しています。

背景: 選択的状態空間シーケンスモデル

SSM は、一次微分方程式を使用して、時間の経過に伴う隠れ状態の進化をモデル化します。線形時間不変 SSM は、いくつかのモダリティの深層学習において有望な結果を示しています。しかし、Mamba の著者 Gu 氏と Dao 氏は最近、これらの方法の一定のダイナミクスには、言語モデリングなどのタスクに必要となる可能性のある、隠れた状態における入力依存のコンテキスト選択が欠けていると主張しました。この目的のために、彼らはMambaを提案しました。これは、入力x(t)∈R、隠れ状態h(t)∈R^n、出力y(t)∈Rが与えられた場合の時刻tにおける時間変化する連続状態ダイナミクスを次のように定義します。

そのパラメータは、対角時間不変システム行列 A∈R^(n×n) と、時間変動入力および出力行列 B(t)∈R^(n×1) および C(t)∈R^(1×n) です。

バイトのような離散時系列をモデル化するには、(1)の連続時間ダイナミクスを離散化によって近似する必要がある。これにより、離散時間の隠れた再帰が起こり、各時間ステップで新しい行列A、B、Cが生成されます。

(2)はリカレントニューラルネットワークの線形バージョンに似ており、言語モデル生成中にこのリカレント形式で適用できることに注意する。離散化では、各入力位置にx[k] = x(t_k)に対応する時間ステップ∆[k]が必要です。離散時間行列A、B、Cは∆[k]から計算できます。図 2 は、Mamba が離散シーケンスをモデル化する方法を示しています。

Mambaでは、SSM項は入力選択的であり、つまり、B、C、および∆は入力x[k]∈R^dの関数として定義されます。

ここで、W_B ∈ R^(n×d) (C の定義も同様)、W_∆ ∈ R^(d×r)、W_R ∈ R^(r×d) (ある r ≪d の場合) は学習可能な重みであり、softplus は正値性を保証します。各入力次元 d について、SSM パラメータ A、B、C は同じですが、時間ステップの数 ∆ は異なることに注意してください。これにより、各時間ステップ k でサイズ n × d の隠し状態が生成されます。

Mamba はこの SSM レイヤーを完全なニューラル ネットワーク言語モデルに埋め込みます。具体的には、このモデルは、以前のゲート SSM にヒントを得た一連のゲート レイヤーを採用しています。図 3 は、SSM レイヤーとゲート付きニューラル ネットワークを組み合わせた Mamba アーキテクチャを示しています。

線形再帰による並列スキャン。トレーニング中、作成者はシーケンス x 全体にアクセスできるため、線形再帰の計算をより効率的に行うことができます。スミスら[2023]は、効率的な並列スキャンを使用して線形SSMの順次再帰を効率的に計算できることを実証しました。 Mamba の場合、著者らはまず e_k = の L 組のシーケンスに再帰をマッピングし、次にとなるような関連付け演算子を定義します。最後に、並列スキャンを適用してシーケンスを計算します。一般に、これはL/2 プロセッサを使用すると時間がかかります。ここで、は行列乗算のコストです。 A は対角行列であり、線形再帰はO(nL) の時間と空間で並列に計算できることに注意してください。対角行列を使用した並列スキャンも非常に効率的に実行され、O (nL) FLOP しかかかりません。

実験結果

表 2 は、各データセットのバイトあたりのビット数 (BPB) を示しています。この実験では、MegaByte758M+262M モデルと MambaByte モデルは、バイトあたり同じ数の FLOP を使用します (表 1 を参照)。著者らは、MambaByte がすべてのデータセットで一貫して MegaByte を上回っていることを発見しました。さらに、著者らは、資金の制約により、MambaByte を 80B 全体でトレーニングすることはできなかったが、それでも MambaByte は計算とトレーニング データを 63% 削減しながらも MegaByte よりも優れたパフォーマンスを発揮したと指摘しています。さらに、MambaByte-353M は Byte-Level Transformer や PerceiverAR よりも優れたパフォーマンスを発揮します。

MambaByte は、非常に少ないトレーニング ステップで、はるかに大規模なモデルよりも優れたパフォーマンスを発揮するのはなぜでしょうか?図 1 では、同じ数のパラメータを持つモデルを調べることで、この関係をさらに詳しく調べています。この図は、同じパラメータ サイズの MegaByte モデルの場合、入力パッチが少ないモデルの方がパフォーマンスが優れていることを示していますが、計算正規化後はパフォーマンスは同様になります。実際、フルレングスの Transformer は、絶対的には遅いものの、計算正規化後は MegaByte と同様のパフォーマンスを発揮します。対照的に、Mamba アーキテクチャに切り替えると、計算使用率とモデルのパフォーマンスが大幅に向上します。

これらの調査結果に基づいて、表 3 では PG19 データセットでこれらのモデルのより大きなバージョンを比較しています。この実験では、著者らはMambaByte-972MをMegaByte-1.3B+350Mおよびその他のバイトレベルモデル、およびいくつかのSOTAサブワードモデルと比較しました。研究者らは、MambaByte-972M がすべてのバイトレベル モデルよりも優れており、わずか 150B バイトでトレーニングした場合でもサブワード モデルと競合できることを発見しました。

テキスト生成。 Transformer モデルでの自己回帰推論ではコンテキスト全体をキャッシュする必要があり、生成速度に大きく影響します。 MambaByte にはこのボトルネックがありません。時間の経過とともに変化するレイヤーごとに 1 つの隠し状態のみを保持するため、各ステップの生成時間は一定です。表 4 は、A100 80GB PCIe GPU 上の MambaByte-972M と MambaByte-1.6B のテキスト生成速度を MegaByte-1.3B+350M と比較したものです。 MegaByte はパッチ適用によって生成コストを大幅に削減しますが、ループ生成を使用しているため、MambaByte は同様のパラメータ設定で前者よりも 2.6 倍高速であることが確認されました。