[Journal club] A Generalization of ViT/MLP-Mixer to Graphs

慶應義塾⼤学 杉浦孔明研究室 M1 和田唯我 / Yuiga Wada

Transcript

1. A Generalization of ViT/MLP-Mixer to Graphs Xiaoxin He1, Bryan Hooi1,2, Thomas Laurent3, Adam Perold4, Yann LeCun5,6, Xavier Bresson1 (1University of Singapore 2National University of Singapore, 3Loyola Marymount University, 4Element, Inc., 5New York University, 6Meta AI ) 慶應義塾⼤学 杉浦孔明研究室 M1 和⽥唯我 Xiaoxin He et al., “A Generalization of ViT/MLP-Mixer to Graphs” in ICML(2023)
2. 概要 2 ü 背景 • GNNは表現⼒が低く⻑距離依存を捉えきれない • TransformerはGNNよりも⾼い表現⼒を持つが，計算コストが⾼い ü 提案⼿法 • MLP-Mixerをグラフ⼀般に適⽤したGraph MLP-Mixerを提案 • 線形時間で⾼い表現⼒を実現し，⻑距離依存を捉えることが可能 ü 結果 • 様々なベンチマークにおいてSOTA⼿法に匹敵する性能
3. 背景: GNNは表現⼒に乏しく⻑距離依存を捉えきれない 3 o GNNは⼆つに⼤別可能 • Message-Passing GNNs (MP-GNNs) → 隣接するノードの情報を集約 • Weisfeiler-Leman GNNs (WL-GNNs) → 𝑘-WL グラフを識別可能なGNN o GNNの⽋点 1. 表現⼒が低い • MP-GNNsでは同型なグラフの識別すら難しい • WL-GNNsで改善したものの，未だ表現⼒は不⼗分 2. ⻑距離依存を捉えきれない • ⼀般に 𝐿-hop先のノードに情報を伝達するためには 𝐿 層必要 • 受容野のサイズが膨⼤に (e.g., ⼆分⽊の場合 𝑂(2!)) → これらの⽋点を改善する⼿法としてGraph Transformersが有望視されている C input layer X1 X2 X3 X4 F output layer Z1 Z2 Z3 Z4 hidden layers Y1 Y4 GCN [Kipf+, ICLR17]
4. 背景: Transformerと同等の性能で⾼速に動作する⼿法が望まれている 4 o Graph Transformers • ノードとエッジの埋め込みをトークンとしてTransformerに⼊⼒ • ⾼い表現⼒を持ち，⻑距離依存を捉えることが可能 • e.g., Graph Transformer [Dwivedi+, AAAI21], Graphormer [Ying+, NeurIPS21] GNN Graph Transformers 表現⼒ 低い ⾼い ⻑距離依存 △ ◦ 計算量 𝑂 (𝑁 + 𝑀) 𝑂 𝑁" / 𝑂 (𝑁 + 𝑀)" 線形時間・⾼い表現⼒・⻑距離依存を捉える⼿法が望まれている
5. 背景: MLP-Mixerをグラフ⼀般に適⽤することは有望 5 o MLP-Mixer [Tolstikhin+, NeurIPS21] • Transformerはchannel-mixingとtoken-mixingに近似されると主張 • MHA → Channel-mixing (チャネル⽅向にmixing) • MLP → Token-mixing (トークン⽅向にmixing) • 𝑂 𝑁 の計算量でVision Transformer [Dosovitskiy+, ICLR21] と同程度の性能 Token Mixer LayerNorm + Channel Mixer LayerNorm + Mixer Layer ch Embedding Global Average Pooling Graph-Level Prediction Fully Connected de PE Fuse Token & Channel Information Mixer Layers Patch PE Graph Embedding MLP-Mixer [Tolstikhin+, NeurIPS21] 画像はグラフの⼀種 → MLP-Mixerをグラフ⼀般に適⽤することは有望
6. 提案⼿法: Graph MLP-Mixer 6 o Graph MLP-Mixer: MLP-Mixerをグラフ⼀般に適⽤ • グラフに対する⾼い表現⼒ • ⻑距離依存を捉えることが可能 • 系列⻑に対して線形の計算量 (speed / memory) Token Mixer LayerNorm + Channel Mixer LayerNorm + Mixer Layer 3 6 7 8 9 1 2 3 4 8 5 6 10 11 12 1 2 3 4 5 6 7 8 9 10 11 12 3 4 5 6 7 10 12 METIS Graph Partitioning 1-Hop Overlapping Patch Extraction Graph Encoder (GNN) Patch Embedding Global Average Pooling Graph-Level Prediction Fully Connected (Generate P new patches at each epoch) 1 2 3 4 5 6 7 8 9 10 11 12 Node PE Fuse Token & Channel Information Mixer Layers Patch PE Graph Embedding
7. 画像とグラフ⼀般の違い / グラフへの応⽤における課題 7 o 画像はグラフの⼀種 • ノード: pixel • パッチ: pixelの集合 系列⻑固定 系列⻑可変 画像を分割 ？ MLPで⼗分 ？？ パッチの配置位置 ？？？ パッチを どのように定義するか Path Encoderはグラフ において何が最適か permutation invariant → 位置の埋め込みは？
8. Step1. METISを⽤いてグラフからパッチを作成 o Patch Extraction • ViTではpixel(ノード)で構成されるパッチを作成 → 同様にノードからパッチ(部分グラフ)を作成 o パッチの作成⽅法 1. METIS [Karypis+, SIAM98] を⽤いて 固定サイズの部分グラフに分割 2. 部分グラフ同⼠で1-hop分の重複が発⽣する ように，隣接ノードを部分グラフへ取り込む 1,2,3 , 4,5,6 , 7,8,9 , 10,11,12 → 1,2,3,4,8 , 4,5,6,3,7,10,12 , 7,8,9,3,6 , 10,11,12,5,6 3 6 7 8 9 1 2 3 4 8 5 6 10 11 12 1 2 3 4 5 6 7 8 9 10 11 12 3 4 5 6 7 10 12 METIS Graph Partitioning 1-Hop Overlapping Patch Extraction Pat (Generate P new patches at each epoch) 1 2 3 4 5 6 7 8 9 10 11 12 No 8 右図での例 カットエッジによる 情報ロスを防⽌
9. 3 6 7 8 9 1 2 3 4 8 5 6 10 11 12 1 2 3 4 5 6 7 8 9 10 11 12 3 4 5 6 7 10 12 METIS Graph Partitioning 1-Hop Overlapping Patch Extraction Graph Encoder (GNN) Patch Embedding Global Average Pooling Graph-Level Prediction Fully Connected (Generate P new patches at each epoch) 1 2 3 4 5 6 7 8 9 10 11 12 Node PE Fuse Token & Channel Information Mixer Layers Patch PE Graph Embedding Step2. GNNによりパッチから埋め込み表現を得る o Patch Embedding • 画像: パッチサイズが⼀定 (画像の解像度により全て同⼀) → 単純な線形層で埋め込みが可能 • グラフ: permutation invariant・パッチサイズが可変・グラフ同型問題 … etc → 埋め込み表現を得るためにGNNを使⽤ o NodePE: ノードの位置埋め込み • RWSE [Dwivedi+, ICLR21] • random-walkにより位置情報を埋め込む • LapPE [Dwivedi+, JMLR20] • Laplacian⾏列の固有ベクトルを活⽤ 9
10. 3 6 7 8 9 1 2 3 4 8 5 6 10 11 12 1 2 3 4 5 6 7 8 9 10 11 12 3 4 5 6 7 10 12 METIS Graph Partitioning 1-Hop Overlapping Patch Extraction Graph Encoder (GNN) Patch Embedding Global Average Pooling Graph-Level Prediction Fully Connected (Generate P new patches at each epoch) 1 2 3 4 5 6 7 8 9 10 11 12 Node PE Fuse Token & Channel Information Mixer Layers Patch PE Graph Embedding Step3. Mixer Layersにより特徴量を抽出 10 Token-mixer: Channel-mixer: o Mixer Layers • MLP-Mixerに則り，Token-mixer / Channel-mixerより特徴量を抽出 o Patch PE: パッチの位置埋め込み • パッチに対して疑似的な隣接⾏列 𝐴#$ % を定義 クラスタ 𝒱! , 𝒱" のカットサイズを計算 𝒱! − 𝒱" 間のエッジの数を計算 (𝐴 ∶ 隣接⾏列)
11. 3 6 7 8 9 1 2 3 4 8 5 6 10 11 12 1 2 3 4 5 6 7 8 9 10 11 12 3 4 5 6 7 10 12 METIS Graph Partitioning 1-Hop Overlapping Patch Extraction Graph Encoder (GNN) Patch Embedding Global Average Pooling Graph-Level Prediction Fully Connected (Generate P new patches at each epoch) 1 2 3 4 5 6 7 8 9 10 11 12 Node PE Fuse Token & Channel Information Mixer Layers Patch PE Graph Embedding Step3. Mixer Layersにより特徴量を抽出 11 Token-mixer: Channel-mixer: o Mixer Layers • MLP-Mixerに則り，Token-mixer / Channel-mixerより特徴量を抽出 o 最終的な出⼒ 𝑦+ • 各パッチに対してaverage pooling → MLP 各パッチ パッチ 𝑝 が空かどうか (1 if 𝑝 is non-empty else 0) ⼩さいグラフの場合，パッチが空になる場合がある
12. Graph MLP-MixerとMLP-Mixerの対応まとめ 12 o 画像はグラフの⼀種 • ノード: pixel • パッチ: pixelの集合 系列⻑固定 系列⻑可変 画像を分割 クラスタリング MLPで⼗分 GNN パッチの配置位置 Node PE / Patch PE パッチを どのように定義するか Path Encoderはグラフ において何が最適か permutation invariant → 位置の埋め込みは？
13. 実験設定: 多様なデータセット・Backboneにより実験 o データセット • Simulated: CSL, EXP, SR25, TreeNeighbourMatch • Real-world(small): ZINC, MNIST, CIFAR10 • Real-world(large): Long Range Graph Benchmark [Dwivedi+, NeurIPS22] o Patch Encoder として以下のGNNを使⽤ • GCN [Kipf+,ICLR17] • GatedGCN [Bresson+, ICLR17] • GINE [Hu+, ICLR20] GINE [Hu+, ICLR20] C input layer X1 X2 X3 X4 F output layer Z1 Z2 Z3 Z4 hidden layers Y1 Y4 GCN [Kipf+, ICLR17] 13 多様な実験で提案⼿法の 優位性を検証
14. 定量的結果: 提案⼿法がGNNの性能を向上させることを確認 14 o Patch Encoderに以下を使⽤ • GCN • GatedGCN • GINE • Graph Transformer o MLP-Mixer / ViTで性能を検証 • ViT: Channel-mixerとして MHAを使⽤ 提案⼿法がGNNの性能を向上させることを確認
15. 定量的結果: 提案⼿法は⻑距依存を捉えることが可能 15 o データセット • Peptides-func / Peptides-struct • Long Range Graph Benchmarkのタスク → ⻑距離依存に関する性能を検証 o Peptides-func → 平均 0.056 ポイントの向上 o Peptides-struct → 平均 0.028 ポイントの向上 提案⼿法は⻑距離依存を捉えられることを確認
16. 定量的結果: ⾼速に動作し，SOTA⼿法とcompetitiveな性能 16 • 各データセットにおいて最も性能の良いGNNを使⽤ (GCN / GatedGCN / GINE) ⾼速に動作 / SOTA⼿法とcompetitiveな性能
17. Ablation Study: Patch Extractionの有効性を確認 17 o Patch Extractionの性能への寄与を調査 • ZINCにおけるMAEにおいて，最⼤ 0.183 ポイントの向上 +0.183 Patch Extractionの有効性を確認 ✘: ノードをパッチ とみなして計算
18. まとめ 19 ü 背景 • GNNは表現⼒が低く⻑距離依存を捉えきれない • TransformerはGNNよりも⾼い表現⼒を持つが，計算コストが⾼い ü 提案⼿法 • MLP-Mixerをグラフ⼀般に適⽤したGraph MLP-Mixerを提案 • 線形時間で⾼い表現⼒を実現し，⻑距離依存を捉えることが可能 ü 結果 • 様々なベンチマークにおいてSOTA⼿法に匹敵する性能
19. Appendix: 各データセットの統計 20
20. Appendix: 各データセットにおけるパッチの統計 21
21. Appendix: METISアルゴリズムの有効性を検証 22
22. Appendix: 既存研究との⽐較 23