Cross-Domain Graph Data Scaling: A Showcase with Diffusion Models

TL;DR

This work tackles the challenge of scaling graph data across diverse domains by introducing UniAug, a universal graph structure augmentor built on a self-conditioned discrete diffusion model trained on thousands of graphs. In downstream tasks, UniAug uses diffusion guidance to generate augmented graph structures while preserving original node features, enabling plug-and-play improvements in graph property prediction, link prediction, and node classification. The authors demonstrate consistent cross-domain gains across 25 datasets, show favorable scaling behavior with data coverage and compute, and show that diffusion guidance and self-conditioning are essential to prevent negative transfer. This approach offers a practical path toward cross-domain graph foundation models by decoupling upstream structure understanding from downstream inductive biases, highlighting the value of diffusion-based structure augmentation for graph learning at scale.

Abstract

Models for natural language and images benefit from data scaling behavior: the more data fed into the model, the better they perform. This 'better with more' phenomenon enables the effectiveness of large-scale pre-training on vast amounts of data. However, current graph pre-training methods struggle to scale up data due to heterogeneity across graphs. To achieve effective data scaling, we aim to develop a general model that is able to capture diverse data patterns of graphs and can be utilized to adaptively help the downstream tasks. To this end, we propose UniAug, a universal graph structure augmentor built on a diffusion model. We first pre-train a discrete diffusion model on thousands of graphs across domains to learn the graph structural patterns. In the downstream phase, we provide adaptive enhancement by conducting graph structure augmentation with the help of the pre-trained diffusion model via guided generation. By leveraging the pre-trained diffusion model for structure augmentation, we consistently achieve performance improvements across various downstream tasks in a plug-and-play manner. To the best of our knowledge, this study represents the first demonstration of a data-scaling graph structure augmentor on graphs across domains.

Figures (5)

• Figure 1: The pipeline of UniAug. We pre-train a diffusion model across domains and perform structure augmentation on the downstream graphs. The augmented graphs consist of generated structures and original node features and are then processed by a downstream GNN.
• Figure 2: Normalized structural properties of Network Repository and Github Star. We enlarge the distribution coverage of our collection by combining both datasets.
• Figure 3: Effects of pre-training data scale on graph classification (left) and link prediction (right). The groups SMA, FUL, and EXT represent SMALL, FULL, and EXTRA data collection.
• Figure 4: Effects of pre-training amount of compute on graph classification (left) and link prediction (right), where one $\text{PF-days} = 10^{15} \times 24 \times 3600 = 8.64 \times 10^{19}$ floating point operations.
• Figure 5: Effects of pre-training data scale (ratio) on graph classification (left) and link prediction (right).