[論文レビュー] Laughing Hyena Distillery: Extracting Compact Recurrences From Convolutions

Stefano Massaroli, Michael Poli+12 more|VU Research Portal|Oct 28, 2023

Advanced Neural Network Applications被引用数 4

ひとこと要約

この論文では、事前学習済みの長距離畳み込みシーケンスモデル（LCSM）から、自己回帰的生成における1トークンあたりO(1)の計算量とメモリ使用量を実現する、コンパクトで再帰的な状態空間モデル（SSM）を抽出するdistillation手法「Laughing Hyena」を紹介する。この手法により、1.3BパラメータでTransformerの10倍のスループットを達成し、Hyenaの1.5倍のスループットを達成するが、distillation後の品質に劣化は見られない。

ABSTRACT

Recent advances in attention-free sequence models rely on convolutions as alternatives to the attention operator at the core of Transformers. In particular, long convolution sequence models have achieved state-of-the-art performance in many domains, but incur a significant cost during auto-regressive inference workloads -- naively requiring a full pass (or caching of activations) over the input sequence for each generated token -- similarly to attention-based models. In this paper, we seek to enable $\mathcal O(1)$ compute and memory cost per token in any pre-trained long convolution architecture to reduce memory footprint and increase throughput during generation. Concretely, our methods consist in extracting low-dimensional linear state-space models from each convolution layer, building upon rational interpolation and model-order reduction techniques. We further introduce architectural improvements to convolution-based layers such as Hyena: by weight-tying the filters across channels into heads, we achieve higher pre-training quality and reduce the number of filters to be distilled. The resulting model achieves 10x higher throughput than Transformers and 1.5x higher than Hyena at 1.3B parameters, without any loss in quality after distillation.

研究の動機と目的

• 標準的な推論におけるO(K)のメモリとO(K²)の計算量というボトル neck を克服し、長距離畳み込みシーケンスモデル（LCSM）において、1トークンあたり定数時間・定数メモリの自己回帰的生成を可能にすること。
• モデル品質を保持しつつ、事前学習済みの畳み込み層から低次元で安定した状態空間モデル（SSM）を抽出するdistillationフレームワークの開発。
• チャネル間で重みを共有するフィルタの再設計により、事前学習の品質とdistillationの効率を向上させること。
• メモリフットプリントを削減し、LCSMにおける再帰的推論を可能にすることで、大バッチでの高スループット生成を実現すること。

提案手法

• 事前学習済みLCSMの畳み込みフィルタから、コンパクトなSSMを抽出するために、有理的補間とモデル次数低減を適用する。
• バーリセントリックおよびPronyに類似した手法をインspireした、因子分解されたモーダルパラメータ化を導入し、SSMの安定性を向上させ、数値的問題を回避する。
• 近似の目的関数として畳み込みフィルタの不一致指標を用い、さまざまな下流タスクとの互換性を保証する。
• SSMの最適な状態次元dを決定するために、ハンケル作用素のスペクトル解析を用いる。
• 有効なフィルタ次元を向上させるとともに、distillationの複雑さを低減するために、チャネル間で重みを共有するようにHyenaブロックを再設計する。
• 誤差境界をグラミアン固有値から導出することで、バランストレンケーションとモーダルトレンケーションを用いてSSMモデル低減を実施する。

実験結果

リサーチクエスチョン

• RQ1 事前学習済みの長距離畳み込みモデルから、1トークンあたりO(1)の推論コストを実現するコンパクトで再帰的なSSMを抽出できるか？
• RQ2 SSMのdistillationにおいて、近似誤差とモデル効率の最適なトレードオフをもたらす状態次元dは何か？
• RQ3 SSMのパラメータ化をどのように改善すれば、distillation中の数値的不安定性を回避し、収束性を向上させられるか？
• RQ4 Hyenaブロックのアーキテクチャ的変更は、事前学習の品質向上とdistillationコストの低減に寄与するか？
• RQ5 LCSMからSSMへのdistillationは、下流タスクのパフォーマンスを保持しつつ、高スループットの生成を可能にするか？

主な発見

• Laughing Hyenaは、1.3Bパラメータで、同等のTransformerの10倍のピークスループットを達成し、Hyenaの1.5倍のスループットを達成するが、distillation後の品質に劣化は見られない。
• 1.3Bパラメータの状況で、同じメモリ制約下において、Transformerに比べて3倍少ないメモリで512トークンを生成できる。
• Kトークンの生成において、メモリがO(d)、時間計算量がO(dK)に保たれるのに対し、kvキャッシュ付きのTransformerはそれぞれO(K)とO(K²)である。
• Hyenaブロックにおけるチャネル間の重み共有は、事前学習のパープレクサリティを向上させるとともに、distillation対象のフィルタ数を削減する。
• バランストレンケーションとモーダルトレンケーションの両手法は、一部の層において誤差低減が非単調となる傾向を示しており、数値的安定性に敏感であることが示唆された。
• 近似の目的関数としてフィルタの不一致指標を用いることで、多様な下流タスクにわたるロバスト性が保証された。