Research map / Probabilistic and scientific modeling / Variational and generative sampling
BoltzNCE: Learning likelihoods for Boltzmann Generation with Stochastic Interpolants and Noise Contrastive Estimation
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.
Trains an energy-based model on Boltzmann-emulator samples using a stochastic-interpolant-based combination of denoising score matching and InfoNCE, giving fast likelihoods that enable reweighting to Boltzmann statistics without costly Jacobian traces.
Problem
Boltzmann Generators need exact CNF likelihoods via costly Jacobian trace integrals, which do not scale to large molecular systems; sampling the Boltzmann distribution is also hard due to metastability.
Approach
Amortize CNF likelihoods by training a time-dependent energy-based model on samples from a stochastic-interpolant Boltzmann emulator, using a hybrid of denoising score matching and InfoNCE (negatives from nearby times), so a single EBM forward pass gives approximate log-likelihoods for reweighting.
Key results
- On 2D synthetic systems, combining score matching and InfoNCE reduced KL divergence to 0.0150 (8-mode GMM) and 0.1987 (checkerboard), versus 0.2199/0.8384 for score matching alone.
- On alanine dipeptide, BoltzNCE-VF gave free energy difference 4.08 ± 0.13 vs umbrella sampling 4.10 ± 0.26 with about 100x faster inference (0.09 h vs 9.37 h for ECNF).
- On 7 unseen dipeptides, BoltzNCE achieved 4.005 h inference vs 24.04 h for MD, with φ ΔF error 0.43 ± 0.21 compared with 0.18 ± 0.22 for MD.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add InfoNCE to score matching (ablation) | EBM training · 8-mode Gaussian mixture | KL divergence: 0.2199 -> 0.0150 | Helped |
| add InfoNCE to score matching (ablation) | EBM training · Checkerboard | KL divergence: 0.8384 -> 0.1987 | Helped |
| add score matching to InfoNCE (ablation) | EBM training · 8-mode Gaussian mixture | KL divergence: 0.2395 -> 0.0150 | Helped |
| add score matching to InfoNCE (ablation) | EBM training · Checkerboard | KL divergence: 3.8478 -> 0.1987 | Helped |
| replace exact Jacobian trace likelihoods with learned EBM likelihoods | Boltzmann generation on alanine dipeptide · Alanine dipeptide, biased | Inference time (h): 9.37 -> 0.09 | Helped |
| use learned EBM likelihoods instead of the Hutchinson estimator | Likelihood estimation · 10,000 samples | Spearman correlation: 0.87 -> 0.93 | Helped |
| use learned EBM likelihoods instead of the Hutchinson estimator | Likelihood estimation · 10,000 samples | Time: 8 min -> 0.5s | Helped |
| use BoltzNCE rather than MD | Boltzmann sampling on unseen dipeptides · 7 unseen dipeptides | Inference time (h): 24.04 -> 4.005 | Helped |
| use BoltzNCE rather than TBG-ECNF | Boltzmann sampling on unseen dipeptides · 7 unseen dipeptides | Inference time (h): 123.07 -> 4.005 | Helped |
| use BoltzNCE rather than TBG-ECNF | Boltzmann sampling on unseen dipeptides · 7 unseen dipeptides | φ ΔF error: 0.13 ± 0.10 -> 0.43 ± 0.21 | Hurt |
| use GVP Endpoint rather than ECNF | Alanine dipeptide Boltzmann emulator · Alanine dipeptide, unbiased emulator | Energy W2: 6.22 ± 0.12 -> 3.11 ± 0.70 | Helped |
| use GVP Endpoint rather than ECNF | Alanine dipeptide Boltzmann emulator · Alanine dipeptide, unbiased emulator | NLL: $-$125.53 ± 0.10 -> $-$92.04 ± 3.24 | Hurt |
Surprising findings
- The EBM likelihoods correlated better with exact likelihoods (Spearman 0.93) than the Hutchinson trace estimator (0.87-0.88), and increasing Hutchinson estimator calls did not improve correlation.
- The GVP endpoint emulator had the best energy W2 among emulators but the worst NLL, because its divergent vector field corrupts Jacobian-based likelihoods.
Limitations
- Evaluated only on dipeptide molecular systems; accuracy in higher-dimensional settings is untested.
- Generalization to unseen dipeptides performs worse, with larger energy W2 error (1.08) than MD and TBG-ECNF baselines.
- Training the EBM requires score matching through gradients, increasing compute; EBM likelihoods are approximate and only defined up to a constant.
Experimental setup
Datasets: alanine dipeptide MD (Amber ff99SBildn, GFN-xTB relaxed) 200 dipeptide MD simulations (~50 ns each) umbrella sampling GFN2-xTB data 7 test dipeptides 1 microsecond MD
Benchmarks: alanine dipeptide free energy difference 8-mode Gaussian mixture checkerboard 7 test dipeptides generalization
Metrics: KL divergence free energy difference energy W2 torsion W2 inference time NLL Spearman correlation
Baselines: ECNF GVP Vector Field GVP Endpoint Hutchinson trace estimator TBG-ECNF MD baseline InfoNCE only score matching only
Builds on: Boltzmann Generators stochastic interpolants InfoNCE equivariant continuous normalizing flows geometric vector perceptrons
Closest papers by meaning
- Scalable Equilibrium Sampling with Sequential Boltzmann Generators (ICML 2025 · similarity 0.90)
- Few-Step Boltzmann Generators via Scalable Likelihood Flow Maps (NeurIPS 2026 · similarity 0.89)
- Energy Matching: Unifying Flow Matching and Energy-Based Models for Generative Modeling (NeurIPS 2025 · similarity 0.89)
- Temperature-Annealed Boltzmann Generators (ICML 2025 · similarity 0.89)
- Transferable Boltzmann Generators (NeurIPS 2024 · similarity 0.89)
- Spherical Boltzmann machines: a solvable theory of learning and generation in energy-based models (NeurIPS 2026 · similarity 0.88)
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.