A Class-Aware Representation Refinement Framework for Graph Classification

TL;DR

CARE introduces a Class-Aware Representation rEfinement framework that injects class-level signals derived from subgraph representations into graph embeddings to improve graph classification. It is a plug-and-play augmentation that works with arbitrary GNN backbones and includes a subgraph selector, a class-aware refiner using DeepSets to build class prototypes, and a class loss to promote intra-class cohesion and inter-class separation. The paper provides VC-dimension-based generalization analysis showing CARE can have a lower upper bound than the backbone for the same parameter count, and validates effectiveness across 9 datasets and 11 backbones, with up to 11% gains in accuracy and improved convergence behavior. Ablation and case studies confirm the value of the refiner, subgraph selector, and cosine-based class loss, demonstrating enhanced class separability and practical efficiency in graph classification tasks.

Abstract

Graph Neural Networks (GNNs) are widely used for graph representation learning. Despite its prevalence, GNN suffers from two drawbacks in the graph classification task, the neglect of graph-level relationships, and the generalization issue. Each graph is treated separately in GNN message passing/graph pooling, and existing methods to address overfitting operate on each individual graph. This makes the graph representations learnt less effective in the downstream classification. In this paper, we propose a Class-Aware Representation rEfinement (CARE) framework for the task of graph classification. CARE computes simple yet powerful class representations and injects them to steer the learning of graph representations towards better class separability. CARE is a plug-and-play framework that is highly flexible and able to incorporate arbitrary GNN backbones without significantly increasing the computational cost. We also theoretically prove that CARE has a better generalization upper bound than its GNN backbone through Vapnik-Chervonenkis (VC) dimension analysis. Our extensive experiments with 11 well-known GNN backbones on 9 benchmark datasets validate the superiority and effectiveness of CARE over its GNN counterparts.

Figures (5)

• Figure 1: Framework of CARE. CARE contains four main components, including a graph encoder, a subgraph selector, a class-aware refiner and a class loss. The former two allow the flexible incorporation of existing GNN methods, while the latter two are newly proposed in our framework. We introduce them sequentially in Section \ref{['subsec:care']}.
• Figure 2: Accuracy curves of CARE-GCN and GCN on PROTEINS dataset.
• Figure 3: Loss curves of CARE-GCN and GCN on PROTEINS dataset. Note that CARE-GCN includes the class loss in the training process. The validation loss only includes the classification loss.
• Figure 4: (a) Class Separability on PROTEINS with GCN Backbone (Training Set). (b) Class Separability on PROTEINS with GCN Backbone (Test Set). The results were obtained by passing the test data once at the end of each training epoch. Note that this process doesn't affect the training in any way as the model parameters/loss are not updated when passing the test data.
• Figure 5: Visualization of Graph Representations Produced by GCN and CARE-GCN on PROTEINS dataset.