[論文レビュー] Scatterbrain: Unifying Sparse and Low-rank Attention Approximation

Beidi Chen, Tri Dao+4 more|arXiv (Cornell University)|Oct 28, 2021

Advanced Neural Network Applications参考文献 66被引用数 9

ひとこと要約

Scatterbrainは、局所性に敏感なハッシュ（LSH）によるスパarsityとカーネル特徴写像による低ランク構造を用いて、効率的なトランスフォーマーにおけるスパースおよび低ランク近似を統合的かつ理論的に裏付けられた方法で組み合わせる。この手法は、ベースラインと比較して最大2.1倍の低い近似誤差を達成し、視覚モデルでは1%の精度低下で注意メモリを98%削減する。スパースまたは低ランク手法を単独で用いる場合よりも、言語および視覚タスクの両方で優れた性能を発揮する。

ABSTRACT

Recent advances in efficient Transformers have exploited either the sparsity or low-rank properties of attention matrices to reduce the computational and memory bottlenecks of modeling long sequences. However, it is still challenging to balance the trade-off between model quality and efficiency to perform a one-size-fits-all approximation for different tasks. To better understand this trade-off, we observe that sparse and low-rank approximations excel in different regimes, determined by the softmax temperature in attention, and sparse + low-rank can outperform each individually. Inspired by the classical robust-PCA algorithm for sparse and low-rank decomposition, we propose Scatterbrain, a novel way to unify sparse (via locality sensitive hashing) and low-rank (via kernel feature map) attention for accurate and efficient approximation. The estimation is unbiased with provably low error. We empirically show that Scatterbrain can achieve 2.1x lower error than baselines when serving as a drop-in replacement in BigGAN image generation and pre-trained T2T-ViT. On a pre-trained T2T Vision transformer, even without fine-tuning, Scatterbrain can reduce 98% of attention memory at the cost of only 1% drop in accuracy. We demonstrate Scatterbrain for end-to-end training with up to 4 points better perplexity and 5 points better average accuracy than sparse or low-rank efficient transformers on language modeling and long-range-arena tasks.

研究の動機と目的

• スパースおよび低ランクの注意近似を統合し、現在は異なる領域で優れた性能を発揮しており、しばしば別個に使用されているものを統合すること。
• 長文書処理における効率性と精度のトレードオフを、両方の近似タイプを組み合わせることで解消すること。
• 二重カウントの問題を回避するため、理論的に正確で偏りのない、かつ効率的な手法を構築すること。
• スパース＋低ランク近似が、画像生成や言語モデリングなどの多様なタスクにおいて、個別の手法を上回ることを実証的に検証すること。
• 事前学習済みモデルにおいて、フル注意を最小限の精度低下と顕著なメモリ削減で即座に置き換えられるようにすること。

提案手法

• Scatterbrainは、フル行列を明示的に生成せずに、ソフトマックス化された注意行列のスパース成分を特定・近似するために局所性に敏感なハッシュ（LSH）を用いる。
• 低ランク成分の近似にはカーネル特徴写像を適用し、ランダム特徴近似を用いて効率的な計算を実現する。
• 注意行列をスパース成分と低ランク成分の和に分解し、Robust PCAにインspiredされた構造的分解により、二重カウントを回避する。
• 低ランクベースライン単体と比較して、近似誤差が厳密に小さいという理論的保証を提供する。
• 標準的な注意層の即座の置き換えとして設計されており、Reformer や Performer などの既存の効率的トランスフォーマー・アーキテクチャと互換性を持つ。
• ソフトマックス温度および注意分布のエントロピーに基づいて、最適なスパースおよび低ランク成分の組み合わせを動的に選択する。

実験結果

リサーチクエスチョン

• RQ1 スパース、低ランク、または統合近似が、ソフトマックスエントロピーに基づくどの領域（例：中間エントロピー領域）で注意行列に対して最も優れた性能を発揮するか？
• RQ2 スパースおよび低ランク近似を統合した統一的手法が、個別に用いた場合よりも低い近似誤差を達成できるか？
• RQ3 二重カウントやバイアスの導入を回避する理論的に妥当で効率的な方法で、スパースおよび低ランク成分を組み合わせられるか？
• RQ4 向上した近似誤差が、言語モデリングや画像分類などの下流タスクにおけるエンドツーエンドの性能にどのように反映されるか？
• RQ5 Scatterbrainは、フル注意の即座の置き換えとして、最小限の精度低下と顕著なメモリ削減を実現できるか？

主な発見

• 事前学習済みの T2T-ViT において、Scatterbrainはトップ1精度が1%低下するのみで注意メモリを98%削減し、SMYRF や Local attention を上回る。
• BigGAN における画像生成タスクでは、Scatterbrainがフル注意の即時置き換えとして用いられた際、ベースライン手法と比較して2.1倍低い近似誤差を達成した。
• 言語モデリングおよび Long Range Arena タスクにおけるエンドツーエンド学習では、Scatterbrainはスパースまたは低ランクベースラインと比較して、最大4ポイント低い perplexity および5ポイント高い平均精度を達成した。
• GLUEベンチマークにおいて、Scatterbrainは、BERTの注意層を置き換えた際、全9つの下流タスクのうち8つで、Performer や Smyrf を含むすべてのベースラインを上回った。
• 可視化分析により、スパース＋低ランク近似が、特に個別手法が失敗する中間エントロピー領域において、全層で最小の誤差を示すことが確認された。
• 本手法は、階層的（例：言語）および非階層的（例：画像）タスクの両方で最先端の性能を達成し、広範な適用可能性を示した。