Graffe: Graph Representation Learning via Diffusion Probabilistic Models

TL;DR

Graffe presents a diffusion-guided, self-supervised approach to graph representation learning by coupling a graph encoder with a diffusion-based decoder conditioned on the encoder output. It provides an information-theoretic foundation, showing the denoising objective lower-bounds the conditional mutual information $I(\mathbf{x}_0; E_{\mathbf{\phi}}(\mathbf{x}_0) | \mathbf{x}_t)$ and introduces the Diff-InfoMax principle, which generalizes InfoMax across diffusion noise levels. The architecture employs a Graph-Unet-style decoder and random masking to avoid shortcut learning, achieving strong linear-probing performance on node- and graph-level tasks, with claims of state-of-the-art results on several datasets. Overall, the work demonstrates that diffusion models can be principled and effective tools for graph representation learning, offering a pathway to richer, high-frequency semantic capture in graphs.

Abstract

Diffusion probabilistic models (DPMs), widely recognized for their potential to generate high-quality samples, tend to go unnoticed in representation learning. While recent progress has highlighted their potential for capturing visual semantics, adapting DPMs to graph representation learning remains in its infancy. In this paper, we introduce Graffe, a self-supervised diffusion model proposed for graph representation learning. It features a graph encoder that distills a source graph into a compact representation, which, in turn, serves as the condition to guide the denoising process of the diffusion decoder. To evaluate the effectiveness of our model, we first explore the theoretical foundations of applying diffusion models to representation learning, proving that the denoising objective implicitly maximizes the conditional mutual information between data and its representation. Specifically, we prove that the negative logarithm of the denoising score matching loss is a tractable lower bound for the conditional mutual information. Empirically, we conduct a series of case studies to validate our theoretical insights. In addition, Graffe delivers competitive results under the linear probing setting on node and graph classification tasks, achieving state-of-the-art performance on 9 of the 11 real-world datasets. These findings indicate that powerful generative models, especially diffusion models, serve as an effective tool for graph representation learning.

Figures (4)

• Figure 1: The overall framework of Graffe. (Left) The input graph has certain nodes corrupted and is subsequently fed into a GNN encoder to obtain node representations as the condition. The decoder then receives both the noisy graph features $\mathbf{x}_t$ and the condition $\mathbf{z}$ as inputs to perform denoising, aiming to restore the original node features $\mathbf{x}_0$. (Right) The diffusion process of graph features and the architecture of GraphU-Net decoder.
• Figure 2: The comparison of denoising losses using different conditions on Cora datasets. (Vanilla) The denoising loss without condition information. (Label) Class label information obtained via linear embedding. (Representation) Learned representations obtained from Graffe.
• Figure 3: The correlation between the negative logarithm of diffusion loss (x-axis) and linear probing accuracy (y-axis) on the Photo dataset.
• Figure 4: The effect of mask ratio $m$ on Cora, Computer and MUTAG dataset.

Theorems & Definitions (17)

• Theorem 1
• Lemma 1
• Lemma 2
• Theorem 2
• proof
• Lemma 3
• Theorem 3
• proof
• Remark 1
• Theorem 1