Source-linked AI summary
Grand Canonical Generators
Andreas Burger, Malte Franke, Luka Mucko, Kjell Jorner, Alan Aspuru-Guzik
TL;DR
Open molecular systems require sampling variable particle numbers as well as configurations, a challenge for fixed-size Boltzmann generators and grand-canonical Monte Carlo. GCG provides joint and factorized designs; experiments report reproduced observables across fluid and adsorption systems, with chemical-potential transfer and importance-sampling correction.
Problem
Grand-canonical sampling must cover variable particle numbers and configurations, while canonical Boltzmann generators fix particle number and GCMC can be slow to explore broad particle-number ranges.
Source: §1 · pp. 1–2Method
GCG offers a joint variable-size generator and a factorized design that combines a particle-number model with a conditional canonical Boltzmann generator and analytically encodes chemical-potential dependence.
Source: §1 · p. 2Results
Takeaways & Limitations
The factorized model’s tractable likelihood supports grand-canonical importance sampling to correct residual model error, and generated samples can accelerate GCMC.
Source: §6 · p. 9The models require equilibrium training samples collected from GCMC; the methane adsorption study treats methane as point-like and the zeolite crystal as rigid.
Source: §6 · p. 9
Abstract
from arXiv · showhide
We introduce Grand Canonical Generators (GCG), a generative framework that extends Boltzmann generators to the grand canonical ensemble. We present two designs. The first conditions a variable-size generative model on the chemical potential, sampling particle number and configuration jointly. The second factorizes the grand canonical distribution into a particle-number distribution and the corresponding canonical Boltzmann density. This factorized formulation can use any existing Boltzmann generator for the canonical component, encodes the known linear chemical-potential dependence analytically, and yields a tractable likelihood that supports self-normalized importance sampling (SNIS). Empirically, GCG accurately reproduces grand canonical observables on a Lennard--Jones fluid and methane adsorption in a zeolite, demonstrating generalization across chemical potentials and correction via SNIS and grand canonical Monte Carlo.
1 INTRODUCTION
Open molecular systems require sampling both particle number and molecular configuration, but conventional grand canonical Monte Carlo can be slow across particle numbers. GCG extends Boltzmann generators to this setting with joint and factorized designs, and is validated on three systems.
1 INTRODUCTION
- GCMC insertion moves often have low acceptance, and sequential exploration makes metastable states and broad particle-number ranges slow to sample.Each particle number must be visited often enough to estimate its statistical weight reliably.Source: §1 · p. 1
- Grand-canonical sampling requires a variable-size state space and particle-number probabilities tied to canonical partition functions, which shift exponentially with chemical potential.Canonical Boltzmann generators fix particle number and do not estimate the partition functions needed for those probabilities.Source: §1 · p. 2
- GCG offers a joint model and a factorized model; the latter analytically encodes chemical-potential dependence, learns partition-function ratios, and provides a tractable joint likelihood for SNIS.The joint model learns chemical-potential dependence from data but lacks a tractable likelihood; stratified SNIS maintains coverage across relevant particle numbers.Source: §1 · p. 2
- GCG recovers ideal-gas behavior, reproduces reference observables for a Lennard–Jones fluid after importance correction, transfers to nearby chemical potentials, and recovers a methane-zeolite adsorption isotherm.The three tests probe variable-size generation without interactions, an interacting fluid, and adsorption in a crystalline host.Source: §1 · p. 2
2 BACKGROUND
The grand canonical ensemble describes configurations with variable particle number, whereas canonical Boltzmann generators model configurations at a fixed number. This section sets out the density and sampling methods that motivate GCG.
2 BACKGROUND
- The grand canonical ensemble assigns a distribution over particle number and coordinates, with chemical potential controlling particle exchange with a reservoir.Its normalization is the grand canonical partition function Ξ, and each particle number has a configuration space of its own.Source: §2 · p. 3
- Canonical Boltzmann generators model configurations at fixed N and can use self-normalized reweighting without estimating the canonical partition function ZN.In the grand canonical ensemble, particle-number probabilities depend on ZN and shift exponentially with chemical potential.Source: §2 · p. 3
- Flow matching trains a continuous normalizing flow from source and target configurations, then maps source samples to generated configurations by integrating a learned velocity field.The endpoint log-density follows from the change-of-variables formula; evaluating it involves the velocity-field Jacobian divergence.Source: §2 · p. 3
- Self-normalized importance sampling uses proposal samples and their target-to-proposal weights to estimate observables when the model log-density is available.It is asymptotically unbiased but has finite-sample bias, and effective sample size measures proposal–target overlap.Source: §2 · p. 3
- Generator Matching extends flow matching to variable-size Markov processes by jointly matching continuous coordinate drift and discrete particle-number transition rates.Because endpoint likelihoods for jump processes generally require integrating over possible paths, they are generally not tractable.Source: §2 · p. 4
3 RELATED WORK
Prior neural equilibrium samplers generally assume fixed particle number, while several variable-size or grand-canonical approaches target lattice systems, design, or optimization rather than continuous-position equilibrium sampling.
3 RELATED WORK
- Neural equilibrium generators with tractable likelihoods support importance reweighting, but existing methods assume fixed particle number.GCG adds a tractable particle-number distribution alongside a spatial generator.Source: §3 · p. 4
- Existing learned grand-canonical approaches include lattice-focused models, partition-function estimation without samples, lattice-site sampling, and continuous-existence optimization rather than equilibrium sampling.The paper instead targets continuous particle positions without lattice constraints and learns a normalized proposal for grand-canonical importance sampling.Source: §3 · p. 4
- Variable-size flows, autoregressive models, and jump processes have been used for point-set generation, molecular generation, or design, but had not been adapted to the grand canonical ensemble.The paper positions these methods as possible components within its GCG framework.Source: §3 · p. 4
4 GRAND CANONICAL GENERATORS
GCG combines a variable-size sampling design with a particle-number model and a conditional configuration generator. The factorized design exploits the ensemble’s known chemical-potential dependence to support transfer and likelihood-based correction.
4 GRAND CANONICAL GENERATORS
- GCG presents joint and factorized constructions with different trade-offs in flexibility and likelihood evaluation, alongside a structured particle-number model and stratified importance correction.The framework also specifies model parameterizations used in the experiments.Source: §4 · p. 4
- Joint GCG generates particle number and configuration in one variable-size process, but must learn chemical-potential dependence from data and has no tractable endpoint likelihood.Its implementation uses graph-level insertion intensity, insertion-coordinate flow matching, periodic coordinates, and periodic optimal-transport couplings.Source: §4 · pp. 5–6
- Factorized GCG separates the particle-number distribution from the conditional canonical configuration distribution, so the latter is independent of chemical potential.It samples N first and then X; any likelihood-evaluable Boltzmann generator conditioned on N can provide the configuration component.Source: §4 · p. 5
- The particle-number model uses categorical probabilities over N and can encode known grand-canonical factors analytically to improve transfer across chemical potentials.Its logit retains βµN and −log(N!) while a learned residual represents log ZN without depending on µ; the residual is anchored at N = 0.Source: §4 · p. 5
- The factorized joint likelihood supports grand-canonical SNIS, while stratified sampling allocates samples across particle numbers to address missed, low-proposal-probability states.GCG samples can also initialize conventional GCMC chains near equilibrium without requiring likelihood evaluations.Source: §4 · pp. 5–6
- The implementation uses Morph with a Diffusion Transformer backbone; the factorized spatial model shares the periodic backbone, while zeolite adsorption adds absolute-position features and local message passing.The zeolite setup removes translation and rotation augmentations and uses independent coupling.Source: §4 · p. 6
5 EXPERIMENTS
Experiments test GCG on an interacting Lennard–Jones fluid, chemical-potential transfer, inference-time correction, and methane adsorption. Factorized GCG closely matches reference observables, though transfer and correction have limits in some settings.
Lennard–Jones fluid
- Both GCG variants recover the Lennard–Jones fluid's radial distribution and energy density; Factorized GCG matches the reference most closely, while Morph misplaces the first radial peak.Morph's first peak is at r ≈ 1.5σ rather than 1.1σ, and its mean potential energy is −1.9ϵ versus −2.7ϵ.Source: §5 · p. 7
Chemical-potential transfer
- At held-out chemical potentials, βµ-tilted qϕ(N) has at most 1.1% relative error, versus 8% for Joint GCG and 6% for unstructured qϕ(N).At βµ = −2.6, the tilted model's relative error is 0.6%; at βµ = −2.4, all models fail, with the tilted model underestimating ⟨N⟩ by 20%.Source: §5 · p. 7
Inference-time correction
- With 1.6 × 10^4 target-energy evaluations, conditional SNIS lowers energy W1 from 0.6 to 0.3 and pair-distance W1 from 1.6×10−3 to 9×10−4, while GC-SNIS raises particle-number W1 from 0.08 to 0.2.GCMC needs 92× and 25× more evaluations to match the conditional-SNIS energy and pair-distance errors, respectively.Source: Figure 4 · p. 8, §5 · p. 8
- After two trial moves per chain, GCG-seeded GCMC has 23× lower particle-number W1 than cold-start GCMC: 1.2 versus 28 for 64 chains.Particle-number and energy errors remain below the cold-start baseline until ∼4 × 10^6 evaluations.Source: §5 · p. 8
Methane adsorption in a zeolite
- Across held-out pressures up to 150 bar, Factorized GCG's methane particle-number error is at most 0.5 molecules (< 0.4%), matching GCMC even in extrapolation.Generated configurations also reproduce methane distributions across the channels and the pair-distance distribution at 100 bar.Source: §5 · p. 8, Figure 5 · p. 9
6 CONCLUSION
Grand Canonical Generators extend Boltzmann generators to the grand canonical ensemble, with Factorized GCG combining a particle-number model and a conditional Boltzmann generator. The conclusion highlights its tractable likelihood and chemical-potential transfer, while noting training and model-scope limitations.
6 CONCLUSION
- Factorized GCG combines a particle-number model with any conditional Boltzmann generator, analytically encodes chemical-potential dependence, and provides a tractable joint likelihood for importance sampling.The model learns canonical partition-function ratios, improving transfer to chemical potentials beyond the training range; generated samples can also accelerate GCMC for models without likelihoods.Source: §6 · p. 9
- Autoregressive architectures could unify joint and factorized GCG; proposed directions also include multiple hosts and particles, flexible hosts, and less reliance on GCMC training data.The conclusion notes that autoregressive models generate variable-size systems particle by particle while providing exact likelihoods.Source: §6 · p. 9
- GCG training requires equilibrium samples collected from GCMC, and the methane study assumes point-like methane and a rigid zeolite crystal.Extending to polar adsorbates or flexible frameworks would require sampling additional degrees of freedom.Source: §6 · p. 9
A.1.1 IDEAL GAS
The ideal gas provides an interaction-free test of variable-size generation and spatial uniformity. Joint and Factorized GCG reproduce the target more faithfully than Morph, with Factorized GCG giving the best overall agreement in the reported error comparison.
IDEAL GAS
- The ideal-gas test uses a unit-volume periodic box and evaluates activities 1, 4, and 8 with 65,536 configurations each.The exact target has uniform one-coordinate marginals and g(r) = 1.Source: §A.1.1 · pp. 12–13
- The evaluation measures coordinate-histogram TV and radial-distribution MAE against ideal-gas targets, with uniform samples estimating finite-sample error.Coordinates are divided into 32 bins; radial distances use 25 bins, and finite-sample error is largest in low-occupancy inner radial shells.Source: §A.1.1 · p. 13
- Morph’s central density is about 13× the target at activity one, while both GCG variants recover nearly uniform marginals at all three activities.Morph also leaves outer bins near the box faces empty at activities four and eight; this is consistent with its non-periodic coordinates and Gaussian source.Source: Figure 6 · p. 12, §A.1.1 · p. 13
- At activity four, Factorized GCG has g(r) MAE 0.052 versus Joint GCG’s 0.216, and gives the best overall agreement.Joint GCG reduces the g(r) MAE from Morph’s 0.509 to 0.216; the uniform-sample finite-sample noise is 0.021.Source: §A.1.1 · p. 13, Figure 7 · p. 13
A.1.2 PARTICLE-NUMBER MODEL PARAMETERIZATION
The particle-number model separates the known ideal-gas contribution from a learned interaction residual. This reference structure improves transfer to unseen chemical potentials, although accuracy remains limited outside the training range.
PARTICLE-NUMBER MODEL PARAMETERIZATION
- The model separates log ZN into a one-particle reference term and a residual representing interactions and deviations from that reference.For the homogeneous Lennard–Jones system, q1 is set to the volume in reduced coordinates; for heterogeneous systems it can represent accessible volume or a host-specific partition function.Source: §A.1.2 · p. 14
- The residual rϕ is an MLP output conditioned on particle number and can additionally use host and adsorbate features for transfer across systems.It is constrained to rϕ(0, β, c) = 0 and initialized so the starting distribution matches the analytic reference.Source: §A.1.2 · p. 14
- At βµ = −2.5, the ideal-gas reference lowers transfer error from 5% to 3% versus the βµ-tilted model; at −2.4, it lowers error from 19% to 14%.For βµ ≥ −2.2, no model is accurate; the ideal-gas model has the smallest error, ranging from 13% to 35%.Source: §A.1.2 · p. 14
A.1.3 EFFECTIVE SAMPLE SIZE
The appendix reports effective-sample-size behavior and compares correction methods, then specifies the systems, metrics, and model settings used in evaluation.
A.1.3 EFFECTIVE SAMPLE SIZE
- Pooled ESS/M is 0.13 for conditional SNIS and 0.12 for GC-SNIS at 16,400 configurations; within-sector values generally decline as N increases.Within-sector ESS/M is ∼0.3–0.5 for N ≤20 and ∼0.1–0.2 for N ≥35; N = 45 and N = 50 fall to ∼0.02, but their model probabilities are only 0.5% and 0.1%.Source: §A.1.3 · p. 15
A.1.4 BENNETT ACCEPTANCE RATIO
- BAR combines insertion and deletion estimates of adjacent canonical partition-function ratios to correct relative particle-number probabilities, whereas conditional SNIS corrects configurations at fixed N.Conditional SNIS weights are capped at their 99th percentile before estimating adjacent particle-number ratios.Source: §A.1.4 · pp. 15–16
- Conditional SNIS and GC-SNIS reach low errors by direct reweighting, while GCMC improves with chain budget and BAR with configurations per particle number and insertion trials.BAR performance depends on proposal budget; Figure 10 compares particle-number, energy, and pair-distance W1 errors against target-energy evaluations, with lower values indicating better agreement.Source: §A.1.4 · p. 16, Figure 10 · p. 16
A.2.1 2D REPELLING PARTICLES IN A TRAP
- The 2D trap example uses Coulomb-repelling particles in an anisotropic harmonic trap, with confinement four times stronger along y than x.The setup favors elongated arrangements; for two or more particles, repulsion separates them along x.Source: §A.2.1 · pp. 16–17
A.2.2 PERIODIC LENNARD–JONES FLUID
- The periodic Lennard–Jones fluid uses a cubic box of side L = 6.5σ at T* = 1.5, with training chemical potentials βµ ∈ {−3.75, −3.25, −3.00, −2.75} and a held-out test at βµ = −3.5.The held-out reference comprises 32,768 configurations from four independent GCMC chains.Source: §A.2.2 · p. 17
A.2.3 METHANE ADSORPTION IN MFI ZEOLITE
- The methane adsorption study uses a rigid all-silica MFI zeolite and a classical force field at 298 K; models are trained at seven pressures from 0.1 to 70 bar and tested at 50–150 bar.The 150-bar condition contains up to 140 methane molecules.Source: §A.2.3 · p. 17
A.2.4 EVALUATION METRICS
- Evaluation uses 1-Wasserstein distance between model and GCMC-reference distributions for scalar observables such as particle number, potential energy, and pair distance.W1 has the observable’s units, bounds the mean error, and is reported as less prone to outliers than W2.Source: §A.2.4 · p. 17
A.3 MODEL AND HYPERPARAMETERS
- Models are trained for 10 to 20 GPU hours on an NVIDIA A100, and Table 1 lists hyperparameters for the Lennard–Jones fluid and MFI zeolite.The supplied table caption identifies its scope but does not provide the parameter values.Source: §A.3 · p. 17, Table 1 · p. 18