[论文解读] Test-Time Training for Graph Neural Networks

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

Advanced Graph Neural Networks被引用 4

一句话总结

该论文提出了GT3，这是首个用于图神经网络（GNNs）的测试时训练框架，旨在提升图分类任务在分布偏移情况下的泛化能力。通过在推理过程中应用自监督学习和一种适应性约束，GT3对每个测试图进行微调，显著缩小了训练集与测试集之间的性能差距，尤其在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的权重在训练过程中保持不变。
• 消融实验证实，自监督学习任务和适应性约束对实现最优性能均不可或缺。

实验结果

研究问题

• RQ1 在分布偏移条件下，测试时训练能否有效缩小图分类任务中训练集与测试集之间的性能差距？
• RQ2 如何为图数据有效设计自监督学习，以指导测试时的适应过程？
• RQ3 适应性约束在防止测试时微调过程中的特征空间畸变中起到何种作用？
• RQ4 共享层的数量以及SSL组件的选择如何影响GT3的整体性能？
• RQ5 测试时训练在理论上和实证上是否合理，能否有效提升GNN模型的泛化能力？

主要发现

• GT3在ogbg-molhiv、D&D、ENZYMES和PROTEINS等基准数据集上显著提升了图分类性能，尤其在分布偏移场景下表现优异。
• 在ogbg-molhiv数据集上，GT3显著缩小了训练集与测试集之间的性能差距，基于GIN的模型在测试集上的ROC-AUC最高达到75.2%，而验证集上为84.1%。
• 消融研究显示，若移除适应性约束、全局对比学习或局部对比学习，性能均会下降，证明所有组件均不可或缺。
• 共享层的数量影响性能表现，最优配置可显著提升准确率——例如在D&D数据集上，共享层的引入优于基线模型。
• 理论与实证分析共同表明，结合自监督学习的测试时训练对GNN而言在理论上合理且在实践中有效，尤其在分布外设置下表现突出。
• GT3在多个数据集和GNN架构上均表现出一致的性能增益，验证了其强大的泛化能力。