[論文レビュー] Stable, Fast and Accurate: Kernelized Attention with Relative Positional Encoding

Shengjie Luo, Shanda Li+7 more|arXiv (Cornell University)|Jun 23, 2021

Advanced Neural Network Applications被引用数 10

ひとこと要約

この論文は、高速フーリエ変換（FFT）を用いてO(n log n)の複雑度を達成する、相対的位置符号化（RPE）を組み込んだ新しいカーネル化されたアテンションメカニズムを提案している。これにより、Transformerにおける長期間シーケンスのモデリングが安定的で高速かつ高精度に可能になる。本手法は、RPE行列のトーペリッツ構造を活用して計算を高速化するとともに、アテンションスコアの分散を低減することで訓練の安定性を向上させる。

ABSTRACT

The attention module, which is a crucial component in Transformer, cannot scale efficiently to long sequences due to its quadratic complexity. Many works focus on approximating the dot-then-exponentiate softmax function in the original attention, leading to sub-quadratic or even linear-complexity Transformer architectures. However, we show that these methods cannot be applied to more powerful attention modules that go beyond the dot-then-exponentiate style, e.g., Transformers with relative positional encoding (RPE). Since in many state-of-the-art models, relative positional encoding is used as default, designing efficient Transformers that can incorporate RPE is appealing. In this paper, we propose a novel way to accelerate attention calculation for Transformers with RPE on top of the kernelized attention. Based upon the observation that relative positional encoding forms a Toeplitz matrix, we mathematically show that kernelized attention with RPE can be calculated efficiently using Fast Fourier Transform (FFT). With FFT, our method achieves $\\mathcal{O}(n\\log n)$ time complexity. Interestingly, we further demonstrate that properly using relative positional encoding can mitigate the training instability problem of vanilla kernelized attention. On a wide range of tasks, we empirically show that our models can be trained from scratch without any optimization issues. The learned model performs better than many efficient Transformer variants and is faster than standard Transformer in the long-sequence regime.

研究の動機と目的

• 標準的な自己アテンションが長期間シーケンスにおいて二次的複雑度のため非効率であるという問題に対処すること。
• 最先端モデルで一般的に採用されている相対的位置符号化（RPE）を用いるTransformerにおける効率的なアテンション計算を可能にすること。
• ドット-次に指数関数を適用するアテンションの枠組みを超える表現力を持つRPEに対応できない既存の近似手法の限界を克服すること。
• 実際の運用でしばしば最適化の不安定性を示すカーネル化アテンションモデルの訓練を安定化させること。
• 微調整に依存せずに、長期間シーケンス領域において計算効率と高い性能を両立すること。

提案手法

• 相対的位置符号化行列のトーペリッツ構造を活用することで、RPEを組み込んだカーネル化アテンションを再定式化する。
• 高速フーリエ変換（FFT）を用いて、RPE行列とカーネル化特徴行列間の行列-ベクトル積を高速化し、時間的複雑度をO(n log n)に低減する。
• 標準的なソフトマックスアテンションの代わりに、特徴写像を用いてアテンション関数を近似するカーネル化アテンションメカニズムを導入する。
• クエリとキーのノルムが大きい場合に特に有効な分散低減と訓練安定性の向上を図るため、正規化アテンション（NPRF）を導入する。
• エンコーダーとデコーダーの両ブロックにカーネル化アテンションメカニズムを適用し、Transformerアーキテクチャへの完全統合を可能にする。
• さまざまな特徴写像関数（例：PRF, TRF, Sphere-PRF, ORF）をサポートし、異なる選択肢において一般化性を示す。

実験結果

リサーチクエスチョン

• RQ1 標準的なドット-次に指数関数を適用する定式化を超えて、相対的位置符号化を組み込んだカーネル化アテンションを効率的に計算できるか？
• RQ2 カーネル化アテンションにRPEを導入することで、通常のカーネル化アテンションと比較して訓練の安定性が向上するか？
• RQ3 RPE行列のトーペリッツ構造を活用することで、アテンション計算における準二次的時間複雑度を達成できるか？
• RQ4 提案手法の性能は、標準的および効率的なTransformerの変種と比較して、長期間シーケンスタスクでどうなるか？
• RQ5 提案手法は、カーネル化アテンションメカニズムにおける異なる特徴写像次元およびタイプに対して頑健か？

主な発見

• 本手法は、RPE行列のトーペリッツ構造を活用し、FFTを用いることでO(n log n)の時間的複雑度を達成し、効率的な長期間シーケンスモデリングを可能にする。
• 正規化カーネル化アテンションとRPEを組み合わせたモデルは、最適化の問題が生じないまま、訓練をスクラッチから安定して実行可能である。これは、通常のカーネル化アテンションとは対照的である。
• ImageNetデータセットでは、RPEを組み込んだNPRF-Transformerがトップ-1正確度80.9%を達成し、標準的なDeiT-base（81.2%）およびPRF変換モデル（79.5%）を上回った。
• IWSLT14ドイツ語-英語翻訳タスクでは、RPEと正規化アテンションを備えたモデルがBLEUスコア32.1を達成し、well-trainedな標準Transformer（32.3）と同等の性能を示した。同時に、線形複雑度を維持した。
• 特徴写像次元に対して本手法は頑健であり、16次元でも性能低下が最小限に抑えられ、さまざまな特徴写像関数（PRF, TRF, Sphere-PRF, ORF）に一般化する。
• RPEの導入により、標準アテンションからカーネル化アテンションに変換する際の近似誤差が低減され、特に正規化と組み合わせた場合に顕著に性能劣化が小さくなる。アブレーションスタディでその傾向が示された。