Amortized Sampling with Transferable Normalizing Flows

TL;DR

Prose introduces a large-scale, transferable all-atom normalizing flow trained on peptide MD data that enables zero-shot, uncorrelated sampling across varying sequence lengths up to eight residues. It combines TarFlow-style autoregressive flows with adaptive conditioning and chemistry-aware permutations to achieve cross-system transfer while retaining efficient likelihood evaluation, and uses SNIS-based inference with self-improvement and temperature-transfer capabilities. Across the ManyPeptidesMD dataset, Prose delivers state-of-the-art sampling performance, surpassing MD under similar compute and providing robust transfer to unseen systems and temperatures. The work is complemented by extensive ablations, multiple sampling strategies, and open-source resources to encourage further exploration of amortized sampling methods in molecular modeling.

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 finetuning 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 finetuning objectives.

Figures (12)

• Figure 1: Prose exceeds the quantitative performance of molecular dynamics on unseen peptide systems. Wasserstein-2 distances on energy, dihedral torus, and TICA projection with respect to reference molecular dynamics (5), for a (1) molecular dynamics baseline and Prose (with SNIS), at a range of energy evaluation (above) and GPU walltime budgets (below). Each value represents the mean over 30 unseen tetrapeptide systems. Prose outperforms the baseline with respect to energy evaluations on all metrics. Whilst comparable on $E\text{‑}\mathcal{W}_2$ for a given time budget, the baseline is significantly inferior on the $\mathcal{T}\text{‑}\mathcal{W}_2$ and $\mathrm{TICA}\text{‑}\mathcal{W}_2$ macrostructure metrics, highlighting long simulation periods were required to traverse the separated metastable states.
• Figure 2: All-atom block-wise autoregressive normalizing flow based on the TarFlow zhai_normalizing_2024. Peptides are encoded via atom types $A$, residue types $R$, residue position $P$, and sequence length $L$. Atom positions in 3D Cartesian coordinates define the system state. The embedding of the peptide is applied as conditioning to the coordinates such that Prose achieves transferability between systems. Within each block the sequence $z_t$ is permuted and passed to a transformer, defining an autoregressive affine update. In the backbone permutation the backbone $[N_i,C_{\alpha,i},C_i, O_i]_{i=1}^L$ of all residues (with associated hydrogens) is updated before any sidechains, providing additional diversity to the causal attention for global structure modeling.
• Figure 3: Prose accurately samples from the Boltzmann distributions of unseen octopepitde system. Empirical results for sampling from DGVAHALS peptide system, not present in training data. Energy histogram (left) for reference MD data, Prose proposal and Prose reweighted using SNIS, demonstrate fine-grained detail accuracy. TICA plots for MD (center) and SNIS-reweighted Prose (right) illustrate mode coverage.
• Figure 4: By drawing uncorrelated proposal samples, Prose achieves greater metastable state coverage than molecular dynamics for the same number of energy evaluations. TICA projection plots for unseen tetrapeptide system (RLMM). After $5 \cdot 10^9$ energy evaluations the reference molecular dynamics (left) has traversed four distinct metastable states, taken to be ground truth. However, with an energy evaluation budget of $10^6$ molecular dynamics explores only a single metastable state (center), highlighting the limitations of simulation-based sampling methods for mode exploration. Prose with SNIS (right) samples all 4 states given the same budget of energy evaluations, indicating successful amortization of the mode exploration problem.
• Figure 5: Scaled prior greatly the ability of Prose to accurately reweight to arbitrary temperatures. Metrics for Prose on RLMM unseen tetrapeptide, targeting temperatures up to 800K. Naively applying SNIS to the target temperature leads to a rapid degradation in energy distribution, and to a lesser extent the dihedral angle distribution. Applying prior scaling (Prose SP) leads to a significant improvement in energy distribution at high temperatures and moderate improvement in dihedral angles. Notably, the TICA distribution improves at higher temperatures irrespective of scaled prior usage, although scaled prior remains more effective.