Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Conditional Synthesis of 3D Molecules with Time Correction Sampler

Hojung Jung, Youngrok Park, Laura Schmid, Jaehyeong Jo, Dongkyu Lee, Bongsang Kim, Se-Young Yun, Jinwoo Shin

TL;DR

Time-Aware Conditional Synthesis (TACS) is presented, a novel approach to conditional generation on diffusion models that integrates adaptively controlled plug-and-play"online"guidance into a diffusion model, driving samples toward the desired properties while maintaining validity and stability.

Abstract

Diffusion models have demonstrated remarkable success in various domains, including molecular generation. However, conditional molecular generation remains a fundamental challenge due to an intrinsic trade-off between targeting specific chemical properties and generating meaningful samples from the data distribution. In this work, we present Time-Aware Conditional Synthesis (TACS), a novel approach to conditional generation on diffusion models. It integrates adaptively controlled plug-and-play "online" guidance into a diffusion model, driving samples toward the desired properties while maintaining validity and stability. A key component of our algorithm is our new type of diffusion sampler, Time Correction Sampler (TCS), which is used to control guidance and ensure that the generated molecules remain on the correct manifold at each reverse step of the diffusion process at the same time. Our proposed method demonstrates significant performance in conditional 3D molecular generation and offers a promising approach towards inverse molecular design, potentially facilitating advancements in drug discovery, materials science, and other related fields.

Conditional Synthesis of 3D Molecules with Time Correction Sampler

TL;DR

Time-Aware Conditional Synthesis (TACS) is presented, a novel approach to conditional generation on diffusion models that integrates adaptively controlled plug-and-play"online"guidance into a diffusion model, driving samples toward the desired properties while maintaining validity and stability.

Abstract

Diffusion models have demonstrated remarkable success in various domains, including molecular generation. However, conditional molecular generation remains a fundamental challenge due to an intrinsic trade-off between targeting specific chemical properties and generating meaningful samples from the data distribution. In this work, we present Time-Aware Conditional Synthesis (TACS), a novel approach to conditional generation on diffusion models. It integrates adaptively controlled plug-and-play "online" guidance into a diffusion model, driving samples toward the desired properties while maintaining validity and stability. A key component of our algorithm is our new type of diffusion sampler, Time Correction Sampler (TCS), which is used to control guidance and ensure that the generated molecules remain on the correct manifold at each reverse step of the diffusion process at the same time. Our proposed method demonstrates significant performance in conditional 3D molecular generation and offers a promising approach towards inverse molecular design, potentially facilitating advancements in drug discovery, materials science, and other related fields.

Paper Structure

This paper contains 78 sections, 2 theorems, 25 equations, 8 figures, 15 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Works
  3. Diffusion models
  4. Conditional molecular generation
  5. Preliminaries
  6. Diffusion models
  7. Online guidance
  8. 3D representations of molecules
  9. Time-Aware Conditional Synthesis
  10. Time Correction Sampler
  11. Time Predictor
  12. Time correction
  13. Unified guidance with time correction
  14. Time clipping
  15. Experiments
  16. ...and 63 more sections

Key Result

Proposition 1

Let $f: \mathbb{R}^{N \times 3} \to \mathbb{R}^{N \times d}$ be an E(3)-equivariant function and $g: \mathbb{R}^{N \times d} \to \mathbb{R}^d$ be a permutation-invariant function. Then, the composition $h = g \circ f: \mathbb{R}^{N \times 3} \to \mathbb{R}^d$ is invariant to E(3) transformations, i.

Figures (8)

  • Figure 1: (a) Overview of Time-Aware Conditional Synthesis (TACS). TACS helps generate high-quality samples that match target condition while following basic properties of the molecules. At each timestep $t$, online guidance is applied to push $x_t$ towards the desired condition. Time Predictor finds the desired timestep $t_p$ for $x_t$ after applying the guidance. Using predicted timestep $t_p$, Tweedie's formula is used to predict the clean molecule $\hat{x}^t_0$. Finally, forward process $q(x_{t-1}|x_0)$ is applied to proceed to the next denoising step of $t-1$. (b) Motivation for TACS. Applying online guidance ($\vec{g}$, purple) can shift the generated samples away from the correct data manifold corresponding to the current timestep. This undesirable deviation (red) can be avoided by using time correction to first measure the deviated timestep $t'$, then adjusting the guidance to get corrected guidance vector ($\vec{g*}$, green), which keeps the generated samples stay on the correct data manifold.
  • Figure 2: Synthetic experiment on $H_3^+$ dataset. TACS is robust in generating samples that (a) match the desired condition and (b) stick to the original data distribution.
  • Figure 2: Quantitative results showing similarity and stability of generated molecules compared to target structures.
  • Figure 3: (a) Comparative analysis of molecule stability and MAE across five distinct methods, highlighting the enhanced stability of our approach at comparable MAE levels. (b) Performance of the time predictor for train and test set on QM9.
  • Figure 4: Conditional generation of target property $\alpha$ on QM9. Visualization of molecules generated by TCS (top), online guidance (middle), and TACS (bottom).
  • ...and 3 more figures

Theorems & Definitions (6)

  • Definition 1: E(3) Equivariance
  • Definition 2: Permutation Invariance
  • Proposition 1
  • proof
  • Theorem 1: Time Predictor Equivariance
  • proof