Macrocanonical Generator Networks: Data-Efficient Neural Surrogates for Amortized Physics Simulation

Niall Jeffrey; Benjamin D. Wandelt

cs.LGNov 3, 2026

Macrocanonical Generator Networks: Data-Efficient Neural Surrogates for Amortized Physics Simulation

Authors: Niall Jeffrey, Benjamin D. Wandelt

Organizations: Dept. of Physics & King’s Institute for Artificial Intelligence, King’s College London · Dept. of Physics and Astronomy & Dept. of Applied Mathematics and Statistics, Johns Hopkins University

Abstract

Many scientific simulations are computationally expensive, limiting their use in simulation-based inference, uncertainty quantification, and decision making. We introduce macrocanonical generator networks, a framework for learning fast, data-efficient neural surrogates for amortized simulation of stationary physical processes by matching multiscale statistics of a target process. Inspired by microcanonical maximum-entropy synthesis, which matches feature statistics per sample by inner-loop optimization, our approach instead trains a neural generator to satisfy the same constraints in expectation, preserving physically realistic sample-to-sample fluctuations. We prove that gradient-descent training induces preconditioned gradient-descent dynamics in sample space that inherit the symmetry-preservation properties of the microcanonical gradient descent of Bruna and Mallat, amortizing its per-sample inner-loop optimization into a single training run after which every sample is generated in one forward pass. We further show that residual generator parameterizations initialized near the identity, trained with a small learning rate, preserve entropy locally during early optimization and mitigate mode collapse. Experiments on cosmology and fluid simulations show that macrocanonical generators recover realistic samples from extremely small training sets, in some cases a single training example, while generating new realizations more than $10^6$ times faster per sample than the reference simulator and more than $10^3$ times faster than microcanonical gradient descent. The trained generators can be reused for downstream pretraining, including simulation-based inference, providing a data-efficient route to multi-fidelity scientific machine learning workflows.

Figures & tables

Figure 1: Macrocanonical generator networks ( centre ) turn a small set of expensive simulator outputs ( left ; turbulence and cosmology targets, Section 5 ) into large ensembles of new realizations matching the chosen statistical features (e.g. wavelet phase harmonics) of the targets, unlocking downstream scientific machine-learning tasks ( right ) that are intractable with the expensive simulator alone.

	Macrocanonical	GAN	Diffusion	Microcanonical
	(ours)		model	synthesis
Training samples needed	$\geq 1$ realization	large dataset	large dataset	$\geq 1$ realization
Inference cost per sample	1 forward pass	1 forward pass	many ( $\sim 10^{2\text{--}3}$ ) steps	inner-loop optimization
Symmetries of moment constraints	exact, by construction	approximate (learned)	approximate (learned)	exact, by construction
Per-realization fluctuations	yes (constraints in expectation)	yes	yes	no (Dirac on $\phi^{\ast}$ )
Training signal	feature-moment matching	adversarial	denoising score matching	per-sample feature matching

Table 1: Macrocanonical generator networks (ours) compared with prevailing deep generative methods and with classical microcanonical synthesis. Properties marked “exact, by construction” follow from architectural equivariance and the loss; properties marked “approximate (learned)” hold only insofar as the training set is large and diverse enough to make them statistically apparent.

Figure 2: Simulation-based inference of the dark-energy fraction $\Omega_{\rm DE}$ from 2-D density images (Section 5 ). (a, b) A validation example before and after the observational effects (i.e. masking). (c) Posterior after fine-tuning on $N$ “real” simulations. (d) Validation NLL against $N$ (inc. standard deviation with 3 seeds) using cosmologies in the $\Omega_{\rm DE}$ training range. The surrogate-pretrained model improves on the prior before any “real” simulation is used and consistently outperforms the baseline.

Figure 3: Inductive-bias ablation on the MHD (top) and cosmic-web (bottom) benchmarks; one target slice and one synthesis from each variant after $100$ optimizer steps. The three failure modes correspond to removing the small step ( $\eta=10^{-2}$ ), the residual shortcut, and (additionally) the circular-padded convolutions; zero-pad means the U-Net is no longer translation-equivariant, so Corollary 1 no longer guarantees stationarity (producing boundary-correlated structure).

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

