Diffusion Twigs with Loop Guidance for Conditional Graph Generation

TL;DR

A novel score-based diffusion framework named Twigs is introduced that incorporates multiple co-evolving flows for enriching conditional generation tasks, underscoring the potential of Twigs in challenging generative tasks such as inverse molecular design and molecular optimization.

Abstract

We introduce a novel score-based diffusion framework named Twigs that incorporates multiple co-evolving flows for enriching conditional generation tasks. Specifically, a central or trunk diffusion process is associated with a primary variable (e.g., graph structure), and additional offshoot or stem processes are dedicated to dependent variables (e.g., graph properties or labels). A new strategy, which we call loop guidance, effectively orchestrates the flow of information between the trunk and the stem processes during sampling. This approach allows us to uncover intricate interactions and dependencies, and unlock new generative capabilities. We provide extensive experiments to demonstrate strong performance gains of the proposed method over contemporary baselines in the context of conditional graph generation, underscoring the potential of Twigs in challenging generative tasks such as inverse molecular design and molecular optimization.

Figures (5)

• Figure 1: Overview of the proposed method (Twigs). We define two types of diffusion processes: (1) multiple Stem processes ($s_{\phi_i}$), which unravel the interactions between graph structure and single properties, and (2) the Trunk process, which orchestrates the combination of the graph structure score from $s_\theta$ with the stem process contributions from $s_{\phi_i}$. During the forward process, the structure $\mathbf{y}_s$ and the properties $\{\mathbf{y}_i\}_k$ co-evolve toward noise. In each step of the reverse process, the structure is first denoised and subsequently used to denoise the properties (indicated by the green-dashed line). Such de-noised properties are then utilized, in turn, to further denoise the structure (red line), in a process that resembles a guidance loop.
• Figure 2: First row: Samples by Twigs for 3D molecules conditioned on single properties on QM9. Second row: KDE and KL divergence results between target and predicted properties.
• Figure 3: Samples of multiple-property conditional molecules by Twigs ($C_v$ and $\mu$) for QM9.
• Figure 4: Molecules generated by Twigs from ZINC250k conditioned on fa7 (top), parp1 (bottom).
• Figure 5: Visualization of Community-small and Enzymes datasets. First and second rows: samples generated by Twigs. Third and fourth rows: KDE plots and corresponding KL divergence values.