Mambaへの注目：クロスアーキテクチャ蒸留のためのレシピ

要旨: Mambaのような状態空間モデル（SSM）は、Attentionベースの対応モデルと比べて生成時のメモリ消費が少なく、スループットが高いため、Transformerモデルの人気の代替として広く受け入れられるようになってきました。一方で、コミュニティはTransformerの学習方法に関する膨大な知見を積み上げており、多くの事前学習済みTransformerモデルがすぐに利用可能です。既存の事前学習済みTransformerを活用しつつSSMの導入を促進するために、AttentionベースのモデルをMambaのようなアーキテクチャへ蒸留するための効果的なレシピを特定することを目指します。しかし、異なるアーキテクチャ間の蒸留に関する先行研究では、TransformerからMambaへの単純な（na\"ive）蒸留手順では、元の教師モデルの性能を保持できないことが示されており、この制限は、AttentionブロックとSSMブロックを組み合わせたハイブリッドな解決策によってしばしば克服されています。私たちの研究の主要な主張は、Mambaに原理に基づく初期化を施すことで、アーキテクチャ間蒸留の全体としてより良いレシピを回復できるという点です。そこで本研究では、原理に基づく二段階アプローチを提案します。まず、カーネル・トリックの適応を用いて、伝統的なTransformerからAttentionの線形化版へ知識を蒸留します。次に、その線形化版を、いかなるAttentionブロックも使用しない適応されたMambaモデルへ蒸留します。全体として、蒸留されたMambaモデルは下流タスクにおいて元のPythia-1B Transformerの性能を保持でき、パープレキシティは教師の13.86に近い14.11を維持します。提案レシピの有効性を示すために、10Bトークンでシーケンスミキサのアーキテクチャを変化させた1Bスケールでの徹底的なアブレーションを行い、モデルサイズと総蒸留トークンに関するスケーリング分析を実施し、さらに各段階間でのトークン配分に関する感度分析も行います。