Coarse-Grained Boltzmann Generators

1 Professorship of Multiscale Modeling of Fluid Materials, Technical University of Munich, Germany
2 Atomistic Modeling Center, Munich Data Science Institute, Technical University of Munich
*Equal contribution, random order
International Conference on Machine Learning (ICML) 2026

Abstract

Sampling equilibrium molecular configurations from the Boltzmann distribution is a longstanding challenge. Boltzmann Generators (BGs) address this by combining exact-likelihood generative models with importance sampling, but practical scalability is limited. Meanwhile, coarse-grained surrogates enable the modeling of larger systems by reducing effective dimensionality, yet often lack a reweighting procedure required to ensure asymptotically correct statistics. In this work, we propose Coarse-Grained Boltzmann Generators (CG-BGs), a framework for reduced-order generative modeling with importance sampling in coarse-grained coordinate space. CG-BGs generate samples using a flow-based model and reweight them using a learned potential of mean force (PMF). We show that the PMF can be learned from rapidly converged trajectories via enhanced sampling force matching. Experiments demonstrate that CG-BGs capture solvent-mediated interactions in highly reduced representations while substantially reducing computational cost relative to atomistic BGs, providing a practical route toward equilibrium sampling of larger molecular systems.

CG-BGs capture solvent-mediated interactions in highly reduced representations

Operating directly in a coarse-grained space, CG-BGs capture solvent-mediated interactions through the learned PMF, surpassing classical implicit-solvent models while substantially reducing the cost relative to explicit MD simulation and atomistic BGs.

Model Overview

Atomistic configurations are mapped to CG coordinates to construct training data. A PMF network learns Uη(R) from rapidly converged data, while a normalizing flow learns a proposal density qθ(R) over CG configurations. Samples generated from the flow are reweighted using the PMF to recover the target equilibrium distribution p(R), enabling unbiased estimation of thermodynamic observables.

CG-BG on the Mueller-Brown potential

Recovering Equilibrium Distributions from Biased and Unbiased Data

After importance reweighting, CG-BG recovers the equilibrium free-energy landscapes of the alanine peptides from flow models trained on either biased or unbiased trajectories, in close agreement with explicit-solvent MD. Across all systems it surpasses the classical implicit-solvent baselines that upper-bound atomistic Boltzmann Generators.

Method JS divergence (↓) PMF error (↓) ESS (↑)
Alanine dipeptide
CG-BG (reweighted)
Heavy Atom0.0048(1)0.2005(63)0.5112(4)
Heavy Atom (Biased)0.0063(1)0.2277(66)0.4115(4)
Core Beta0.0052(1)0.2210(65)0.5528(4)
Core Beta (Biased)0.0057(1)0.2093(58)0.4818(4)
Implicit solvent (GB)
OBC10.0157(2)0.3709(92)
OBC20.0182(2)0.4028(95)
Alanine tripeptide
CG-BG (reweighted)
Heavy Atom0.0056(1)0.1957(52)0.3201(4)
Core Beta0.0060(1)0.2112(51)0.4212(5)
Implicit solvent (GB)
OBC20.0932(3)1.0274(65)
Alanine hexapeptide
CG-BG (reweighted)
Core Beta0.0100(1)0.3646(81)0.1231(3)
Implicit solvent (GB)
OBC20.1652(3)1.8401(70)

GB: generalized Born implicit-solvent model; OBC1/OBC2 are its two parameterizations.

Effect of Coarse-Graining Resolution on Accuracy and Efficiency

As shown in the above table, reducing the resolution from Heavy Atom to Core Beta raises the effective sample size and lowers training and inference cost, while in general showing comparable accuracy of quantitative results.

Stage Core Beta Heavy Atom All Atom
Training0.45 h0.80 h2.55 h
Inference0.95 min3.78 min14.91 min

Training and inference time on alanine dipeptide; inference is measured over 104 samples.
All Atom generates the full solute configuration without explicit solvent.

Simulation-Free Evaluation of Learned PMFs

The flow model generates CG configurations from a proposal distribution approximating the equilibrium ensemble, and importance reweighting maps these samples to the Boltzmann distribution induced by a given PMF. This enables rapid, simulation-free assessment of candidate CG potentials, in contrast to traditional validation pipelines that require separate MD simulations for each model.

Simulation-free comparison of learned PMFs for alanine dipeptide

Simulation-free PMF benchmarking. Reweighting a single set of CG-BG proposals under potentials trained on unbiased (PMFU) and rapidly converged biased (PMFB) data. PMFB matches the explicit-solvent MD reference, while PMFU misestimates the metastable populations along φ.

BibTeX

@article{chen2026cgbg,
  title={Coarse-Grained Boltzmann Generators},
  author={Chen, Weilong and Zhao, Bojun and Eckwert, Jan and Zavadlav, Julija},
  journal={arXiv preprint arXiv:2602.10637},
  year={2026}
}