Amortized Sampling with Transferable Normalizing Flows

TL;DR

This paper introduces Prose, a large-scale transferable normalizing flow model that enables efficient, zero-shot sampling of molecular conformations across different peptide systems, significantly advancing amortized sampling methods.

Contribution

The paper presents Prose, a 285M parameter transferable normalizing flow trained on peptide MD trajectories, demonstrating zero-shot transferability and competitive sampling performance.

Findings

Prose achieves zero-shot transfer across peptide sequence lengths.

Prose enables efficient likelihood evaluation for diverse systems.

Fine-tuning with importance sampling improves sampling quality.

Abstract

Efficient equilibrium sampling of molecular conformations remains a core challenge in computational chemistry and statistical inference. Classical approaches such as molecular dynamics or Markov chain Monte Carlo inherently lack amortization; the computational cost of sampling must be paid in full for each system of interest. The widespread success of generative models has inspired interest towards overcoming this limitation through learning sampling algorithms. Despite performing competitively with conventional methods when trained on a single system, learned samplers have so far demonstrated limited ability to transfer across systems. We demonstrate that deep learning enables the design of scalable and transferable samplers by introducing Prose, a 285 million parameter all-atom transferable normalizing flow trained on a corpus of peptide molecular dynamics trajectories up to 8 residues in length. Prose draws zero-shot uncorrelated proposal samples for arbitrary peptide systems, achieving the previously intractable transferability across sequence length, whilst retaining the efficient likelihood evaluation of normalizing flows. Through extensive empirical evaluation we demonstrate the efficacy of Prose as a proposal for a variety of sampling algorithms, finding a simple importance sampling-based fine-tuning procedure to achieve competitive performance to established methods such as sequential Monte Carlo. We open-source the Prose codebase, model weights, and training dataset, to further stimulate research into amortized sampling methods and objectives.