A survey of probabilistic generative frameworks for molecular simulations

TL;DR

A taxonomy of probabilistic generative frameworks and numerical results may guide model selection for a wide range of molecular tasks, with no one framework being the best for all purposes.

Abstract

Generative artificial intelligence is now a widely used tool in molecular science. Despite the popularity of probabilistic generative models, numerical experiments benchmarking their performance on molecular data are lacking. In this work, we introduce and explain several classes of generative models, broadly sorted into two categories: flow-based models and diffusion models. We select three representative models: Neural Spline Flows, Conditional Flow Matching, and Denoising Diffusion Probabilistic Models, and examine their accuracy, computational cost, and generation speed across datasets with tunable dimensionality, complexity, and modal asymmetry. Our findings are varied, with no one framework being the best for all purposes. In a nutshell, (i) Neural Spline Flows do best at capturing mode asymmetry present in low-dimensional data, (ii) Conditional Flow Matching outperforms other models for high-dimensional data with low complexity, and (iii) Denoising Diffusion Probabilistic Models appears the best for low-dimensional data with high complexity. Our datasets include a Gaussian mixture model and the dihedral torsion angle distribution of the Aib\textsubscript{9} peptide, generated via a molecular dynamics simulation. We hope our taxonomy of probabilistic generative frameworks and numerical results may guide model selection for a wide range of molecular tasks.

Figures (5)

• Figure 1: Model accuracy results for the 4-modal Gaussian mixture dataset. a) KL divergence ($D_{KL}$) comparison with analytical benchmark as a function of dimensionality. b) KL divergence comparison as a function of training dataset size at a fixed dimensionality of 50.
• Figure 2: Free energy difference estimation accuracy on asymmetric bimodal distributions. A free energy difference of zero represents two equal Gaussian modes, while higher free energy indicates a higher level of asymmetry. $r^2$ is computed with the residuals of each model from the plotted line indicating training free energy difference. We impose a free energy cutoff of 0.0374 kJ/mol and note the data dimensionality is fixed at 50.
• Figure 3: Speed and model size results for the 4-modal Gaussian mixture dataset. a) Single sample generation time comparison as a function of dimensionality. b) Model capacity comparison measured by total number of learnable parameters as a function of dimensionality.
• Figure 4: Model accuracy results and generated data for the Aib9 peptide. a) KLD performance comparison as a function of residue index for the complete Aib9 torsion angle dataset. b) KLD performance comparison for the three models and Gaussian baseline fit as a function of training dataset size for the Aib9 torsion angle data distribution at residue 5. c) d) e) f) $\{ \Phi,\Psi \}$ show free energy contour plots for torsion angle distributions at residue 1 for training data and generated data for NS, CFM, and DDPM, respectively.
• Figure 5: Hyperparameter tuning for each model. a) Tuning of the 'layers' parameter of NS. b) c) Tuning of the 'model dimension' parameter of CFM and DDPM respectively.