Table of Contents
Fetching ...

Swallowing the Bitter Pill: Simplified Scalable Conformer Generation

Yuyang Wang, Ahmed A. Elhag, Navdeep Jaitly, Joshua M. Susskind, Miguel Angel Bautista

TL;DR

This work reframes molecular conformer generation as learning a distribution over conformer fields on molecular graphs, using a diffusion model that denoises 3D coordinates directly. By encoding graphs with Laplacian eigenvectors and employing a scalable PerceiverIO-based score network, the method MCF achieves state-of-the-art results without imposing torsion-based or roto-translation inductive biases. The authors demonstrate that scaling the model size yields substantial generalization gains across QM9, DRUGS, and GEOM-XL, and show that equivariance can be traded for scale. The results suggest a promising, domain-agnostic path for applying diffusion models to complex scientific problems and potential extensions to larger biomolecules and conditional design tasks.

Abstract

We present a novel way to predict molecular conformers through a simple formulation that sidesteps many of the heuristics of prior works and achieves state of the art results by using the advantages of scale. By training a diffusion generative model directly on 3D atomic positions without making assumptions about the explicit structure of molecules (e.g. modeling torsional angles) we are able to radically simplify structure learning, and make it trivial to scale up the model sizes. This model, called Molecular Conformer Fields (MCF), works by parameterizing conformer structures as functions that map elements from a molecular graph directly to their 3D location in space. This formulation allows us to boil down the essence of structure prediction to learning a distribution over functions. Experimental results show that scaling up the model capacity leads to large gains in generalization performance without enforcing inductive biases like rotational equivariance. MCF represents an advance in extending diffusion models to handle complex scientific problems in a conceptually simple, scalable and effective manner.

Swallowing the Bitter Pill: Simplified Scalable Conformer Generation

TL;DR

This work reframes molecular conformer generation as learning a distribution over conformer fields on molecular graphs, using a diffusion model that denoises 3D coordinates directly. By encoding graphs with Laplacian eigenvectors and employing a scalable PerceiverIO-based score network, the method MCF achieves state-of-the-art results without imposing torsion-based or roto-translation inductive biases. The authors demonstrate that scaling the model size yields substantial generalization gains across QM9, DRUGS, and GEOM-XL, and show that equivariance can be traded for scale. The results suggest a promising, domain-agnostic path for applying diffusion models to complex scientific problems and potential extensions to larger biomolecules and conditional design tasks.

Abstract

We present a novel way to predict molecular conformers through a simple formulation that sidesteps many of the heuristics of prior works and achieves state of the art results by using the advantages of scale. By training a diffusion generative model directly on 3D atomic positions without making assumptions about the explicit structure of molecules (e.g. modeling torsional angles) we are able to radically simplify structure learning, and make it trivial to scale up the model sizes. This model, called Molecular Conformer Fields (MCF), works by parameterizing conformer structures as functions that map elements from a molecular graph directly to their 3D location in space. This formulation allows us to boil down the essence of structure prediction to learning a distribution over functions. Experimental results show that scaling up the model capacity leads to large gains in generalization performance without enforcing inductive biases like rotational equivariance. MCF represents an advance in extending diffusion models to handle complex scientific problems in a conceptually simple, scalable and effective manner.
Paper Structure (28 sections, 6 equations, 13 figures, 7 tables, 2 algorithms)

This paper contains 28 sections, 6 equations, 13 figures, 7 tables, 2 algorithms.

Figures (13)

  • Figure 1: Overview of proposed MCF. The structure of molecular graph is encoded through eigenvectors of Laplacian eigen-decomposition $\varphi(\mathcal{V})$ and atomic features $\mathcal{A}$. MCF directly operates on atom coordinates in 3D space and trains the diffusion model to denoise the function in 3D coordinates. The score network is developed with attention-based PerceiverIO architecture. Context pairs $\mathbf{C}_t$ attend to a latent array of learnable parameters via cross attention and the latent array goes through several self attention blocks. Finally, the query pairs $\mathbf{Q}_t$ cross-attend to the latent array to produce the final noise prediction $\hat{\epsilon}_q$ in 3D space.
  • Figure 2: Training
  • Figure 3: Sampling
  • Figure 4: (a) Recall coverage and (b) precision coverage as a function of the threshold distance. MCF outperforms Torsional Diff. across the full spectrum of thresholds. (c) Averaged AMR of recall and precision as a function of the number of atoms in molecules.
  • Figure 5: (a) Mean Coverage and (b) mean AMR of different rotation augmentation strategies on GEOM-QM9 when compared with training on original dataset.
  • ...and 8 more figures