Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Breaking the Bottlenecks: Scalable Diffusion Models for 3D Molecular Generation

Adrita Das, Peiran Jiang, Dantong Zhu, Barnabas Poczos, Jose Lugo-Martinez

TL;DR

This work reframes Directly Denoising Diffusion Models for 3D molecular generation within the Reverse Transition Kernel (RTK) framework, unifying deterministic and stochastic diffusion under a probabilistic formalism. It shows that deterministic denoising can be interpreted as solving a small number of well-conditioned subproblems via a transport-map between noisy and clean samples, yielding faster, more stable inference with SE(3) equivariance. The proposed SE(3)-equivariant state-space diffusion framework (GraphGPS-based) uses input-dependent node selection and decouples local and global information to scale to large molecules, achieving strong fidelity on GEOM-Drugs and the GEOM-LongRange benchmark. Empirically, RTK-guided denoising outperforms stochastic baselines in convergence speed and structural validity while preserving chemical validity, with code and data publicly available for reproducibility and broader adoption.

Abstract

Diffusion models have emerged as a powerful class of generative models for molecular design, capable of capturing complex structural distributions and achieving high fidelity in 3D molecule generation. However, their widespread use remains constrained by long sampling trajectories, stochastic variance in the reverse process, and limited structural awareness in denoising dynamics. The Directly Denoising Diffusion Model (DDDM) mitigates these inefficiencies by replacing stochastic reverse MCMC updates with deterministic denoising step, substantially reducing inference time. Yet, the theoretical underpinnings of such deterministic updates have remained opaque. In this work, we provide a principled reinterpretation of DDDM through the lens of the Reverse Transition Kernel (RTK) framework by Huang et al. 2024, unifying deterministic and stochastic diffusion under a shared probabilistic formalism. By expressing the DDDM reverse process as an approximate kernel operator, we show that the direct denoising process implicitly optimizes a structured transport map between noisy and clean samples. This perspective elucidates why deterministic denoising achieves efficient inference. Beyond theoretical clarity, this reframing resolves several long-standing bottlenecks in molecular diffusion. The RTK view ensures numerical stability by enforcing well-conditioned reverse kernels, improves sample consistency by eliminating stochastic variance, and enables scalable and symmetry-preserving denoisers that respect SE(3) equivariance. Empirically, we demonstrate that RTK-guided deterministic denoising achieves faster convergence and higher structural fidelity than stochastic diffusion models, while preserving chemical validity across GEOM-DRUGS dataset. Code, models, and datasets are publicly available in our project repository.

Breaking the Bottlenecks: Scalable Diffusion Models for 3D Molecular Generation

TL;DR

This work reframes Directly Denoising Diffusion Models for 3D molecular generation within the Reverse Transition Kernel (RTK) framework, unifying deterministic and stochastic diffusion under a probabilistic formalism. It shows that deterministic denoising can be interpreted as solving a small number of well-conditioned subproblems via a transport-map between noisy and clean samples, yielding faster, more stable inference with SE(3) equivariance. The proposed SE(3)-equivariant state-space diffusion framework (GraphGPS-based) uses input-dependent node selection and decouples local and global information to scale to large molecules, achieving strong fidelity on GEOM-Drugs and the GEOM-LongRange benchmark. Empirically, RTK-guided denoising outperforms stochastic baselines in convergence speed and structural validity while preserving chemical validity, with code and data publicly available for reproducibility and broader adoption.

Abstract

Diffusion models have emerged as a powerful class of generative models for molecular design, capable of capturing complex structural distributions and achieving high fidelity in 3D molecule generation. However, their widespread use remains constrained by long sampling trajectories, stochastic variance in the reverse process, and limited structural awareness in denoising dynamics. The Directly Denoising Diffusion Model (DDDM) mitigates these inefficiencies by replacing stochastic reverse MCMC updates with deterministic denoising step, substantially reducing inference time. Yet, the theoretical underpinnings of such deterministic updates have remained opaque. In this work, we provide a principled reinterpretation of DDDM through the lens of the Reverse Transition Kernel (RTK) framework by Huang et al. 2024, unifying deterministic and stochastic diffusion under a shared probabilistic formalism. By expressing the DDDM reverse process as an approximate kernel operator, we show that the direct denoising process implicitly optimizes a structured transport map between noisy and clean samples. This perspective elucidates why deterministic denoising achieves efficient inference. Beyond theoretical clarity, this reframing resolves several long-standing bottlenecks in molecular diffusion. The RTK view ensures numerical stability by enforcing well-conditioned reverse kernels, improves sample consistency by eliminating stochastic variance, and enables scalable and symmetry-preserving denoisers that respect SE(3) equivariance. Empirically, we demonstrate that RTK-guided deterministic denoising achieves faster convergence and higher structural fidelity than stochastic diffusion models, while preserving chemical validity across GEOM-DRUGS dataset. Code, models, and datasets are publicly available in our project repository.
Paper Structure (47 sections, 4 theorems, 80 equations, 2 figures, 9 tables, 2 algorithms)

This paper contains 47 sections, 4 theorems, 80 equations, 2 figures, 9 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Background and Related Work
  4. Notations and Preliminaries
  5. Diffusion Models
  6. DDPMs
  7. Implicit Models
  8. Remark (Implicit Forward Process).
  9. Definitions
  10. Diffusion as Approximate Sampling of Transition Kernels
  11. Problem Formulation
  12. Probability Flow ODE and Deterministic Reverse Diffusion
  13. Euler Discretization and Directly Denoising Diffusion Models
  14. Reverse Transition Kernel Perspective
  15. Proxy Reverse Kernels and Energy-Based Interpretation
  16. ...and 32 more sections

Key Result

Proposition 1

(Log-Concavity of Reverse Subproblem(s)). DDDM reformulates the reverse process as a deterministic iterative refinement procedure, where each iteration corresponds to solving a strongly log-concave reconstruction objective. Let, $p(z) \propto \exp(-g(z))$ denote the target distribution of a reverse Hence, if $(1 - L(z))^2 > R(z)B_{\mathrm{sq}}(z)$, the energy function $g$ is $m$-strongly convex w

Figures (2)

  • Figure 1: Overview of the diffusion generative process and reverse denoising dynamics. The top panel shows the forward noising trajectory from the data distribution $p_{\text{data}}$ to the noise prior $p_{\text{prior}}$, along with the corresponding reverse kernel approximation used during training. The bottom panel illustrates sampling via direct denoising from $x_T$ back to the clean data $x_0$.
  • Figure : Reverse Diffusion via RTKs

Theorems & Definitions (6)

  • Proposition 1
  • Proposition 2
  • Lemma 1
  • proof
  • Corollary 1: One-hidden-layer explicit bounds
  • proof