Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture

要約

機械学習モデルは、より長いコンテキストとより優れたパフォーマンスを実現するために、シーケンスの長さとモデルの次元の両方でますますスケーリングされています。
ただし、Transformers などの既存のアーキテクチャは、これらの両方の軸に沿って二次的にスケールします。
私たちは、シーケンスの長さとモデルの次元に沿って二次二次的にスケールできるパフォーマンスの高いアーキテクチャはあるのか、と尋ねます。
シーケンス長とモデル次元の両方に沿って同じ二次二次プリミティブを使用する新しいアーキテクチャである Monarch Mixer (M2) を紹介します。Monarch 行列は、多くの線形変換をキャプチャする表現力豊かな構造化行列の単純なクラスであり、GPU で高いハードウェア効率を実現します。
そして二次二次的にスケールします。
概念実証として、非因果的な BERT スタイルの言語モデリング、ViT スタイルの画像分類、および因果的な GPT スタイルの言語モデリングの 3 つのドメインで M2 のパフォーマンスを調査します。
非因果的 BERT スタイルのモデリングの場合、M2 は、最大 27% 少ないパラメータでダウンストリーム GLUE 品質において BERT ベースおよび BERT-large に匹敵し、シーケンス長 4K で最大 9.1$\times$ 高いスループットを達成します。
ImageNet では、M2 はパラメータが半分でありながら、精度において ViT-b よりも 1% 優れています。
因果 GPT スタイルのモデルには技術的な課題があります。マスキングによって因果関係を強制すると、二次関数のボトルネックが発生します。
このボトルネックを軽減するために、多変量多項式の評価と内挿に基づいたモナーク行列の新しい理論的ビューを開発します。これにより、準二次関数を維持しながら M2 を因果関係になるようにパラメータ化できます。
このパラメータ化を使用して、M2 は、The PILE の事前トレーニングのパープレキシティで 3 億 6000 万のパラメーターで GPT スタイルの Transformers と一致します。これは、注意や MLP なしで Transformer の品質を一致させることが可能である可能性があることを初めて示しています。

要約(オリジナル)

Machine learning models are increasingly being scaled in both sequence length and model dimension to reach longer contexts and better performance. However, existing architectures such as Transformers scale quadratically along both these axes. We ask: are there performant architectures that can scale sub-quadratically along sequence length and model dimension? We introduce Monarch Mixer (M2), a new architecture that uses the same sub-quadratic primitive along both sequence length and model dimension: Monarch matrices, a simple class of expressive structured matrices that captures many linear transforms, achieves high hardware efficiency on GPUs, and scales sub-quadratically. As a proof of concept, we explore the performance of M2 in three domains: non-causal BERT-style language modeling, ViT-style image classification, and causal GPT-style language modeling. For non-causal BERT-style modeling, M2 matches BERT-base and BERT-large in downstream GLUE quality with up to 27% fewer parameters, and achieves up to 9.1$\times$ higher throughput at sequence length 4K. On ImageNet, M2 outperforms ViT-b by 1% in accuracy, with only half the parameters. Causal GPT-style models introduce a technical challenge: enforcing causality via masking introduces a quadratic bottleneck. To alleviate this bottleneck, we develop a novel theoretical view of Monarch matrices based on multivariate polynomial evaluation and interpolation, which lets us parameterize M2 to be causal while remaining sub-quadratic. Using this parameterization, M2 matches GPT-style Transformers at 360M parameters in pretraining perplexity on The PILE–showing for the first time that it may be possible to match Transformer quality without attention or MLPs.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google