Research map / Probabilistic and scientific modeling / Variational and generative sampling
Amortized Sampling with Transferable Normalizing Flows
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces PROSE, a 285M-parameter all-atom transferable TarFlow-based normalizing flow trained on ManyPeptidesMD to draw zero-shot proposals for unseen peptides up to 8 residues, used with importance sampling and self-improvement fine-tuning.
Problem
Amortized equilibrium sampling of molecular (peptide) conformations from the Boltzmann distribution that transfers across systems of different sequence length and composition, with efficient likelihood evaluation.
Approach
Extends the TarFlow autoregressive transformer flow with variable-length masking, adaptive system conditioning on atom/residue types and length, and chemistry-aware backbone-first permutations, trained on a large peptide MD corpus so that self-normalized importance sampling transfers zero-shot to unseen peptides; self-improvement via resampled likelihood fine-tuning and scaled-prior temperature tr
Key results
- On 8-residue peptides with a 10^4 energy-evaluation budget, PROSE with SNIS achieves E-W2 of 2.456 and TICA-W2 of 0.988 vs. MD-style baselines such as BioEmu (E-W2 193.873 at 8AA) and TimeWarp (only on 2AA/4AA), and ECNF++ (E-W2 10.032 on 4AA).
- Ablation shows removing backbone-first permutation or adaptive transition worsens 4AA E-W2 from 0.646 to 0.816 and 0.880 respectively.
- With a 10^6 energy budget, SNIS is competitive with SMC variants, and self-improvement improves 4AA E-W2 to 0.568 (best in table).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Prose instead of TarFlow as the SNIS proposal | SNIS · 2AA | E-W2: 0.452 -> 0.371 | Helped |
| Use Prose instead of TarFlow as the SNIS proposal | SNIS · 2AA | T-W2: 0.193 -> 0.210 | Hurt |
| Use Prose instead of TarFlow as the SNIS proposal | SNIS · 4AA | E-W2: 1.260 -> 0.932 | Helped |
| Use Prose instead of TarFlow as the SNIS proposal | SNIS · 8AA | TICA-W2: 1.087 -> 0.988 | Helped |
| Add backbone-first permutation (ablation) | Prose · 4AA | ESS: 0.051 -> 0.071 | Helped |
| Add backbone-first permutation (ablation) | Prose · 8AA | TICA-W2: 1.044 -> 0.960 | Helped |
| Add transition block (ablation) | Prose · 2AA | T-W2: 0.282 -> 0.177 | Helped |
| Add transition block (ablation) | Prose · 4AA | E-W2: 0.880 -> 0.646 | Helped |
| Use self-improvement with SNIS instead of SNIS alone | Prose sampling · 4AA | E-W2: 0.665 -> 0.568 | Helped |
| Use self-improvement with SNIS instead of SNIS alone | Prose sampling · 8AA | T-W2: 2.012 -> 2.071 | Hurt |
| Use discrete SMC instead of SNIS | Prose sampling · 8AA | E-W2: 9.386 -> 7.672 | Helped |
| Use discrete SMC instead of SNIS | Prose sampling · 8AA | T-W2: 2.012 -> 2.524 | Hurt |
Surprising findings
- Given a sufficiently strong proposal, plain SNIS without tuning is competitive with continuous and discrete SMC.
- Scaled-prior temperature transfer works for a non-volume-preserving flow despite the Jacobian term being nonzero.
Limitations
- Relies on the assumption that target systems lie in a structured energy-function space (peptides); not tested on more diverse chemistry such as OMol25.
- Scaled-prior temperature transfer works for higher temperatures but precise transfer to lower temperatures needs further development.
- ECNF++ could only be trained on sequences up to length 4 due to prohibitive likelihood cost.
Experimental setup
Datasets: ManyPeptidesMD
Benchmarks: ManyPeptidesMD evaluation set (2, 4, 8 residue peptides)
Metrics: ESS E-W2 T-W2 TICA-W2 TICA-k-JSD T-k-JSD
Baselines: TimeWarp BioEmu UniSim ECNF ECNF++ TarFlow MD baseline SMC continuous SMC discrete
Builds on: TarFlow Boltzmann generators Sequential Monte Carlo Transferable Boltzmann generator SBG
Closest papers by meaning
- JAMUN: Bridging Smoothed Molecular Dynamics and Score-Based Learning for Conformational Ensemble Generation (NeurIPS 2025 · similarity 0.88)
- Scalable Equilibrium Sampling with Sequential Boltzmann Generators (ICML 2025 · similarity 0.88)
- Consistent Sampling and Simulation: Molecular Dynamics with Energy-Based Diffusion Models (NeurIPS 2025 · similarity 0.87)
- AlphaFold Meets Flow Matching for Generating Protein Ensembles (ICML 2024 · similarity 0.87)
- Generative Modeling of Molecular Dynamics Trajectories (NeurIPS 2024 · similarity 0.87)
- Transition Path Sampling with Improved Off-Policy Training of Diffusion Path Samplers (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.