蒸留されたMambaモデルは、Pythia-1B Transformer教師モデルのパープレキシティ13.86にほぼ匹敵する14.01を達成しました。Apple Machine Learning Researchによって発表されたこの知見は、生成過程において優れた効率を提供する状態空間モデル（SSM）へと、事前学習済みTransformerの膨大な知識ベースを移行するための実行可能な経路を実証するものです。
アーキテクチャのギャップを埋める
Mambaのような状態空間モデル（SSM）は、生成時のメモリ消費を抑え、より高いスループットを提供するため、Transformerの代替として人気を集めています。しかし、機械学習コミュニティはすでにTransformerアーキテクチャのトレーニングと洗練に多大な投資を行っており、膨大な事前学習済み知識のレポジトリが構築されています。
SSMを採用する際の大きな障害は、「素朴な」蒸留、つまりTransformerからMambaモデルへ直接知識を転送する手法が、教師モデルの性能を維持できないことが多いことです。この問題を解決するための過去の試みは、純粋なMambaアーキテクチャに依存するのではなく、通常、AttentionとSSMブロックの両方を組み合わせたハイブリッドモデルを必要としていました。
「全体として、蒸留されたMambaモデルは元のPythia-1B Transformerの性能の大部分を維持しており、教師モデルの13.86に近い14.01のパープレキシティに達しています。」
— Attention to Mamba: A Recipe for Cross-Architecture Distillation
2段階蒸留レシピ
これを解決するために、研究者たちは原理に基づいた2段階の蒸留アプローチを提案しました。第1段階では、従来のTransformerから線形化されたバージョンのAttentionへと知識を蒸留します。このステップでは、「カーネルトリック」の適応を利用して、AttentionメカニズムをSSM構造により適合する形式に簡略化します。
第2段階では、この線形化されたAttentionモデルを、Attentionブロックを含まない適応されたMambaモデルへと蒸留します。この中間ステップを使用することで、研究者たちはMambaに「原理に基づいた初期化」を装備することができ、直接蒸留の際に失われた性能を回復させることが可能になりました。
性能と限界
結果は、蒸留されたMambaモデルがPythia-1B Transformerの性能の大部分を維持していることを示しています。パープレキシティのスコアは非常に近いものの、研究者らは、インコンテキスト・リコールに大きく依存する特定のタスクにおいて、いくつかの性能の差が現れることを指摘しました。
HN Digestによる図（machinelearning.apple.comの報告数値に基づく。数値はソースの表記通り）。
この研究では、異なるシーケンスミキサーアーキテクチャをテストするために、1Bスケールで10Bトークンを用いた徹底的なアブレーションを実施しました。また、プロセスを最適化するために、モデルサイズに関するスケーリング分析、および2つの蒸留段階の間でトークンをどのように割り当てるかに関する感度分析も行いました。