[論文レビュー] Test-Time Training for Graph Neural Networks

Yiqi Wang, Chaozhuo Li+5 more|arXiv (Cornell University)|Oct 17, 2022

Advanced Graph Neural Networks被引用数 4

ひとこと要約

本稿では、グラフ分類における分布シフト下での一般化性能を向上させるために、GNNのテスト時学習フレームワークであるGT3を提案する。自己教師あり学習と適応制約を推論時に適用することで、GT3は各テストグラフ上でGNNをファインチューニングし、特に ogbg-molhiv や D&D のような分布シフトが著しいデータセットにおいて、訓練集合とテスト集合の間の性能差を顕著に低減する。

ABSTRACT

Graph Neural Networks (GNNs) have made tremendous progress in the graph classification task. However, a performance gap between the training set and the test set has often been noticed. To bridge such gap, in this work we introduce the first test-time training framework for GNNs to enhance the model generalization capacity for the graph classification task. In particular, we design a novel test-time training strategy with self-supervised learning to adjust the GNN model for each test graph sample. Experiments on the benchmark datasets have demonstrated the effectiveness of the proposed framework, especially when there are distribution shifts between training set and test set. We have also conducted exploratory studies and theoretical analysis to gain deeper understandings on the rationality of the design of the proposed graph test time training framework (GT3).

研究の動機と目的

• グラフ分類において、分布シフト下でも継続的に顕在する訓練集合とテスト集合の間の性能ギャップを是正すること。
• 再訓練を一切行わず、事前学習済みGNNを個々のテストサンプルに適応させるテスト時学習フレームワークを設計すること。
• テスト時適応における特徴空間の歪みを防ぐために制約機構を組み込むことで、モデルの頑健性を確保すること。
• 実証的および理論的分析を通じて、グラフ領域におけるテスト時学習の有効性と妥当性を検証すること。

提案手法

• GT3は、局所的およびグローバルなグラフ構造とノード属性を活用して、テスト時適応のための監視信号を生成する自己教師あり学習（SSL）タスクを採用する。
• フレームワークは、テスト時学習におけるモデル更新を正則化する適応制約を導入し、有害な特徴空間の歪みを防ぐ。
• 主タスクとSSLタスクの間に共有層を設け、一般化性能と適応能力のバランスを最適化するために共有層の数を調整する。
• 同じグラフの増幅ビューを比較するコントラスト学習目的関数を用いて、各テストグラフに対して反復的な更新を実行する。
• この手法は推論時のみに適用され、事前学習済みGNNの重みは学習段階で変更されない。
• アブレーションスタディにより、SSLタスクと適応制約の両方が最適性能を達成するために不可欠であることが確認された。

実験結果

リサーチクエスチョン

• RQ1 分布シフト下において、テスト時学習が訓練集合とテスト集合の間の性能ギャップを効果的に低減できるか？
• RQ2 グラフデータに適した自己教師あり学習をどのように設計すれば、テスト時適応を効果的に導けるか？
• RQ3 適応制約は、テスト時ファインチューニングにおける特徴空間の歪みを防ぐために果たす役割は何か？
• RQ4 共有層の数やSSLコンポonentsの選択が、GT3全体の性能に与える影響は何か？
• RQ5 GNNにおけるテスト時学習は、理論的および実証的に一般化性能の向上に妥当であると裏付けられるか？

主な発見

• GT3は、ogbg-molhiv や D&D、ENZYMES、PROTEINS などのベンチマークデータセットにおいて、特に分布シフト下でグラフ分類性能を顕著に向上させる。
• ogbg-molhiv では、GT3により訓練集合とテスト集合の性能ギャップが縮小され、GINベースのモデルがテストセットで最大75.2%のROC-AUCを達成したのに対し、バリデーションセットでは84.1%であった。
• アブレーションスタディにより、適応制約、グローバルコントラスト学習、ローカルコントラスト学習を削除すると性能が低下し、すべてのコンponentsが最適性能を達成するために不可欠であることが示された。
• 共有層の数は性能に影響を与え、最適な設定により精度が向上する—例として、D&Dでは共有層の導入によりベースラインを上回る性能が得られた。
• 理論的および実証的分析により、SSLを用いたテスト時学習がGNNにおいて理にかなっており、特に分布外設定において有効であることが確認された。
• GT3は複数のデータセットおよびGNNアーキテクチャにわたり一貫した性能向上を示し、その一般化能力を裏付けた。