Table of Contents
Fetching ...

Accelerating the Generation of Molecular Conformations with Progressive Distillation of Equivariant Latent Diffusion Models

Romain Lacombe, Neal Vaidya

TL;DR

Diffusion-based generation of 3D molecular conformations is computationally expensive due to iterative inference. The authors extend GeoLDM by introducing Equivariant Latent Progressive Distillation (ELPD), performing progressive distillation in latent space with both DDIM and DDPM sampling. They systematically evaluate speed versus stability on QM9, reporting up to 7.5× sampling speed with minimal loss in molecular stability, and show stochastic DDPM-based distillation yields the best trade-offs. The results indicate strong potential for high-throughput in silico screening in computational biochemistry and drug discovery, with avenues for scaling to larger molecules and incorporating energy-based criteria.

Abstract

Recent advances in fast sampling methods for diffusion models have demonstrated significant potential to accelerate generation on image modalities. We apply these methods to 3-dimensional molecular conformations by building on the recently introduced GeoLDM equivariant latent diffusion model (Xu et al., 2023). We evaluate trade-offs between speed gains and quality loss, as measured by molecular conformation structural stability. We introduce Equivariant Latent Progressive Distillation, a fast sampling algorithm that preserves geometric equivariance and accelerates generation from latent diffusion models. Our experiments demonstrate up to 7.5x gains in sampling speed with limited degradation in molecular stability. These results suggest this accelerated sampling method has strong potential for high-throughput in silico molecular conformations screening in computational biochemistry, drug discovery, and life sciences applications.

Accelerating the Generation of Molecular Conformations with Progressive Distillation of Equivariant Latent Diffusion Models

TL;DR

Diffusion-based generation of 3D molecular conformations is computationally expensive due to iterative inference. The authors extend GeoLDM by introducing Equivariant Latent Progressive Distillation (ELPD), performing progressive distillation in latent space with both DDIM and DDPM sampling. They systematically evaluate speed versus stability on QM9, reporting up to 7.5× sampling speed with minimal loss in molecular stability, and show stochastic DDPM-based distillation yields the best trade-offs. The results indicate strong potential for high-throughput in silico screening in computational biochemistry and drug discovery, with avenues for scaling to larger molecules and incorporating energy-based criteria.

Abstract

Recent advances in fast sampling methods for diffusion models have demonstrated significant potential to accelerate generation on image modalities. We apply these methods to 3-dimensional molecular conformations by building on the recently introduced GeoLDM equivariant latent diffusion model (Xu et al., 2023). We evaluate trade-offs between speed gains and quality loss, as measured by molecular conformation structural stability. We introduce Equivariant Latent Progressive Distillation, a fast sampling algorithm that preserves geometric equivariance and accelerates generation from latent diffusion models. Our experiments demonstrate up to 7.5x gains in sampling speed with limited degradation in molecular stability. These results suggest this accelerated sampling method has strong potential for high-throughput in silico molecular conformations screening in computational biochemistry, drug discovery, and life sciences applications.
Paper Structure (19 sections, 2 equations, 4 figures, 2 tables, 1 algorithm)

This paper contains 19 sections, 2 equations, 4 figures, 2 tables, 1 algorithm.

Figures (4)

  • Figure 1: Geometric latent diffusion model. In the GeoLDM model xu2023geometric, an equivariant denoising process is trained to generate point clouds in latent space by inverting a diffusion process. A decoder then transforms point clouds into 3-dimensional molecular conformations in atomic space. Image adapted from xu2023geometric and neurips23diffusionworkshop.
  • Figure 2: Equivariant latent progressive distillation: the denoising process in latent space is iteratively distilled by teaching a student model to sample two steps at a time from the teacher model. The algorithm initializes from the original denoising model, and progresses by successive halvings.
  • Figure 3: Molecular stability of generated samples for various diffusion steps. GeoLDM-DDIM is the original model sampled for a varying number of DDIM steps. GeoLDM-ProgDist-DDIM and GeoLDM-ProgDist-DDPM are sequences of progressively distilled models, using DDIM and DDPM samplers at train time. Note: $x$ axis is logarithmic; distillation progresses from right to left.
  • Figure 4: Examples of conformation generations, sampled from: (a, b) original 1000-step GeoLDM model; (c, d) 125-step model distilled with DDPM sampler after 3 halvings; (e, f) 16-step model distilled with DDPM sampler after 6 halvings.