Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models

要約

Transformer アーキテクチャは、言語モデリングなどの分野で主流のパラダイムとなっていますが、二次時間の自己注意のため、多くの推論設定で問題が発生します。
Mamba など、最近提案された二次二次アーキテクチャは有望であることが示されていますが、最も強力な Transformer モデルよりも大幅に少ない計算リソースで事前トレーニングされています。
この研究では、事前トレーニングされた Transformer アーキテクチャを状態空間モデル (SSM) などの代替アーキテクチャに抽出できる方法を紹介します。
私たちのアプローチの重要なアイデアは、トランスフォーマーと SSM の両方を、トークン シーケンスに対してさまざまな形式の混合行列を適用しているとみなすことができるということです。
したがって、SSM でさまざまな粒度を一致させることによって、Transformer アーキテクチャを段階的に抽出できます。最初に混合行列自体を一致させ、次に各ブロックの隠れユニットを一致させ、最後にエンドツーエンドの予測を一致させます。
MOHAWK と呼ばれる私たちのメソッドは、3B トークンのみを使用して Phi-1.5 アーキテクチャに基づく Mamba-2 バリアント (Phi-Mamba) を抽出することができ、5B トークンを使用してハイブリッド バージョン (Hybrid Phi-Mamba) を抽出することができます。
通常、モデルをゼロからトレーニングするために使用されるトレーニング データの 1% 未満しか使用していないにもかかわらず、Phi-Mamba は、過去のすべてのオープンソースの非 Transformer モデルと比較して、大幅に優れたパフォーマンスを誇ります。
MOHAWK を使用すると、SSM などのモデルが、Transformer ベースのアーキテクチャのトレーニングに投資された計算リソースを活用できるようになり、そのようなモデルを構築するための新しい手段が強調されます。

要約(オリジナル)

Transformer architectures have become a dominant paradigm for domains like language modeling but suffer in many inference settings due to their quadratic-time self-attention. Recently proposed subquadratic architectures, such as Mamba, have shown promise, but have been pretrained with substantially less computational resources than the strongest Transformer models. In this work, we present a method that is able to distill a pretrained Transformer architecture into alternative architectures such as state space models (SSMs). The key idea to our approach is that we can view both Transformers and SSMs as applying different forms of mixing matrices over the token sequences. We can thus progressively distill the Transformer architecture by matching different degrees of granularity in the SSM: first matching the mixing matrices themselves, then the hidden units at each block, and finally the end-to-end predictions. Our method, called MOHAWK, is able to distill a Mamba-2 variant based on the Phi-1.5 architecture (Phi-Mamba) using only 3B tokens and a hybrid version (Hybrid Phi-Mamba) using 5B tokens. Despite using less than 1% of the training data typically used to train models from scratch, Phi-Mamba boasts substantially stronger performance compared to all past open-source non-Transformer models. MOHAWK allows models like SSMs to leverage computational resources invested in training Transformer-based architectures, highlighting a new avenue for building such models.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google