ハイエナが次世代トランスフォーマーになる？ StripedHyena-7B オープンソース: 最大 128k の入力、トレーニング速度が 50% 向上

近年発表されたAIモデル、例えば言語、視覚、音声、生物学など各分野の大規模モデルは、Transformerアーキテクチャと切り離せないものです。しかし、そのコアモジュールである「アテンションメカニズム」の計算量は、「入力シーケンスの長さ」の2乗で増大します。この機能により、Transformerの長いシーケンスへの応用は大幅に制限されます。たとえば、本全体を一度に処理したり、ギガピクセルレベルの画像を処理したりすることは不可能です。

GPT-4 のような強力なアルゴリズムであっても、この欠陥から逃れることはできません。

最近、Together Research は、「長いコンテキスト」用の新しいアーキテクチャを採用した新しい言語モデル、StripedHyena をオープンソース化しました。最大 128k トークンの長いコンテキストを処理でき、トレーニングと推論における Transformer アーキテクチャのパフォーマンスを向上させ、現在の主流のアーキテクチャに代わるソリューションを提供します。

オープンソースリンク: https://github.com/togethercomputer/stripedhyena

StripedHyena は、ショートコンテキストとロングコンテキストの両方の評価において、同じモデルサイズで最高のオープンソース Transformer モデルと同等のパフォーマンスを達成した最初のモデルでもあります。OpenLLM ベンチマーク タスクでは Llama-2、Yi、Mistral 7B と同等のパフォーマンスを達成し、ロングコンテキストの要約ではより優れたパフォーマンスを発揮します。

StripedHyena は、Hyena ブロック内に配置されたマルチヘッド、グループ化されたクエリ アテンション、ゲート畳み込みで構成されるハイブリッド アーキテクチャです。これは、畳み込みを状態空間モデル (SSM) または一定のメモリ デコードによる切り捨てフィルターとして Hyena ブロック内に表現する点で、従来のデコーダーのみの Transformer とは異なります。

実験結果によると、StripedHyena は、32k トークン、64k トークン、128k トークン シーケンスのエンドツーエンドのトレーニングにおいて、従来のトランスフォーマーよりも 30%、50%、100% 以上高速です。

SH7B のもう 1 つの利点は、自己回帰生成中のメモリ フットプリントが Transformers と比較して 50% 以上削減されることです。Transformers では、各レイヤーのキーと値が事前入力フェーズ中にキャッシュされるため、再計算を回避し、増分デコードを高速化できます。

ハイエナブロック

低ランクおよびスパース近似に基づく既存のサブ二次手法は、Transformer に一致するように密な注意層と組み合わせる必要があり、つまり両者の間には表現力のギャップがあることになります。

つまり、注意メカニズムは言語処理においてその二次パワーのほんの一部しか活用していないため、研究上の疑問は、大規模にトレーニングした場合に注意メカニズムのパフォーマンスに匹敵できる二次以下の演算子が存在するかどうかです。

今年 2 月、スタンフォード大学とモントリオール大学 (MILA と Université de Montréal) の研究者が、準二次降下法の注意の代替手段である Hyena を提案しました。数千から数十万のトークンのシーケンスの想起と推論のタスクにおいて、Hyena は状態空間やその他の暗黙的および明示的な方法に依存する演算子よりも 50 ポイント以上精度を向上させ、注意ベースのモデルと一致します。

論文リンク: https://arxiv.org/abs/2302.10866

研究者らは、標準データセット (WikiText 103 および The Pile) での言語モデリング用に新しい高密度アテンションフリー アーキテクチャを構築し、シーケンス長 2k で必要なトレーニング計算を 20% 削減しながら Transformer 品質を達成しました。シーケンス長 8k では、Hyena 演算子は高度に最適化されたアテンションの 2 倍の速度で、シーケンス長 64k では 100 倍の速度です。

研究者らは、要素ごとの乗算（ゲーティング）や長い畳み込み（つまり、フィルタのサイズが入力と同じ長さになる畳み込み）などの効率的な二次プリミティブを組み合わせ、最終的に実験結果から肯定的な答えを得ました。

