Amortized Sampling with Transferable Normalizing Flows

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Tan, Charlie B., Hassan, Majdi, Klein, Leon, Syed, Saifuddin, Beaini, Dominique, Bronstein, Michael M., Tong, Alexander, Neklyudov, Kirill
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866917248425263104
author Tan, Charlie B.
Hassan, Majdi
Klein, Leon
Syed, Saifuddin
Beaini, Dominique
Bronstein, Michael M.
Tong, Alexander
Neklyudov, Kirill
author_facet Tan, Charlie B.
Hassan, Majdi
Klein, Leon
Syed, Saifuddin
Beaini, Dominique
Bronstein, Michael M.
Tong, Alexander
Neklyudov, Kirill
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2508_18175
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Amortized Sampling with Transferable Normalizing Flows
Tan, Charlie B.
Hassan, Majdi
Klein, Leon
Syed, Saifuddin
Beaini, Dominique
Bronstein, Michael M.
Tong, Alexander
Neklyudov, Kirill
Machine Learning
Artificial Intelligence
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.
title Amortized Sampling with Transferable Normalizing Flows
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2508.18175