# AugWard: Augmentation-Aware Representation Learning for Accurate Graph Classification

arXiv: https://arxiv.org/abs/2503.21105 · HTML: https://mlsanity.com/paper/2503.21105/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 35 references that are in the map.
- First posted on arXiv 27 Mar 2025.

## Builds on (most cited first, 18 of 18)

- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [Graph Contrastive Learning with Augmentations](https://mlsanity.com/paper/2010.13902/) (2020)
- [MoleculeNet: A Benchmark for Molecular Machine Learning](https://mlsanity.com/paper/1703.00564/) (2017)
- [Strategies for Pre-training Graph Neural Networks](https://mlsanity.com/paper/1905.12265/) (2019)
- [TUDataset: A collection of benchmark datasets for learning with graphs](https://mlsanity.com/paper/2007.08663/) (2020)
- [GraphMAE: Self-Supervised Masked Graph Autoencoders](https://mlsanity.com/paper/2205.10803/) (2022)
- [InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation Learning via Mutual Information Maximization](https://mlsanity.com/paper/1908.01000/) (2019)
- [Graph Random Neural Network for Semi-Supervised Learning on Graphs](https://mlsanity.com/paper/2005.11079/) (2020)
- [Optimal Transport for structured data with application on graphs](https://mlsanity.com/paper/1805.09114/) (2018)
- [Motif-based Graph Self-Supervised Learning for Molecular Property Prediction](https://mlsanity.com/paper/2110.00987/) (2021)
- [Spectral Augmentation for Self-Supervised Learning on Graphs](https://mlsanity.com/paper/2210.00643/) (2022)
- [Graph Transplant: Node Saliency-Guided Graph Mixup with Local Structure Preservation](https://mlsanity.com/paper/2111.05639/) (2021)
- [Model-Agnostic Augmentation for Accurate Graph Classification](https://mlsanity.com/paper/2202.10107/) (2022)
- [Automated Data Augmentations for Graph Classification](https://mlsanity.com/paper/2202.13248/) (2022)
- [Graph Mixup with Soft Alignments](https://mlsanity.com/paper/2306.06788/) (2023)
- [Fused Gromov-Wasserstein Graph Mixup for Graph-level Classifications](https://mlsanity.com/paper/2306.15963/) (2023)
- [Accurate Link Prediction for Edge-Incomplete Graphs via PU Learning](https://mlsanity.com/paper/2405.11911/) (2024)

## Built on by (most cited first, 2 of 2)

- [Zero-shot Quantization: A Comprehensive Survey](https://mlsanity.com/paper/2505.09188/) (2025)
- [pyMEAL: A Multi-Encoder Augmentation-Aware-Learning Toolbox for Robust Medical Image Translation](https://mlsanity.com/paper/2505.24421/) (2025)

## Abstract (arXiv)

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.