SwinGNN: Rethinking Permutation Invariance in Diffusion Models for Graph Generation

TL;DR

SwinGNN rethinks permutation invariance in graph diffusion by embracing a non-invariant backbone that uses a 2-WL-inspired edge-to-edge transformer with shifted-window attention, coupled with an invariant-sampling trick via random permutations. The authors show that invariant losses induce target distributions with exponentially many modes, making learning harder, and demonstrate that invariant sampling can be achieved without invariant losses. Through SGD-based diffusion with EDM preconditioning and self-conditioning, SwinGNN achieves state-of-the-art results on synthetic, protein, and molecule graph generation tasks, while remaining scalable to graphs with hundreds of nodes. The combination of architectural design, training/sampling techniques, and a simple permutation-based sampling fix yields practical, high-quality graph generation with strong empirical validation and broad applicability.

Abstract

Diffusion models based on permutation-equivariant networks can learn permutation-invariant distributions for graph data. However, in comparison to their non-invariant counterparts, we have found that these invariant models encounter greater learning challenges since 1) their effective target distributions exhibit more modes; 2) their optimal one-step denoising scores are the score functions of Gaussian mixtures with more components. Motivated by this analysis, we propose a non-invariant diffusion model, called $\textit{SwinGNN}$, which employs an efficient edge-to-edge 2-WL message passing network and utilizes shifted window based self-attention inspired by SwinTransformers. Further, through systematic ablations, we identify several critical training and sampling techniques that significantly improve the sample quality of graph generation. At last, we introduce a simple post-processing trick, $\textit{i.e.}$, randomly permuting the generated graphs, which provably converts any graph generative model to a permutation-invariant one. Extensive experiments on synthetic and real-world protein and molecule datasets show that our SwinGNN achieves state-of-the-art performances. Our code is released at https://github.com/qiyan98/SwinGNN.

Figures (12)

• Figure 1: Data distribution and target distribution for a 3-node tree graph. For permutation matrix $\boldsymbol{P}_i$ and adjacency matrix $\boldsymbol{A}_{i}$, filled/blank cells mean one/zero. The probability mass function (PMF) highlights the difference in modes. Our example also shows graph automorphism (e.g., $\boldsymbol{P}_1$ and $\boldsymbol{P}_2$).
• Figure 2: Invariant models perform significantly worse than non-invariant models when the number of applied permutations ($l$) is small.
• Figure 3: The overall architecture of our SwinGNN.
• Figure 4: Qualitative results on plain graph and molecule datasets.
• Figure 5: Different configurations of $TV(q, p_{\rm{data}})$, given fixed $\mathcal{A}$ (the support of $p_{\rm{data}}$) and its induced $\mathcal{A}^*$. Here, we modify $\mathcal{Q}$ (the support of $q$). Left: maximal TV, $\mathcal{Q}$ is disjoint from $\mathcal{A}/\mathcal{A}^*$. Middle: intermediate TV, intersecting $\mathcal{Q}$ and $\mathcal{A}/\mathcal{A}^*$. Right: minimal TV, $\mathcal{Q} = \mathcal{A}^*$.

Theorems & Definitions (16)

• Lemma 4.0
• Lemma 4.0
• Lemma A.0
• proof
• Lemma A.0
• proof
• Lemma A.0
• proof
• Lemma A.0
• Corollary A.0