$N_{\text{train}}$	WPH distance	KS mean- $p$	NN ratio
$1$	$0.20$ ( $0.053$ )	$0.08$ ( $0.54$ )	$0.89$ ( $1.00$ )
$10$	$0.075$ ( $0.056$ )	$0.27$ ( $0.47$ )	$0.98$ ( $0.96$ )
$100$	$0.086$ ( $0.037$ )	$0.31$ ( $0.50$ )	$0.99$ ( $1.01$ )

Appendix

Table 2: Effect of training-set size (

64{\times}64

MHD suite of Appendix F , scored against

100

held-out fields; parentheses: held-out fields scored against themselves).

Method	$1$	$10$	$100$	ms/sample	Outcome at $N_{\text{train}}=1$
	NN ratio at $N_{\text{train}}$
Macrocanonical generator (ours)	$0.89$	$0.98$	$0.99$	$0.9$	new fields
Microcanonical synthesis	$0.87$	–	–	$9200$	statistics pinned (no fluctuations)
DDPM	$0.32$	$0.58$	$0.86$	$107$	memorizes training field
WGAN-GP	$0.010$	$0.33$	$0.84$	$0.09$	memorizes training field
WGAN-GP + translation aug.	$0.24$	–	–	$0.10$	memorizes shifts of it

Appendix

Table 3: Baseline comparison on the

64{\times}64

MHD suite of Appendix F : all methods trained on the same fields and scored against

100

held-out fields. NN ratios much less than

1

mean memorization; held-out fields scored against themselves give

1.00

0.96

and

1.01

N_{\text{train}}=1

10

and

100

. Sampling time is per field on one A100.

$\eta$	Median per-step change	Input-Jacobian deviation, max (final)	Diversity, min (final)	Loss
$10^{-5}$	$0.03$	$0.40$ ( $0.39$ )	$0.94$ ( $1.02$ )	$2.0$
$10^{-4}$ (paper)	$0.085$	$0.65$ ( $0.63$ )	$0.93$ ( $0.98$ )	$1.2$
$10^{-3}$	$0.14$	$0.70$ ( $0.69$ )	$0.83$ ( $0.99$ )	$0.76$
$10^{-2}$	$0.46$	$322$ ( $0.90$ )	$0.70$ ( $0.98$ )	$0.83$

Appendix

Table 4: Learning-rate sweep (

64{\times}64

MHD,

600

iterations, all other settings fixed). The per-step change is the median relative output change on a fixed probe batch; the linearization of Theorem 1 requires it to be small. The input-Jacobian deviation and diversity are given as the extreme value over training and, in parentheses, the value at the end of training; the loss is the mean over the last

50

iterations.

$\sigma$	$\\|F(x)-x\\|/\\|x\\|$ at initialization	Input-Jacobian deviation	Diversity	Loss
$10^{-5}$	$<10^{-11}$	$0.81$	$0.96$	$0.51$
$10^{-4}$ (paper)	$5\times 10^{-5}$	$0.61$	$0.94$	$0.70$
$10^{-3}$	$0.046$	$0.86$	$0.96$	$0.44$
$10^{-2}$	$0.44$	$2.1$	$0.93$	$0.36$
$3\times 10^{-2}$	$1.3$	$4.9$	$0.91$	$0.49$
$10^{-1}$	$4.5$	$11.1$	$0.90$	$0.85$

Appendix

Table 5: Initialization-scale sweep (

64{\times}64

MHD,

\eta=10^{-4}

600

iterations, all other settings fixed). The loss is the mean over the last

50

iterations; the diversity and input-Jacobian deviation are measured at the end of training.

Benchmark	$J$	$L$	$\Delta j$	$\Delta\ell$	$\Delta n$	$p$ -list
Cosmological cosmic web (Quijote)	5	4	4	2	0	$\emptyset$
MHD turbulent fluid simulation	5	4	4	2	1	pywph default

Appendix

Table 6: WPH summary-statistic configuration per benchmark; symbols match [ 23 ] .

