GNNのUPDATE関数を深掘りする|オーバースムージングとスキップ接続
GNNのメッセージパッシングはAGGREGATEとUPDATEの2ステップで構成されます。この記事ではUPDATEに注目します。メッセージパッシングを重ねると顕著になる、オーバースムージング(ノード表現が均質化する現象)を緩和するため、UPDATE関数の設計でも様々な工夫が凝らされています。本記事ではスキップ接続・ゲート付き更新・Jumping Knowledge Connectionsの3手法を、理論的な裏付けとともに解説します。
オーバースムージング問題|層を重ねるほど表現が均質化する理由
オーバースムージングとは何か
GNNでメッセージパッシングを繰り返すと、各ノードの埋め込みは周辺ノードの情報を段階的に取り込みます。GNNでは敢えてこの動作を設計する訳ですが、層数 を大きくしすぎるとグラフ内のすべてのノード表現が均質化してしまいます。この現象をオーバースムージング(over-smoothing)と呼びます。
流体シミュレーションの非構造格子を例に挙げると、初期状態では各節点が局所的な速度・圧力の情報を保持しています。しかしGNNで多数の更新を繰り返すと、すべての節点が「グラフ全体の平均的な状態」を表すようになり、衝撃波の位置や渦の中心といった局所的な特徴が失われてしまいます。
ノード間の影響度を定量化する
Xu et al. [2018] は、あるノード の初期特徴量がノード の最終埋め込みにどれだけ影響するかを定量化する影響度 を定義しました。
ここで はノード の 層目の埋め込み、 はノード の初期埋め込み(入力特徴量)、 は全要素が1のベクトルです。ヤコビアン行列 の各成分は「 の初期表現の変化が の最終表現の各次元に与える影響」を表します。例えば、埋め込みを長さ2のベクトルだとすれば、ヤコビアンは2x2の行列になるので、 式(1)の右辺は となり、ヤコビアンの全成分の和が と定義されていることがわかります。
自己ループGNNにおける定理
Xu et al. [2018] は、自己ループ型GNN(隣接行列に自己ループを加えた上で正規化集約を使うGCNスタイル)に対して次の定理を証明しています。
定理:自己ループを用いる正規化集約(GCNスタイル)では、
が成り立ちます。ここで はノード を起点とした長さ のランダムウォークでノード に到達する確率です。
この定理は、 の極限、つまり無限回ランダムウォークすれば は均一な分布に収束するのだから、 を大きくするとどのノードへの影響度も均一化してしまうことを示唆しています。つまり、GCNスタイルのGNNでは層を深くするほどオーバースムージングが生じます。
この問題は自己ループ型に限らず基本GNNにも及びます。各層で が成り立つ場合(近傍情報を自身よりも強く取り込む場合)、オーバースムージングが生じることが知られています。
スキップ接続によるUPDATE関数の改善
スキップ接続の基本アイデア
オーバースムージングの根本原因は、UPDATE関数が前の層の埋め込み の情報を近傍からのメッセージ に「徐々に上書き」してしまうことにあります。これを防ぐ最も直接的な方法は、ノード自身の前の埋め込みを明示的に保持することです。
最も単純なスキップ接続は、ベースのUPDATE関数の出力にノード自身の前の埋め込みを連結することで実現されます。
ここで はベクトルの連結です。この連結により、メッセージパッシングの各ステップで「前の層における自分の表現」が明示的に保持されます。
この手法はGraphSAGE [Hamilton et al., 2017] で提案されたもので、畳み込みニューラルネットワーク(CNN)における残差接続(ResNet [He et al., 2016])のGNN版として位置づけられます。
線形補間による重み付きスキップ接続
連結よりも柔軟なアプローチとして、ゲートベクトル を用いた線形補間があります。
ここで 、 は要素ごとの積(Hadamard積)を表します。 が0に近いほど前の自分の表現を維持し、1に近いほど近傍情報を強く取り込みます。ゲートベクトル は学習パラメータとして直接最適化することも、現在の隠れ状態を入力とするMLPや1層GNNの出力として生成することもできます [Pham et al., 2017]。
流体計算との対応で考えると、例えば非定常流れの予測において、ゲートベクトルは現在の時間ステップの状態をどれだけ保持するかを調節する役割を担うと言えます。
スキップ接続が有効な場面
スキップ接続(連結・線形補間)には特に次のような特徴があります。
- メッセージパッシングが2〜5層のGNNに向いている
- ノード分類が得意
- 局所的な特徴量情報を近傍の影響から守ってくれる
ゲート付き更新とJumping Knowledge|深いGNNを実現する2つの戦略
ゲート付き更新:RNNの仕組みをGNNに導入する
スキップ接続による改善は効果的ですが、10層以上の深いGNNでは依然としてオーバースムージングが生じやすいです。これに対して、リカレントニューラルネットワーク(RNN)の学習安定化技術をUPDATE関数に直接適用するアプローチがあります。
GNNのメッセージパッシングをRNNの視点で解釈すると、各ノードの埋め込み が「隠れ状態」、近傍から集約されたメッセージ が「観測」に相当します。この対応関係を活かして、GRU(Gated Recurrent Unit)[Cho et al., 2014] をUPDATE関数として使うと次のように定義されます。
GRUはリセットゲートと更新ゲートの2つのゲート機構を持ち、前の状態のどの成分を引き継ぎ、どの成分を近傍情報で更新するかを自動的に学習します。GRUの代わりにLSTM(Long Short-Term Memory)を用いることもあります [Selsam et al., 2019]。
重要な点は、パラメータがすべてのノードで共有されることです。グラフ全体で同一のGRU/LSTMセルを使って各ノードを更新し、層(メッセージパッシングの反復)をまたいでもパラメータは共有されます。
ゲート付き更新が有効な場面
ゲート付き更新は次のような場面で特に威力を発揮します。
- 10層以上の深いGNNが必要なタスク
- グラフの大域的な構造についての推論が必要なタスク
- プログラム解析 [Li et al., 2015] や組み合わせ最適化 [Selsam et al., 2019] への応用
流体力学の応用では、複雑な配管ネットワークの流量計算がこれに相当します。局所的な節点情報だけでなく、グラフ全体を通じた圧力分布・流量収支を正確に捉えるためには遠距離のノード間の依存関係を維持する必要があり、ゲート付き更新が適しています。
Jumping Knowledge Connections:全層の埋め込みを活用する
スキップ接続もゲート付き更新も、層ごとの更新設計を工夫することでオーバースムージングを抑制します。これらとは別のアプローチとして、各層の埋め込みを最終出力に直接利用する Jumping Knowledge (JK) connections [Xu et al., 2018] があります。
標準的なGNNでは最終層 の埋め込みのみを出力として使います。
これに対しJK connectionsでは、全層の埋め込みを集約して最終表現を構成します。
ここで は任意の微分可能関数であり、全層の埋め込みを入力として最終表現を生成します。 の具体的な選択肢は以下のとおりです。
- 全層の埋め込みをそのまま連結する。( が恒等関数の場合に相当)
- 各次元で全層にわたる最大値を取る
- LSTMを用いて各層の埋め込みに重みを付けて集約する
JK connectionsの直感的な意味は、「浅い層(メッセージパッシングの回数が少ないとき)の局所的な情報も深い層(回数が多いとき)の大域的な情報も、どちらも最終出力で利用できるようにする」ことです。ノード埋め込みは1ホップ近傍の局所的な特徴を保持しており、深い層の埋め込みはより広い近傍の情報を含みます。 が全層の中から適切な情報を選択することで、タスクに応じた最適な「視野の広さ」を自動的に学習できます。
Xu et al. [2018] は幅広いタスクにわたってJK connectionsが安定した性能向上をもたらすことを実験的に示しており、汎用的に有効な戦略として広く用いられています。
3手法の比較まとめ
| 手法 | 適切な層数 | 主な用途 | 計算コスト |
|---|---|---|---|
| スキップ接続(連結・補間) | 2〜5層 | ノード分類(同類性の強いグラフ) | 低 |
| ゲート付き更新(GRU/LSTM) | 10層以上 | 大域的推論が必要なタスク | 中 |
| JK connections | 任意 | 汎用的な性能向上 | 低〜中 |
これら3つの手法は互いに排他的ではありません。ゲート付き更新とJK connectionsを組み合わせることも可能であり、タスクや問題規模に応じた組み合わせを選ぶことが実践的なアプローチとなります。
さらに深く学ぶための書籍
本記事のテーマをより深く学びたい方には、以下の書籍をおすすめします。
- Graph Representation Learning(William L. Hamilton, Morgan & Claypool, 2020):本記事の内容の原典となる教科書です。GNNのAGGREGATE/UPDATE設計の詳細を理論的な裏付けとともに解説しており、GNNを体系的に学ぶための最良の一冊です。