Attention to Mamba: クロスアーキテクチャ蒸留のためのレシピ
Attention to Mamba: A Recipe for Cross-Architecture Distillation (machinelearning.apple.com)
要約
State Space Models (SSMs)であるMambaは、Transformerの代替として注目されています。本研究では、TransformerモデルからMambaライクなアーキテクチャへの知識蒸留を、プリンシプルに基づいた初期化手法を用いることで効率化する2段階のアプローチを提案しています。この手法により、元のTransformerモデルの性能の大部分を維持したMambaモデルの構築が可能になり、特に大規模モデルでの応用が期待されます。
全文翻訳
Attention to Mamba: クロスアーキテクチャ蒸留のためのレシピ
著者: Abhinav Moudgil†
**, Ningyuan Huang‡
**, Eeshan Gunesh Dhekane, Pau Rodríguez Lopez, Luca Zappella, Federico Danieli
State Space Models (SSMs) であるMambaは、メモリ消費量の削減と生成時のスループット向上により、Transformerモデルの代替として人気を集めています。一方で、Transformerのトレーニングに関するコミュニティの知識は豊富であり、多くの事前学習済みTransformerモデルが利用可能です。SSMの採用を促進しつつ、既存の事前学習済みTransformerを活用するために、AttentionベースのモデルからMambaライクなアーキテクチャへの知識を蒸留するための効果的なレシピを特定することを目指します。しかし、これまでのクロスアーキテクチャ蒸留の研究では、TransformerからMambaへの単純な蒸留手順では元の教師モデルの性能を維持できないことが示されており、この限界はAttentionとSSMブロックを組み合わせたハイブリッドソリューションで克服されることがよくありました。本研究の鍵となる主張は、Mambaにプリンシプルに基づいた初期化を施すことで、クロスアーキテクチャ蒸留のための全体的により良いレシピを回復できるということです。この目的のため、プリンシプルに基づいた2段階のアプローチを提案します。まず、カーネルトリックの適応を用いて、従来のTransformerからAttentionの線形化バージョンへ知識を蒸留します。次に、Attentionブロックを一切使用しない適応型Mambaモデルへ、線形化されたバージョンを蒸留します。全体として、蒸留されたMambaモデルは、元のPythia-1B Transformerの性能の大部分を維持し、教師モデルの13.86に近い14.01のパープレキシティを達成しました。これは下流タスクにも大部分引き継がれますが、インコンテキストリコールにより依存するタスクではいくつかのギャップが見られます。本研究のレシピの有効性を示すため、1Bスケールで10Bトークンを使用し、シーケンスミキサーアーキテクチャの変更、モデルサイズと総蒸留トークンに関するスケーリング分析、およびステージ間のトークン割り当てに関する感度分析を実施しました。
- † Concordia University and MILA Research Institute
- ‡ Johns Hopkins University and Flatiron Institute
** Work done while at Apple
関連文献とアップデート。
State-Space Models (SSMs)、特にMambaは最近、Transformerの有望な代替として登場しました。MambaはSSMレイヤー（S6）にインプットセレクティビティを導入し、ブロック定義に畳み込みとゲーティングを組み込んでいます。これらの変更はMambaのSSM前身モデルに対する性能を向上させますが、Mambaがインプットセレクティビティによって提供される追加機能をどのように活用しているかは、依然として不明瞭な点が多く、また…
我々は、計算予算と、学生モデルと教師モデルへのその割り当てに基づいて、蒸留モデルの性能を推定する蒸留スケーリング則を提案します。我々の発見は、学生モデルの性能を最大化するために教師モデルと学生モデルの両方に対して計算効率の良い割り当てを可能にすることで、大規模蒸留に伴うリスクを軽減します。我々は、教師モデルが既に存在する場合と、…