Reimagining Graph Classification from a Prototype View with Optimal Transport: Algorithm and Theorem

Chen Qian,Huayi Tang,Hong Liang,Yong Liu
DOI: https://doi.org/10.1145/3637528.3671696
2024-01-01
Abstract:Recently, Graph Neural Networks (GNNs) have achieved inspiring performances in graph classification tasks. However, the message passing mechanism in GNNs implicitly utilizes the topological information of the graph, which may lead to a potential loss of structural information. Furthermore, the graph classification decision process based on GNNs resembles a black box and lacks sufficient transparency. The non-linear classifier following the GNNs also defaults to the assumption that each class is represented by a single vector, thereby limiting the diversity of intra-class representations. To address these issues, we propose a novel prototype-based graph classification framework that introduces the Fused Gromov-Wasserstein (FGW) distance in Optimal Transport (OT) as the similarity measure. In this way, the model explicitly exploits the structural information on the graph through OT while leading to a more transparent and straightforward classification process. The introduction of prototypes also inherently addresses the issue of limited within-class representations. Besides, to alleviate the widely acknowledged computational complexity issue of FGW distance calculation, we devise a simple yet effective NN-based FGW distance approximator, which can enable full GPU training acceleration with a marginal performance loss. In theory, we analyze the generalization performance of the proposed method and derive an O (1 over N) generalization bound, where the proof techniques can be extended to a broader range of prototype-based classification frameworks. Experimental results show that the proposed framework achieves competitive and superior performance on several widely used graph classification benchmark datasets. The code is avaliable at https://github.com/ChnQ/PGOT.
What problem does this paper attempt to address?