BoltzNCE: Learning Likelihoods for Boltzmann Generation with Stochastic Interpolants and Noise Contrastive Estimation

TL;DR

This work addresses the bottleneck of sampling from the Boltzmann distribution $p(x) \propto \exp(-U(x)/K_B T)$ for large molecular systems by introducing BoltzNCE, an amortized likelihood framework that couples a Boltzmann Emulator with an Energy-Based Model (EBM). The EBM learns a likelihood surrogate using a hybrid of score matching and Noise Contrastive Estimation within a stochastic interpolant scheme, enabling fast, simulation-free evaluation of $\log \hat{\rho}_1(x)$ and thus rapid reweighting to unbiased Boltzmann statistics. The approach achieves up to $100\times$ speedups over exact likelihood computations on alanine dipeptide and at least $6\times$ speedups over MD on unseen dipeptide systems, with transfer learning demonstrated across multiple peptides. By decoupling emulator-based sampling from likelihood computation and leveraging amortized EBMs, BoltzNCE enables scalable, accurate Boltzmann statistics at scale, broadening the applicability of Boltzmann-based generative modeling in molecular science.

Abstract

Efficient sampling from the Boltzmann distribution given its energy function is a key challenge for modeling complex physical systems such as molecules. Boltzmann Generators address this problem by leveraging continuous normalizing flows to transform a simple prior into a distribution that can be reweighted to match the target using sample likelihoods. Despite the elegance of this approach, obtaining these likelihoods requires computing costly Jacobians during integration, which is impractical for large molecular systems. To overcome this difficulty, we train an energy-based model (EBM) to approximate likelihoods using both noise contrastive estimation (NCE) and score matching, which we show outperforms the use of either objective in isolation. On 2d synthetic systems where failure can be easily visualized, NCE improves mode weighting relative to score matching alone. On alanine dipeptide, our method yields free energy profiles and energy distributions that closely match those obtained using exact likelihoods while achieving $100\times$ faster inference. By training on multiple dipeptide systems, we show that our approach also exhibits effective transfer learning, generalizing to new systems at inference time and achieving at least a $6\times$ speedup over standard MD. While many recent efforts in generative modeling have prioritized models with fast sampling, our work demonstrates the design of models with accelerated likelihoods, enabling the application of reweighting schemes that ensure unbiased Boltzmann statistics at scale. Our code is available at https://github.com/RishalAggarwal/BoltzNCE.

Figures (10)

• Figure 1: Overview. BoltzNCE offers an accelerated alternative to exact model likelihood computation. Samples from a prior are first transformed to a distribution of conformers by a Boltzmann Emulator, which is easy to sample from but difficult to evaluate likelihoods for. The generated samples are then reweighted with likelihoods estimated by an energy-based model (EBM), which we train to approximate the emulator distribution with a hybrid score matching and noise contrastive estimation scheme. This EBM gives access to likelihoods in a single function call, enabling us to reweight to the target Boltzmann distribtion up to $100\times$ faster than exact computation.
• Figure 2: Low-dimensional Systems: Qualitative Results. EBM density learned on synthetic two-dimensional systems. (Above) An 8-mode Gaussian mixture. (Below) The checkerboard distribution. In both cases, the true density is shown in the leftmost column, and the results obtained with different methods are shown to the right. Using both objectives (right) provides the best performance.
• Figure 3: BoltzNCE: Qualitative Results. Results for BoltzNCE on alanine dipeptide trained on the biased dataset. We use a GVP vector field as the Boltzmann Emulator. BoltzNCE successfully captures the energy distribution and the free energy projection. (Left) Ramachandran plot of generated samples. (Middle) Energy histogram along with BoltzNCE reweighting. (Right) Calculated free energy surfaces for the angle $\varphi$ on the right.
• Figure 4: BoltzNCE Results on NY dipeptide. BoltzNCE inference results for NY dipeptide (top left) after fine-tuning. Energy distribution (top right), free energy surfaces along the $\varphi$ angle (bottom left) and the first TICA component (bottom right). BoltzNCE successfully captures the right energy distribution and free energy projections for the dipeptide.
• Figure 5: Energy Based Model training workflow