[論文レビュー] KDEformer: Accelerating Transformers via Kernel Density Estimation

Amir Zandieh, In‐Su Han+2 more|arXiv (Cornell University)|Feb 5, 2023

Advanced Neural Network Applications被引用数 5

ひとこと要約

KDEformerは、ドット積注意機構をカーネル密度推定（KDE）問題に再定式化することで、証明可能なスペクトルノルム誤差バウンドを伴う部分二次的計算を可能にし、Transformerの計算を高速化する。ImageNetでは最大18.3倍の高速化を達成し、精度低下は0.5%未満であり、BigGANではFIDスコアが正確な注意機構を上回り、FLOPsが4.14倍少ない。

ABSTRACT

Dot-product attention mechanism plays a crucial role in modern deep architectures (e.g., Transformer) for sequence modeling, however, naïve exact computation of this model incurs quadratic time and memory complexities in sequence length, hindering the training of long-sequence models. Critical bottlenecks are due to the computation of partition functions in the denominator of softmax function as well as the multiplication of the softmax matrix with the matrix of values. Our key observation is that the former can be reduced to a variant of the kernel density estimation (KDE) problem, and an efficient KDE solver can be further utilized to accelerate the latter via subsampling-based fast matrix products. Our proposed KDEformer can approximate the attention in sub-quadratic time with provable spectral norm bounds, while all prior results merely provide entry-wise error bounds. Empirically, we verify that KDEformer outperforms other attention approximations in terms of accuracy, memory, and runtime on various pre-trained models. On BigGAN image generation, we achieve better generative scores than the exact computation with over $4 imes$ speedup. For ImageNet classification with T2T-ViT, KDEformer shows over $18 imes$ speedup while the accuracy drop is less than $0.5\%$.

研究の動機と目的

• Transformerの自己注意機構における二次的計算およびメモリのボトル neck 問題を解決すること。
• 単なる要素ごとの誤差バウンドではなく、証明可能なスペクトルノルム誤差保証を持つ注意近似手法を開発すること。
• 計算複雑性の低減により、長文脈モデルにおける効率的な推論および学習を可能にすること。
• 従来の近似注意機構と比較して、言語およびビジョンモデルにおける生成性能と効率性を向上させること。

提案手法

• 対角スケーリング行列Dおよび注意行列Aの両方の計算を、一般化されたカーネル密度推定（KDE）問題に還元する。
• KDEに基づく確率を用いた重要度サンプリングにより、近似注意計算に使用する列の少数を効率的に選択する。
• 指数カーネル密度問題をガウスKDE問題に変換する変換を導入し、最先端のKDEソルバーを活用する。
• 完全な注意行列Aの明示的計算を回避するため、KDEを用いて対角スケーリング行列Dを推定する。
• 推定されたサンプリング行列Πおよび対角行列eDを用いて、部分サンプリングされた行列積を適用し、出力を効率的に計算する。
• 注意出力の安定性と忠実性を保証するスペクトルノルム誤差バウンド（式1）を提供する。

実験結果

リサーチクエスチョン

• RQ1 注意機構をカーネル密度推定問題に再定式化することで、部分二次的計算が可能になるか？
• RQ2 KDEに基づくサンプリングおよびスケーリングは、計算複雑性を低減しつつモデル性能を維持できるか？
• RQ3 注意近似においてスペクトルノルム誤差バウンドを達成でき、要素ごとのバウンドよりも強い理論的保証が得られるか？
• RQ4 提案手法は、長文脈タスクにおいて、既存の近似注意機構と比較して精度、速度、メモリ効率で優れているか？
• RQ5 KDEformerは、FLOPsを削減しながらBigGANのようなモデルの生成品質を向上させられるか？

主な発見

• BigGAN画像生成タスクにおいて、KDEformerは正確な注意機構よりもFIDスコアが低く（品質が向上）、FLOPsが4.14倍少ない。
• T2T-ViTを用いたImageNet分類では、KDEformerはトップ1精度82.08%を達成し、正確な計算と比較して0.47%の低下にとどまり、FLOPsを18.3倍削減した。
• Long Range Arenaにおけるエンドツーエンド学習では、KDEformerはテキストデータセットで正確な注意機構と比較して8倍高速であり、精度低下は最小限だった。
• LRAのデータセット全体で平均して2.28 GBのピークメモリを消費し、正確な計算（9.21 GB）および競合他手法と比較して顕著に少ない。
• LRAにおける平均テスト精度は57.15%で、全近似手法の中でも2番目に高く、最も高速かつ最もメモリ効率が良かった。
• 実験的結果から、KDEformerは複数のベンチマークでReformer、Performer、ScatterBrainをすべて上回る精度と効率性を達成した。