Scalable Boltzmann Generators for equilibrium sampling of large-scale materials

TL;DR

This work presents a scalable, transferable approach to equilibrium sampling of large-scale materials by introducing local augmented coupling flows that leverage graph neural networks. The method learns environment-conditioned, coordinate-transformations based on local neighborhoods, enabling one-shot generation of accurate equilibrium ensembles for systems with thousands of atoms while keeping training costs manageable. Through augmented variables and TFEP, the framework delivers reliable free-energy estimates (Helmholtz and Gibbs) across sizes and thermodynamic states, often outperforming global-flow architectures in both efficiency and scalability. The approach is demonstrated on Lennard–Jones crystals, Stillinger–Weber parametrizations, and silicon phases, enabling size-transferable sampling, volume fluctuations, and phase-diagram analysis with high fidelity to MD+MBAR references.

Abstract

The use of generative models to sample equilibrium distributions of many-body systems, as first demonstrated by Boltzmann Generators, has attracted substantial interest due to their ability to produce unbiased and uncorrelated samples in `one shot'. Despite their promise and impressive results across the natural sciences, scaling these models to large systems remains a major challenge. In this work, we introduce a Boltzmann Generator architecture that addresses this scalability bottleneck with a focus on applications in materials science. We leverage augmented coupling flows in combination with graph neural networks to base the generation process on local environmental information, while allowing for energy-based training and fast inference. Compared to previous architectures, our model trains significantly faster, requires far less computational resources, and achieves superior sampling efficiencies. Crucially, the architecture is transferable to larger system sizes, which allows for the efficient sampling of materials with simulation cells of unprecedented size. We demonstrate the potential of our approach by applying it to several materials systems, including Lennard-Jones crystals, ice phases of mW water, and the phase diagram of silicon, for system sizes well above one thousand atoms. The trained Boltzmann Generators produce highly accurate equilibrium ensembles for various crystal structures, as well as Helmholtz and Gibbs free energies across a range of system sizes, able to reach scales where finite-size effects become negligible.

Figures (8)

• Figure 1: Scheme of the size-transferable augmented flow. Left: A GNN learns environment-dependent particle embeddings, $\mathbf{h}_i$, in a small system and is transferable across system size via each particle’s local neighborhood, $\mathcal{N}_i$. Right: Physical and auxiliary particles are updated sequentially using the learned embeddings. In each step, either all auxiliary particles or all physical variables are updated .
• Figure 2: Radial distribution functions (left) in units of the interaction length $\sigma$ and reduced energy histograms (right) for FCC LJ with $N = 1372$ (top) and cubic mW ice with $N = 1728$ (bottom) as obtained from MD, the base distributions, and the local BGs. The local BG results were computed using models trained with $N = 216$ (mW) and $N = 256$ (LJ), with the red line indicating the corresponding half box length during training. No reweighting was applied.
• Figure 3: Effective sample size (ESS) plotted against the number of training steps for global and local BGs of cubic mW ice with $N = 64$ and $N = 216$ (left and middle panels), and FCC LJ with $N = 256$ (right panel). The ESS was evaluated every 1k steps using 1k samples. Solid lines represent running averages of ESS across five independent training runs, with thin, light-colored lines indicating the raw ESS values from individual runs. For the local BGs, the joint efficiency is shown. Learning rate reductions were applied after 250k and 500k steps.
• Figure 4: Left: Absolute reduced Helmholtz free energy estimates per particle for the cubic ice systems against the particle number as obtained from the local BG. Blue and orange lines indicate the evaluation of joint and marginal densities, respectively. Reference MD+MBAR values are shown as a gray hatched area corresponding to their mean $\pm 10^{-3}k_B T$. The red hatched are corresponds to the MD+MBAR results evaluated at $N=4096$. BG results were obtained using a model trained $N=216$. Uncertainties were estimated by training four independent models and evaluating each of them five times with $50$k samples using different random seeds. Right: Reduced free energy difference between cubic and hexagonal ice, $\Delta f = f_{\rm hex} - f_{\rm cubic}$, against the particle number. Hatched areas have the same meaning as for the left plot.
• Figure 5: Left: Reduced Gibbs free energy difference $\Delta g = g_{\rm HCP} - g_{\rm FCC}$ per particle between HCP and FCC crystal structures with 1080 particles in the LJ potential as a function of the cutoff radius. Shown are BG-based predictions (joint density estimates) alongside reference results from MD+MBAR. The red vertical line denotes the cutoff radius applied during training with 180 particles (quantities marked with an asterisk are expressed in LJ units, see SI). Right: Densities $\rho^*$ as obtained from the BGs (circles) and mean densities from MD $NPT$ simulations (crosses) for FCC (upper panel) and HCP (lower panel). Error bars for both BG and MBAR are smaller than the plotted marker size (see Fig. \ref{['fig:df_transferable']}).