Research map / Probabilistic and scientific modeling / Variational and generative sampling
Temperature-Annealed Boltzmann Generators
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 temperature-annealed Boltzmann generators: reverse-KLD training of a normalizing flow at high temperature (1200 K), then iterative reweighting-based forward-KLD annealing down to 300 K, evaluated on alanine di/tetra/hexapeptide.
Problem
Variational sampling of Boltzmann distributions of molecules with normalizing flows suffers from mode collapse under reverse KLD, so flows fail to cover the full configurational space (metastable states).
Approach
Reverse-KLD training is mode-collapse-free at high temperature where energy barriers are lower; the flow is then iteratively annealed to the target temperature by importance-reweighting buffered samples and performing forward-KLD training at each intermediate temperature, followed by a final fine-tuning step.
Key results
- On alanine hexapeptide TA-BG achieves NLL -501.511 vs FAB -501.268 and Ram KLD 8.63e-3 vs 2.09e-2, using 3.08e8 target energy evaluations vs FAB's 4.2e8.
- On alanine dipeptide TA-BG uses 7.56e7 energy evaluations (about one third of FAB's 2.13e8) with best NLL -213.665 and ESS 95.61%.
- Reverse KLD at 300 K collapses for the hexapeptide (ESS 22.22%, Ram KLD 0.541), while TA-BG covers all metastable states.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use TA-BG instead of FAB | Boltzmann generator · Alanine dipeptide | NLL: -213.653(0.000)e0 -> -213.665(0.002)e0 | Helped |
| Use TA-BG instead of FAB | Boltzmann generator · Alanine tetrapeptide | NLL: -330.100(0.002)e0 -> -330.114(0.004)e0 | Helped |
| Use TA-BG instead of FAB | Boltzmann generator · Alanine hexapeptide | Ram KLD: 2.09(0.02)e-2 -> 8.63(0.66)e-3 | Helped |
| Use TA-BG instead of FAB | Boltzmann generator · Alanine hexapeptide | PE EVALS: 4.2e8 -> 3.08e8 | Helped |
| Add final fine-tuning iteration (ablation) | TA-BG · Alanine dipeptide | NLL: -213.666 -> -213.667 | Helped |
| Add final fine-tuning iteration (ablation) | TA-BG · Alanine tetrapeptide | ESS: 57.57 -> 62.57 | Helped |
| Add final fine-tuning iteration (ablation) | TA-BG · Alanine hexapeptide | ESS: 13.37 -> 15.38 | Helped |
| Add intermediate fine-tuning iterations after each annealing iteration (ablation) | TA-BG · Alanine hexapeptide | ESS | Helped |
| Use direct importance sampling from 1200 to 300 kelvin rather than multistep annealing (ablation) | TA-BG · Molecular systems | ESS | Hurt |
| Use a linear rather than geometric temperature schedule (ablation) | TA-BG · Alanine dipeptide and hexapeptide | Buffer ESS | Hurt |
| Train a temperature-conditioned flow continuously over the temperature range rather than iterative buffered annealing (ablation) | TA-BG workflow · Molecular systems | Accuracy | Hurt |
| Train with reverse KLD directly at 300 kelvin rather than high-temperature pretraining and annealing (ablation) | Normalizing flow · Alanine tetrapeptide and hexapeptide | Ram KLD | Hurt |
Surprising findings
- Reverse KLD training at 300 K for alanine dipeptide yields only partial mode collapse, performing well despite low-probability transition regions.
- Direct importance sampling from 1200 K to 300 K gives very poor overlap, so a multistep annealing is necessary.
Limitations
- Uses internal coordinates, which are not transferable between systems and ignore permutation invariance.
- Reverse ESS only measures support of the flow and does not capture mode collapse; forward ESS deemed unreliable.
- Hexapeptide requires intermediate fine-tuning that increases target energy evaluations; results only on three small peptides with a force field.
Experimental setup
Datasets: MD simulations of alanine dipeptide, tetrapeptide, hexapeptide (AMBER force fields, OpenMM) REMD ground truth for tetrapeptide and hexapeptide
Benchmarks: alanine dipeptide alanine tetrapeptide alanine hexapeptide
Metrics: NLL reverse ESS Ramachandran KLD Ramachandran KLD with reweighting target energy evaluations
Baselines: forward KLD reverse KLD FAB
Builds on: Boltzmann generators FAB (flow annealed importance sampling bootstrap) rational-quadratic neural spline flows bgflow
Closest papers by meaning
- Scalable Equilibrium Sampling with Sequential Boltzmann Generators (ICML 2025 · similarity 0.91)
- Transferable Boltzmann Generators (NeurIPS 2024 · similarity 0.89)
- BoltzNCE: Learning likelihoods for Boltzmann Generation with Stochastic Interpolants and Noise Contrastive Estimation (NeurIPS 2025 · similarity 0.89)
- Conditional Normalizing Flows for Active Learning of Coarse-Grained Molecular Representations (ICML 2024 · similarity 0.88)
- Path Gradients after Flow Matching (NeurIPS 2025 · similarity 0.87)
- Fast training and sampling of Restricted Boltzmann Machines (ICLR 2025 · similarity 0.87)
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.