技術要約：グラフのための回転位置エンコーディング (WIRE)

問題提起

Transformerは、トークン間の空間的またはトポロジカルな関係性を捉えるために、位置エンコーディング（PE）に大きく依存しています。Rotary Position Encodings (RoPE) は、トークンの回転を通じて相対的な位置情報を注入できる能力により、大規模言語モデル（LLM）やVision Transformer（ViT）における標準となっていますが、グラフ構造データへの適用については未だ十分に探索されていません。

グラフには標準的な座標系が存在しないため、「位置」や「距離」の定義は非自明です。既存のグラフPEは、主に以下の2つのカテゴリに分類されます：

1. 絶対位置エンコーディング (APE): 多くの場合、グラフ・ラプラシアンの固有ベクトルに基づいています。これらは構造情報を捉えますが、一般に相対的な手法よりも性能が低く、特定の不変特性を欠いています。
2. 相対位置エンコーディング (RPE): 通常、アテンション・ロジットへのバイアスとして実装されます（例：最短経路距離や有効抵抗に基づく）。これらは効果的ですが、$N \times N$ のアテンション行列全体をインスタンス化する必要があるため、線形アテンション機構（$\mathcal{O}(N)$ でスケールする）との互換性がなく、効率的なKVキャッシュも妨げます。

本論文は次のように問いかけます：回転位置エンコーディングを、どの程度グラフ構造データに適用できるか？

手法：波誘導型回転エンコーディング (WIRE)

著者らは、グラフのスペクトルを利用してRoPEをグラフに適応させる手法である WIRE (Wave-Induced Rotary Encodings) を提案しています。

コア・アルゴリズム

1. スペクトル特徴抽出: グラフ・ラプラシアン $\mathbf{L}$ の最小 $m$ 個の固有ベクトル ($\mathbf{u}_k$) と固有値 ($\lambda_k$) を計算する。
2. 座標の定義: これらの固有ベクトルにおける各ノード $i$ の値を構成要素とするベクトル $\mathbf{r}_i$ を、スペクトル特徴として定義する（例：$\mathbf{r}_i = [\mathbf{u}_k[i]]_{k=0}^{m-1}$）。
3. 回転変換: これらのスペクトル座標を用いて、クエリ ($\mathbf{q}_i$) およびキー ($\mathbf{k}_i$) ベクトルに標準的なRoPE回転を適用する：
   $$\mathbf{z}_i \to \text{RoPE}(\mathbf{r}_i)\mathbf{z}_i$$
   ここで、$\text{RoPE}(\mathbf{r}_i)$ は、トークンベクトルの2Dブロックを、スペクトル座標によって決定される角度で回転させる。

理論的特性

本論文では、WIREに関する以下の理論的保証を確立しています：

• 置換等変性 (Permutation Equivariance): WIREは、ノードの並べ替えに対して（符号反転や退化した部分空間における回転を除いて）等変であり、これはグラフ学習において必要な特性です。
• 標準的なRoPEの回復: グリッドグラフ（パスグラフの直積）において、WIREはLLMやViTで使用される標準的なRoPEの挙動を回復します。グリッド上のスペクトル座標は空間座標と単調に対応しており、それらを $(-1, 1)$ の範囲に効果的に正規化します。
• 有効抵抗への依存性: 理論的には、WIREの周波数がランダムにサンプリングされた場合、期待される変換は、ノード間の有効抵抗に比例してクエリ・キーのアテンション・スコアを減衰させます。有効抵抗は、最短経路距離の下界となる指標です。
• 線形アテンションとの互換性: バイアスベースのRPEとは異なり、WIREはアテンション・ロジットを直接修正するのではなく、トークン（クエリおよびキー）を直接修正します。これにより、WIREはフルアテンション行列をインスタンス化することなく、線形アテンション機構（Performerなど）で使用でき、$\mathcal{O}(N)$ のスケーリングを維持できます。

主な貢献

1. WIREの導入: グラフのラプラシアン固有ベクトルを位置座標として用いることで、RoPEを一般化した新しい位置エンコーディングを提案。
2. 理論的分析:
   • WIREがグリッドグラフ上で標準的なRoPEを回復することを証明。
   • WIREが有効抵抗に基づいてアテンションを暗黙的に変調し、原理的なトポロジー的帰納バイアスを提供することを実証。
   • 以前は抵抗ベースの相対的エンコーディングでは不可能であった、線形アテンションとの互換性を確立。
3. 実験的検証: 合成タスク、点群データ、および標準的なグラフベンチマークにおける広範な実験。

実験結果

著者らは、WIREを評価するために200以上のTransformerモデルを訓練しました：

• 合成タスク:
  • 単色部分グラフ (Monochromatic Subgraphs): WIREは、特にグラフのトポロジーが複雑になる（グリッドからエッジが削除される）につれて、WIREのないベースラインと比較してテストRMSEを大幅に減少させました。
  • 最短経路予測: Watts-Strogatzグラフにおいて、WIREは有効抵抗への依存性を活用することで、APEのみのベースラインと比較してテストRMSEをほぼ半分に抑えました。
• 点群トランスフォーマー (Point Cloud Transformers): ModelNet40（分類）およびShapeNet（セグメンテーション）に適用。WIREは標準的なPCTベースラインに対して一貫した利点を提供し、SE(3)不変性を維持しながら、Cartesian RoPEと同等またはそれを上回る性能を示しました。
• グラフベンチマーク (GNNタスク):
  • MNIST、CIFAR-10、PATTERN、CLUSTER、およびOGBデータセット（例：ogbg-molhiv）を含むデータセットで評価。
  • Performer（線形アテンション）を用いた際、WIREはベースラインに対して一貫して性能を向上させました。
  • 特筆すべきは、WIREは、効率的な線形アテンションモデルと高価なフルランク・ソフトマックス・トランスフォーマーの間の性能ギャップをしばしば埋めたことです。例えば、MalNet-Tinyにおいて、WIREを導入したPerformerは、フルランクのトランスフォーマーに匹敵する精度を達成しつつ、単一のT4 GPUで訓練可能でした。
• 効率性: WIREのパラメータは全モデルパラメータの1%未満です。この手法は、ソフトマックス、線形（Performer）、SGFormer、BigBirdを含む様々なアテンション機構と互換性があります。

重要性と主張

本論文は、WIREがシーケンスや画像データにおけるRoPEの成功と、グラフ構造データのニーズとの間の溝を埋めることに成功したと主張しています。その主な意義は以下の通りです：

1. スケーラビリティ: 強力な構造的帰納バイアス（有効抵抗のような）を、以前はバイアスベースのRPEでは不可能であった線形アテンションモデルで使用することを可能にします。
2. パフォーマンス: グラフタスクにおけるTransformerの精度を向上させ、効率的な線形モデルが、より計算コストの高いソフトマックスベースのモデルの性能に匹敵することをしばしば可能にします。
3. 理論的根拠: グリッド上で既知の挙動を回復しつつ、一般的なグラフにおける有効抵抗への新たな洞察を提供する、グラフのトポロジーをアテンション機構に注入するための数学的に原理に基づいた方法を提供します。

著者らは、WIREが既存のグラフ・トランスフォーマー・アーキテクチャに対するドロップイン（そのまま追加可能な）な拡張であり、軽量で効果的な構造的帰納バイアスを提供し、回転エンコーディングを他のデータモダリティへ一般化するさらなる研究を促すものであると結論付けています。