Figure 4: Across-sample distributions of four randomly chosen WPH coefficients from $100$ macrocanonical generator outputs ( blue ) and $100$ held-out Quijote simulator slices ( orange ). Microcanonical synthesis [ 4 , 35 ] would instead collapse each panel to a Dirac mass at $\phi^{\ast}$ .

Cosmology

See all 8 papers

Sep 30, 2026

Grand Canonical Generators

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.

Andreas Burger, Malte Franke, Luka Mucko +2

University of Toronto NVIDIA Vector Institute Toronto, Canada · ETH Zurich NCCR Catalysis Zurich, Switzerland · University of Toronto NVIDIA Vector Institute Acceleration Consortium CIFAR Toronto, Canada

Generative Modeling Molecular Generation Conditional Generative Modeling Boltzmann Sampling

Ask PDF

Feb 23, 2026

Exact Discrete Stochastic Simulation with Deep-Learning-Scale Gradient Optimization

Exact stochastic simulation of continuous-time Markov chains (CTMCs) is essential when discreteness and noise drive system behavior, but the hard categorical event selection in Gillespie-type algorithms blocks gradient-based learning. We eliminate this constraint by decoupling forward simulation from backward differentiation, with hard categorical sampling generating exact trajectories and gradients propagating through a continuous massively-parallel Gumbel-Softmax straight-through surrogate. Our approach enables accurate optimization at parameter scales over four orders of magnitude beyond existing simulators. We validate for accuracy, scalability, and reliability on a reversible dimerization model (0.09% error), a genetic oscillator (1.2% error), a 203,796-parameter gene regulatory network achieving 98.4% MNIST accuracy (a prototypical deep-learning multilayer perceptron benchmark), and experimental patch-clamp recordings of ion channel gating (R^2 = 0.987) in the single-channel regime. Our GPU implementation delivers 1.9 billion steps per second, matching the scale of non-differentiable simulators. By making exact stochastic simulation massively parallel and autodiff-compatible, our results enable high-dimensional parameter inference and inverse design across systems biology, chemical kinetics, physics, and related CTMC-governed domains.

Jose M. G. Vilar, Leonor Saiz

Biofisika Institute (CSIC, UPV/EHU), University of the Basque Country (UPV/EHU), P.O. Box 644, 48080 Bilbao, Spain · IKERBASQUE, Basque Foundation for Science, 48011 Bilbao, Spain · Department of Biomedical Engineering, University of California, 451 E. Health Sciences Drive, Davis, CA 95616, USA

Straight-Through Estimator Markov Chains Differentiable Physics

Ask PDF

Aug 7, 2026

Leveraging generative models to assist Monte Carlo sampling

Sampling high-dimensional probability distributions is a central task in scientific computing, with applications ranging from Bayesian inference to statistical physics and molecular simulation. Despite decades of methodological developments, two major challenges remain: scaling to high dimensions and efficiently exploring multimodal distributions characterized by metastable states. Classical approaches such as Markov chain Monte Carlo, tempering methods, or enhanced sampling based on collective variables have achieved major successes, but they also face intrinsic limitations. This tutorial review explores a new paradigm that has recently emerged at the interface of machine learning and computational statistical physics: the use of generative models as tools for sampling. In this context, models such as normalizing flows and diffusion models are not used in their traditional data-driven setting, but rather as flexible probabilistic models that can assist the sampling of distributions known only up to a normalization constant. This manuscript reviews the early development of this rapidly evolving field and discusses several methodological directions, including exact samplers based on generative models and strategies to train such models in the absence of data. While an exhaustive survey of the literature is not attempted, we present a selection of key ideas and methods, along with a discussion of their strengths and limitations. The review is intended to be an accessible tutorial for both physics and machine learning audiences, and it aims to provide a starting point for researchers interested in exploring this exciting area of research.

Marylou Gabrié

Laboratoire de Physique Statistique, École normale supérieure, PSL Research University, 24 rue Lhomond, 75005 Paris, France

Normalizing Flows Generative Modeling Diffusion Models

Ask PDF

Abstract

Cosmology

Explore similar work

Grand Canonical Generators

Exact Discrete Stochastic Simulation with Deep-Learning-Scale Gradient Optimization

Leveraging generative models to assist Monte Carlo sampling