はじめに

Graph Warp Module: an Auxiliary Module for Boosting the Power of Graph Neural Networksを読んだのでメモ．

気持ち

GNNの表現力の低さを問題視した論文．この問題は以前読んだICLR2019のHow Powerful are Graph Neural Networks?でも議論されている．GNNは（データとタスクによるが）一般的なNNに比べ学習データにoverfitすることすらできないことがある．そのためこの論文ではGNNの表現力を向上させるgraph warp module (GWM)を提案するというもの．GWMは既存のGNNに追加モジュールとして組み合わせることが可能．

GWMのポイントはグラフの持つグローバルな特性を捉えるためのvirtual supernodeの導入と，正確なノード間の情報のやりとりを達成するためのGRUとattention利用した構造の2点とのこと．

notation

グラフを $G=(V,E)$ とし， $V,E$ はそれぞれノードとエッジの集合を表す．ノードのラベルを $i=1,2,\dots,|V|$ として，各エッジをノードのペアとして表現する．隣接行列 $\mathcal{A}\in\mathbb{R}^{|V|\times |V|}$ は重み付きとし，各ノードには特徴ベクトル $x_i$ を割り振る．ここでは多層構造のGNNを，平滑化関数 $\mathcal{F}_l$ を利用して再帰的に出力を計算するものとして考える．初期値を $x_j=h_{0,j}$ とし $h_{l,i}=\mathcal{F}_{l-1,i}(h_{l-1,j};j\in V)\in\mathbb{R}^d$ をGNNの $l$ 層目によって $i$ 番目のノードに割り当てられた特徴ベクトルとする．最終的に得られた特徴ベクトルの集合 $\{ h_{L,i};i\in V\}$ を集約することで出力を得る．

Graph Warp Module

GWMはsuper node, transmitter unit, warp gate unitの3つのブロックから構成されていて，GWM付きのGNNは $l-1$ 層目からのメッセージ $\mathcal{F}_{l-1,i}(h_{l-1,i};j\in V)\in\mathbb{R}^d$ をsupernodeの値と合わせ混んでl層目の出力 $h_l$ として返す．

Supernode

Supernodeは全てのノードと接続を持つノードで，GNNの各層に用意される．Supernodeは全てのノードと接続を持つため，グローバルな情報を伝達する助けとなるとのこと． $l$ 層目のsupernodeの特徴ベクトルを $g_l$ とすると，各層においてtransmitterと呼ばれるunitはsupernodeから $l+1$ 層目に渡すメッセージ $\mathcal{g}_l(g_l)$ を要求し，メインのGNNにメッセージを送る．ただし， $\mathcal{g}_l$ はsmooth function．

$l$ 層目のsupernodeの特徴ベクトルは $l-1$ 層目のsupernodeの特徴ベクトルを利用して作られるため，初期値 $g_0$ は何らかの値で初期化されている必要がある．ここでは初期値の例としてノードやエッジの数，入力ノードの特徴量の平均やヒストグラムなどを挙げている．

Transmitter Unit

Transmitter unitはGNNとGWM間のメッセージのやり取りを扱うユニット．ここでは複数のタイプのメッセージを扱うために，attentionを利用する工夫を入れている．GNNからsupernodeにメッセージを送る前に，transmitterは $K$ -head attention mechanismを使ってメッセージをタイプごとに集約する．処理の流れとしては，GNNからsupernodeへと送るメッセージを作成 $m_{l,k}^{main\rightarrow super}$ ，伝送 $h_l^{main\rightarrow super}$ ．その後supernodeからのメッセージをGNNへ伝送 $g_l^{super\rightarrow main}$ の流れ．細かい計算は次のようになる．

$\displaystyle h_l^{main\rightarrow super}=\mathrm{tanh}\left(W_lm_{l,1:k}^{main\rightarrow super}\right)\in\mathbb{R}^{D'}\\ m_{l,k}^{main\rightarrow super}=\sum_i\alpha_{l,i,k}U_{l,k}h_{l-1,i}\in\mathbb{R}^{D'}\\ g_{l}^{super\rightarrow main}=\mathrm{tanh}(F_lg_{l-1})\in\mathbb{R}^D \alpha_{l,i,k}=\mathrm{softmax}(a(h_{l-1},g_{l-1};A_{l,k}))\in(0,1)\\ a(h_{l-1},g_{l-1};A_{l,k}):=h_{l-1}^T,A_{l,k}g_{l-1}$

$\alpha_{l,i,k}$ の部分がattention．

Warp Gate

Warp geteは送られてきたメッセージをマージしてその結果をGRUを通してsupernodeとGNNに送るというもの．構成要素は以下．

・ $h_l^0$ : $l$ 層目のGNNにメッセージを伝送するGRUの入力・ $g_l^0$ : $l$ 層目のsupernodeにメッセージを伝送するGRUの入力・ $\hat{g}_l$ : $l-1$ 層目のsupernodeからのメッセージ $\mathcal{g}_{l-1}(g_{l-1}\in\mathbb{R}^{D'}$ ・ $\hat{h}_{l,i}$ : $l-1$ 層目のGNNからのメッセージ $\mathcal{F}_{l-1,i};k\in V)\in\mathbb{R}^D ・[tex:z_{l,i}$ :supernodeからGNNへの伝送のためのwarp gate coefficients ・ $z_{l,i}^{(S)}$ :GNNからsupernodeへの伝送のためのwarp gate coefficients

各変数の具体的な定義は以下．

$\displaystyle h_{l,i}^0=(1-z_{l,i})\odot\hat{h}_{l-1,i}+z_{l,i}\odot g_l^{super\rightarrow main}\in\mathbb{R}^D\\ z_{l,i}=\sigma\left(H_l\tilde{h}_{l,i}+G_lg_l^{super\rightarrow main}\right)\\ g_l^0=z_l^{(S)}\odot h_l^{main\rightarrow super}+(1-z_l^{(S)})\odot\hat{g}_l\in\mathbb{R}^{D'}\\ z_l^{(S)}=\sigma\left(H_l^{(S)}h_l^{main\rightarrow super}+G_l^{(S)}\hat{g}_l\right)$

これらをGRUに通してGNNとsupernodeのメッセージを組み合わせた値を返す．

$\displaystyle h_{l,i}=\mathrm{GRU}(h_{l-1,i},h_{l,i}^0)\in\mathbb{R}^D\\ g_l=\mathrm{RGU}(g_{l-1},g_l^0)\in\mathbb{R}^{D'}$

まとめ

モチベーションははHow Powerful are Graph Neural Networks?と同じだがこの論文は問題の解決方法がpracticalで個人的にはHow Powerful are Graph Neural Networks?の方が好き．ただ，既存の枠組みに導入するだけで性能改善可能という点は扱いやすく，良い手法．