Research map / Probabilistic and scientific modeling / Variational and generative sampling
Conditional Normalizing Flows for Active Learning of Coarse-Grained Molecular Representations
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.
Proposes active learning in a coarse-grained space using a conditional normalizing flow trained by energy to sample fine-grained coordinates, enabling PMF estimation without long all-atom MD, and shows large speedups on alanine dipeptide.
Problem
Efficiently learning the potential of mean force (PMF) for coarse-grained molecular simulations normally requires long all-atom MD trajectories, and normalizing flows trained by energy alone suffer from mode collapse.
Approach
Separate fine-grained and coarse-grained degrees of freedom: a normalizing flow conditioned on CG coordinates is trained by reverse KL using the energy, and the PMF is estimated by reweighting flow samples; an ensemble of PMF models drives an active learning loop that samples high-error CG points via Metropolis MC, avoiding constrained MD.
Key results
- On alanine dipeptide, AL workflow reaches forward KLD 9.29e-4 with 3.35e7 potential energy evaluations vs 2e8 for Flow AIS Bootstrap (KLD 2.51e-3).
- Speedup over MD at fixed KLD 2.5e-3 is about 15.9 for AL and 216.2 for grid conditioning, versus 4.5 for Flow AIS Bootstrap.
- Grid conditioning (short) reaches KLD 1.68e-3 with 4.85e6 potential energy evaluations, while reverse-KLD flow shows mode collapse (reverse ESS 54%, KLD 3.15).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use active learning with the conditional normalizing flow instead of Flow AIS Bootstrap | Alanine dipeptide PMF estimation · Alanine dipeptide | Forward KLD PMF: 2.51(39)e-3 -> 9.29(24)e-4 | Helped |
| Use grid conditioning instead of Flow AIS Bootstrap | Alanine dipeptide PMF estimation · Alanine dipeptide | Forward KLD PMF: 2.51(39)e-3 -> 6.32(28)e-4 | Helped |
| Use grid conditioning instead of AL | Proposed conditional-flow method · Alanine dipeptide | Forward KLD PMF: 9.29(24)e-4 -> 6.32(28)e-4 | Helped |
| Use grid conditioning instead of Flow AIS Bootstrap | Alanine dipeptide PMF estimation · Alanine dipeptide | Potential energy evaluations: 2e8 -> 2.40e7 | Helped |
| Use active learning instead of long MD | Alanine dipeptide PMF estimation · Alanine dipeptide | Forward KLD PMF: 2.32e-3 -> 9.29(24)e-4 | Helped |
| Use active learning instead of short MD | Alanine dipeptide PMF estimation · Alanine dipeptide | Forward KLD PMF: 1.87e-2 -> 9.29(24)e-4 | Helped |
| Use the conditional-flow active-learning workflow instead of all-atom MC | Müller-Brown PMF estimation · Müller-Brown potential | Forward KLD PMF: 1.02e-2 -> 2.04(23)e-4 | Helped |
| Use grid conditioning instead of AL | Proposed conditional-flow method · Alanine dipeptide | Reverse ESS / %: 43.83(3.81) -> 83.92(1.15) | Helped |
| Use grid conditioning instead of Flow AIS Bootstrap | Alanine dipeptide all-atom sampling · Alanine dipeptide | Reverse ESS / %: 92.8(1) -> 83.92(1.15) | Hurt |
| Remove the highest 5 loss values from each training batch (ablation) | Conditional normalizing flow energy training · Alanine dipeptide | Training stability | Helped |
| Replace active learning with grid conditioning (ablation) | Proposed conditional-flow method · Alanine dipeptide | Forward KLD PMF: 9.29(24)e-4 -> 6.32(28)e-4; short grid: 1.68(13)e-3 | Mixed |
| Use the alternative PMF expectation formulation instead of the paper's direct reweighting formulation (ablation) | Conditional-flow PMF estimation · Alanine dipeptide | PMF accuracy | Hurt |
Surprising findings
- Reweighting-based PMF from the conditional flow remains accurate in regions with no ground-truth MD data, so the reported KLD may underestimate accuracy.
- The alternative PMF formula used in prior work (Zhang et al.) performed worse than the proposed estimator in practice.
Limitations
- Evaluated only on low-dimensional CG spaces (2D for alanine dipeptide, 1D for Muller-Brown); grid sampling does not scale to higher dimensions.
- Reverse KLD training is mode-seeking, so it may only fit approximately unimodal conditional distributions, a concern for protein side-chain conditionals.
- Ground-truth reference data is unreliable in high-energy regions, so accuracy cannot be fully quantified.
Experimental setup
Datasets: alanine dipeptide replica-exchange MD test dataset (Stimper et al. 2022; Midgley et al. 2023) Muller-Brown potential (synthetic) 5 microsecond MD trajectory generated by the authors
Benchmarks: alanine dipeptide PMF (phi, psi) Muller-Brown potential PMF
Metrics: forward KL divergence of PMF all-atom log-likelihood reverse effective sample size speedup at fixed KLD potential energy evaluations
Baselines: Flow AIS Bootstrap Flow with reverse KLD MD (short) MD (long)
Builds on: Boltzmann Generators Flow AIS Bootstrap multiscale force-matching flow-matching coarse-graining rational-quadratic spline coupling flows
Closest papers by meaning
- Path Gradients after Flow Matching (NeurIPS 2025 · similarity 0.89)
- Temperature-Annealed Boltzmann Generators (ICML 2025 · similarity 0.88)
- Stein Kernelized Molecular Dynamics for Active Learning of Interatomic Potentials (NeurIPS 2026 · similarity 0.87)
- Scalable Equilibrium Sampling with Sequential Boltzmann Generators (ICML 2025 · similarity 0.87)
- Generative Modeling of Molecular Dynamics Trajectories (NeurIPS 2024 · similarity 0.86)
- BoltzNCE: Learning likelihoods for Boltzmann Generation with Stochastic Interpolants and Noise Contrastive Estimation (NeurIPS 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.