メカニズムの解釈可能性（想起や帰納法など）に関する最近の研究を踏まえて、私たちは一連の対象を絞った推論タスクを策定し、注意のパフォーマンスと既存の二次方程式法との品質ギャップに相関する注意の 3 つの特性を抽出します。

1. データ管理

アテンション メカニズムは、単一のブロックに線形関数のファミリ全体をエンコードする、表現力豊かなデータ制御線形演算子を実装します。

2. サブ線形パラメータスケーリング

注意層のパラメータの数をシーケンスの長さから切り離すと、Transformer はフィードフォワード ニューラル ネットワーク (FFN) などの注意層間の他の場所にさらに多くのパラメータを分散できるようになります。

3. 制限のないコンテキスト

特定の入力に対して、アテンションは無制限のコンテキストを持ち、自己回帰モデルなどのマスクを使用する場合を除き、局所性などの任意の制限なしに、任意の 2 つの入力間の依存関係を近似できます。

ハイエナの階層

上記の 3 つの発見に基づいて、研究者らは、2 つの効率的な二次プリミティブ再帰的に定義された演算子 (長い畳み込みと要素ごとの乗算ゲーティング) で構成される Hyena 階層を提案しました。

再帰の指定された深さ (つまり、ステップ数) によって演算子のサイズが制御されます。短い再帰の場合、既存のモデルを特別なケースとして使用できます。

研究者らは、ハイエナ再帰の各ステップを対応する行列形式にマッピングすることで、ハイエナ演算子は、データ制御行列、つまりエントリが入力関数である行列の分解として同等に定義できることを発見しました。

さらに、研究者らは、高速畳み込みアルゴリズムを使用して、完全な行列を具体化せずにハイエナ演算子を効率的に評価する方法を示しました。

経験的に、Hyena 演算子は大規模な注意によって品質ギャップを大幅に埋めることができ、混合注意を必要とせずに、より少ない計算コストで同様の困惑度と下流のパフォーマンスを達成できます。

能力ギャップを埋める

Hyena は、標準的な高密度アテンションとサブ二次演算子の間に品質のギャップがあり、このギャップは大規模言語モデリングのパフォーマンスに関連する推論タスクを通じて識別できるという考えに基づいて設計されました。

研究者らは、機械的な解釈可能性のベンチマークの基本セット（誘導と想起）を拡張し、タスクの複雑さが増すと（語彙のサイズが大きくなるなど）、モデルのパフォーマンスがどれだけ急速に低下するかを調べるために追加のタスクを追加しました。

さらに、Hyena における長い畳み込みの最適なパラメータ化についても研究されています。

数十万の項を含む最も困難な設定では、暗黙的なパラメータ化スキームにより、状態空間、周波数領域パラメータ化、または標準畳み込みを使用する他の演算子よりも 50% 以上精度が向上します。

言語と視覚の拡張

研究者らはまた、推論ベンチマークスイートでのランキングが大規模な品質を予測することを検証し、10億未満のパラメータ規模でHyena自己​​回帰言語モデリングをテストし、標準データセット（WikiText103およびThe Pile）の非高密度アテンションアーキテクチャに新しいsotaを実装して、Transformerに匹敵する品質を達成しました。

3 億 3,500 万のパラメータを持つ The Pile データセットでは、システムは浮動小数点演算 (FLOP) を 20% 削減しながら、Transformer に匹敵するパープレキシティ メトリックを達成しました。

拡張として、研究者らは、大規模な画像認識と Visual Transformer (ViT) における注意の代替として、Hyena 演算子の一般化可能性をテストしました。

画像分類では、Hyena は ImageNet-1k データセットで最初からトレーニングした場合、Attention に匹敵する精度を実現します。

より長い文脈

研究者らは、長いシーケンスにおけるハイエナの効率性もベンチマークした。長さ 8192 のシーケンスでは、測定された速度は Dense Self-Attention の 5 倍、高度に最適化された FlashAttention の 2 倍ですが、長さ 64k のシーケンスでは、測定された速度は FlashAttention の 100 倍です。