[論文レビュー] From block-Toeplitz matrices to differential equations on graphs: towards a general theory for scalable masked Transformers

Krzysztof Choromański, Han Lin+8 more|arXiv (Cornell University)|Jul 16, 2021

Quantum optics and atomic interactions被引用数 7

ひとこと要約

本論文は、マスクをグラフ上の位相的変調としてモデル化することにより、スケーラブルなマスク付きアテンションの一般理論的枠組みを導入する。これにより、ブロック・トーペリッツ行列構造と高速アルゴリズムを用いて、二次未満の計算が可能となる。従来の手法（線形因果アテンションや対数線形RPE-アテンション）を統合・拡張し、d次元RPEおよびグラフカーネルマスクのための新たなスケーラブルなメカニズムを提案し、その効率性を証明可能とする。

ABSTRACT

In this paper we provide, to the best of our knowledge, the first comprehensive approach for incorporating various masking mechanisms into Transformers architectures in a scalable way. We show that recent results on linear causal attention (Choromanski et al., 2021) and log-linear RPE-attention (Luo et al., 2021) are special cases of this general mechanism. However by casting the problem as a topological (graph-based) modulation of unmasked attention, we obtain several results unknown before, including efficient d-dimensional RPE-masking and graph-kernel masking. We leverage many mathematical techniques ranging from spectral analysis through dynamic programming and random walks to new algorithms for solving Markov processes on graphs. We provide a corresponding empirical evaluation.

研究の動機と目的

• 多様なマスキング機構を二次的でない複雑性を伴わずにトランスフォーマー・アーキテクチャに統合する一般的でスケーラブルな理論の構築。
• 線形因果アテンションや対数線形RPE-アテンションといった既存の効率的アテンション機構を、単一の理論的枠組みで統合すること。
• スケーラブルなアテンションの適用範囲を、d次元位置符号化およびグラフベースのカーネルマスクといった新たな設定に拡張すること。
• スペクトル解析、動的計画法、およびグラフ上のマルコフ過程に基づく数学的基盤を提供し、効率的なマスク付きアテンション計算を実現すること。
• 新しいモデル、Graph Kernel Attention Transformer (GKAT) を用いた実験的検証を通じて、本フレームワークの有効性を、最先端のGNNと比較して検証すること。

提案手法

• マスキングを、マスクがグラフ構造または距離関数から導かれる、未マスクアテンションの位相的（グラフベースの）変調として形式化する。
• マスク行列が二次未満の行列-ベクトル乗算をサポートするという性質により、効率的なマスキングを特徴づける。これにより、スケーラブルな低ランクアテンションが可能となる。
• d次元グリッド（例：画像、動画）におけるブロック・トーペリッツ構造を活用し、1次元シーケンスを超えた相対的位置符号化の一般化を実現する。
• 最短経路や拡散カーネルに特に適した、グラフベースのマスクの高速化に、高速フーリエ変換（FFT）およびハンケル行列技術を適用する。
• 木構造グラフ上で、下位から上位へのパスと上位から下位へのパスを用いた再帰的分割統治アルゴリズムを設計し、マスク付きアテンションを効率的に計算する。
• パディングおよびシーケンスパッキングを、修正されたグラフトポロジーのモデル化と木分解上の効率的計算を用いて、フレームワークに統合する。

実験結果

リサーチクエスチョン

• RQ1 どのような条件下で、マスキング機構を二次未満の時間計算量でアテンション機構に組み込むことができるか？
• RQ2 d次元相対的位置符号化は、構造的行列の性質を用いて、画像および動画モデリングにおいてスケーラブルに可能となるか？
• RQ3 最短経路や拡散カーネルに基づくグラフカーネルマスクは、低ランクアテンション設定において効率的に計算可能か？
• RQ4 因果アテンションや対数線形RPEといった既存の効率的アテンション機構は、単一の理論的枠組みで統合可能か？
• RQ5 本理論から導出された新しいスケーラブルなアテンション機構、Graph Kernel Attention Transformer (GKAT) は、既存のGNNを上回る性能を示せるか？

主な発見

• 本論文は、スケーラブルなマスク付きアテンションの一般的条件を確立した。すなわち、マスク行列が二次未満の行列-ベクトル乗算をサポートするならば、マスク付き低ランクアテンションは二次未満の時間で計算可能となる。
• 本フレームワークは、1次元RPEおよび因果アテンションをd次元RPEおよびグラフベースのマスキングの特殊ケースとして一般化し、画像や動画などの2次元・3次元入力における効率的アテンションを可能にする。
• FFTを用いた高速計算を可能にする、新たなd次元RPEメカニズムが導入され、視覚および動画タスクにおけるスケーラブルなアテンションを実現する。
• 最短経路距離および拡散カーネルを用いたグラフカーネルマスクは、FFTおよび木分解技術を用いて効率的に計算可能であることが示された。
• 木構造グラフに対する提案アルゴリズムは、$O(|V( au)| imes \text{diam}( au))$ の時間計算量を達成し、再帰的分解と巡回シフト操作を用いる。
• 実験的評価において、Graph Kernel Attention Transformer (GKAT) モデルは、9つの最先端GNNを上回る性能を示し、本理論フレームワークの実用的有効性を裏付けた。