必要なのは注意だけです。

少なくともマトリックスの場合はそうです。

Mamba アーキテクチャの最新の進歩: 新しいモデルは、わずか 1% の計算コストで最先端のパフォーマンスを実現します。

これはTransformerのおかげで可能になりました。

Transformer モデルから Mamba などの代替アーキテクチャに知識を効果的に移行することにより、モデルは計算コストを低く抑えながらパフォーマンスを向上させることができます。

これは、Mamba の開発者の 1 人である Albert Gu 氏が主導した最新の成果です。

この方法は、Mamba 以外の Transformer 以外のアーキテクチャにも適用されることに注意してください。

TransformerからSSMへ

Transformer は、二次的な自己注意メカニズムに依存しているため、大量の計算を必要とします。

二次自己注意メカニズムにより、モデルはシーケンシャルデータを処理する際に、シーケンス内の長距離依存関係を効果的に捕捉できます。しかし、二次時間計算量（入力サイズが2倍になると、モデルに必要な計算時間は4倍に増加する）により、長いシーケンスの処理にかかる計算コストは​​非常に高くなります。

この問題に対処するために、学術界は、微調整と推論のコストが低い Mamba や RWKV などの多くの新しいアーキテクチャを提案してきました。

Transformer モデルの事前トレーニングにはすでに大量の計算リソースが必要であることを考えると、研究者たちはなぜその基盤の上に構築できないのか疑問に思いました。

そこで本研究では、Transformerの事前学習済みモデルを使用してSSMモデルを学習するMOHAWKと呼ばれる蒸留手法を提案した。

その中核には、線形アテンション、Mambaの構造化マスクアテンション（SMA）などのアテンションメカニズムがあり、これらはすべて入力長次元にわたるシーケンス変換です。そのため、それぞれにソフトマックスなどの独自の行列ミキサーが備わっています。

アテンションと SSM を、異なるクラスのマトリックスを適用して異なるトークン埋め込みをブレンドするシーケンス変換と見なすと、シーケンス モデル アーキテクチャは、独立したシーケンス ブレンディング ブロックとチャネル ブレンディング ブロックに分解できます。

たとえば、Transformer はアテンション (シーケンス ミキサー) ブロックと MLP (チャネル ミキサー) ブロックで構成されており、この分解によってモデルの各要素を抽出できます。

蒸留プロセスは3つの段階に分かれています。

フェーズ1：マトリックスの方向付け。シーケンス変換マトリックス自体の位置合わせを行います。

第2段階：隠れ状態のアライメント。事前に学習した表現を犠牲にすることなく、ネットワークの各層の隠れ状態の表現をアライメントします。

第3段階：重みの転送と知識の蒸留。エンドツーエンドのトレーニングフェーズを通じて重みが転送され、最終的に、トレーニングデータのごく一部のみを使用してネットワークの最終出力が蒸留されます。

この方法は、Phi-Mamba などのモデルを実際に変更するために使用できます。

Mamba-2とPhi-1.5を組み合わせたものです。

MOHAWK メソッドを使用すると、モデルは事前トレーニング済みの Transformer モデルから学習し、状態空間モデルとして、長いシーケンスを処理する際に従来の Transformer アーキテクチャよりも効率的になります。

このモデルは蒸留に 30 億トークンのみを使用し、ゼロからトレーニングされたモデルの1% のデータ量しかありませんが、そのパフォーマンスはオープンソースの非 Transformer アーキテクチャの最先端 (SOTA) に達しています。

実験により、隠れた状態をより適切に整列させることで、後続のステージのパフォーマンスが向上することが示されました。

研究チームは、50億トークン蒸留を採用したハイブリッドモデル「Phi-Mamba-1.5B」も公開しました。このモデルは類似のハイブリッドモデルと同等の性能を発揮しますが、注意層は4層のみを使用します。

この蒸留方法は Mamba にのみ適用できるわけではないことに注意する必要があります。

この研究は、CUM の助教授であり、Cartesia AI の共同設立者兼主任科学者でもある Albert Gu 氏が主導しました。

昨年、彼と FlashAttention の作者 Tri Dao は Mamba を提案しました。Mamba は、Transformer のパフォーマンスに真に匹敵する最初の線形時系列モデルとなりました。

論文の宛先:
https://arxiv.org/abs/2408.10189