Amortized Sampling with Transferable Normalizing Flows

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.

💡 Research Summary

The paper tackles the long‑standing problem of efficiently sampling equilibrium molecular configurations, which is essential for tasks such as protein folding, ligand binding, and crystal structure prediction. Classical approaches—Molecular Dynamics (MD) and Markov Chain Monte Carlo (MCMC)—are inherently non‑amortized: each new system requires a full, costly simulation. Recent generative‑model based Boltzmann Generators (BGs) have shown that a one‑time training phase can produce a proposal distribution that, after self‑normalized importance sampling (SNIS), yields unbiased estimates of thermodynamic quantities. However, existing BGs and the more recent Transferable Boltzmann Generator (TBG) only generalize across very small chemical variations (e.g., dipeptides) and suffer from heavy inference costs due to continuous normalizing flow (CNF) integration.

The authors introduce PROSE (Probabilistic REgularized Sampling Engine), a 285 million‑parameter autoregressive normalizing flow built on the TarFlow architecture. The key innovations enabling transferability are:

1. Variable‑length handling – By masking padded tokens and excluding them from Jacobian‑determinant calculations, the model can be trained on peptides of differing lengths (2–8 residues) while still using a fixed‑dimensional flow.
2. Sinusoidal positional embeddings – Replacing learned position embeddings with periodic sin/cos embeddings improves extrapolation to longer sequences.
3. System‑conditional encoding – Atom type, residue type, and residue index are embedded and injected through multi‑head attention throughout the transformer layers, allowing the flow to condition on arbitrary amino‑acid sequences and temperatures.
4. Maximum‑likelihood training – The flow is trained to maximize log‑likelihood of MD trajectory frames, i.e., maxθ Eₛ Eₓ∼p(x|s) log qθ(x), which is computationally cheaper and more stable than reverse‑KL objectives used in TBG.

A new dataset, ManyPeptidesMD, is compiled: 21 700 distinct peptide sequences (2–8 residues) each simulated for 200 ns, totaling 4.3 ms of MD data. This large, diverse corpus provides the statistical power needed for the model to learn a universal representation of peptide conformational space.

After training, PROSE can generate zero‑shot, uncorrelated samples for any unseen peptide within the length range. These raw samples are re‑weighted using SNIS: weights wᵢ = p(xᵢ)/qθ(xᵢ) where p is the Boltzmann density. Because SNIS is unbiased in the limit of many samples, no additional MCMC tuning is required. The authors also evaluate PROSE as a proposal distribution for Sequential Monte Carlo (SMC), showing that the combination further improves effective sample size and reduces variance of free‑energy estimates.

Empirical results are extensive. On 30 unseen tetrapeptides, PROSE (with SNIS) outperforms a 1 µs MD baseline across three Wasserstein‑2 metrics: energy distribution, dihedral‑torus distribution, and TICA‑projected macro‑structure. For a given GPU‑wall‑time budget, PROSE achieves 4–5× lower error than MD, and its sampling speed is ~4 × 10³ times faster than the TBG CNF approach (which needs ~4 GPU‑days for 30 k samples). Moreover, a simple fine‑tuning step—training the flow on a few thousand SNIS‑weighted samples from a new peptide using a reverse‑KL loss—rapidly adapts the model to novel systems, demonstrating practical transferability.

Limitations are acknowledged: the current work is restricted to short peptides (≤8 residues) and to a single solvent/temperature regime. Scaling to larger proteins, diverse environments, and incorporating dynamic adaptation (e.g., reinforcement‑learning‑based proposal updates) are identified as future directions.

In summary, PROSE establishes a scalable, transferable, and highly efficient amortized sampler for peptide conformations. By marrying a large‑scale autoregressive flow with careful conditioning and masking, and by leveraging SNIS and SMC for unbiased correction, the method delivers orders‑of‑magnitude speed‑ups over traditional MD while maintaining or improving thermodynamic accuracy. This work represents a significant step toward making learned samplers a practical tool in computational chemistry and statistical physics.