AugWard: Augmentation-Aware Representation Learning for Accurate Graph Classification

2025-03-27

AI Summary

Graph neural networks (GNNs) for graph classification suffer from overfitting, while existing graph augmentation methods lack sufficient representation robustness. Method: This paper proposes an augmentation-aware representation learning framework. Its core innovations include: (i) modeling augmentation discrepancy as a learnable graph distance prediction task, jointly constraining structural and feature-level differences between augmented and original graphs; (ii) introducing an augmentation-aware supervised training paradigm that explicitly accounts for varying augmentation strengths—overcoming a key limitation of conventional contrastive learning; and (iii) incorporating multi-level consistency regularization to enhance classifier robustness against diverse augmented representations. Contribution/Results: The framework achieves state-of-the-art performance across supervised, semi-supervised, and cross-domain graph classification benchmarks, demonstrating significant improvements in generalization and robustness.

Technology Category

• Machine Learning: Graph-based Machine Learning
• Computer Vision: Representation Learning for Vision
• Data Mining & Knowledge Management: Graph Mining, Social Network Analysis & Community

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

How can we accurately classify graphs? Graph classification is a pivotal task in data mining with applications in social network analysis, web analysis, drug discovery, molecular property prediction, etc. Graph neural networks have achieved the state-of-the-art performance in graph classification, but they consistently struggle with overfitting. To mitigate overfitting, researchers have introduced various representation learning methods utilizing graph augmentation. However, existing methods rely on simplistic use of graph augmentation, which loses augmentation-induced differences and limits the expressiveness of representations. In this paper, we propose AugWard (Augmentation-Aware Training with Graph Distance and Consistency Regularization), a novel graph representation learning framework that carefully considers the diversity introduced by graph augmentation. AugWard applies augmentation-aware training to predict the graph distance between the augmented graph and its original one, aligning the representation difference directly with graph distance at both feature and structure levels. Furthermore, AugWard employs consistency regularization to encourage the classifier to handle richer representations. Experimental results show that AugWard gives the state-of-the-art performance in supervised, semi-supervised graph classification, and transfer learning.

Problem

• Mitigating overfitting in graph neural networks for classification
• Enhancing graph representation diversity via augmentation-aware training
• Improving classification accuracy with consistency regularization techniques

Innovation

• Augmentation-aware training for graph distance prediction
• Aligns representation difference with graph distance
• Uses consistency regularization for richer representations