Diffusion-HMC: Parameter Inference with Diffusion-model-driven Hamiltonian Monte Carlo
Authors
Nayantara Mudur, Carolina Cuesta-Lazaro, Douglas P. Finkbeiner
Abstract
Diffusion generative models have excelled at diverse image generation and reconstruction tasks across fields. A less explored avenue is their application to discriminative tasks involving regression or classification problems. The cornerstone of modern cosmology is the ability to generate predictions for observed astrophysical fields from theory and constrain physical models from observations using these predictions. This work uses a single diffusion generative model to address these interlinked objectives -- as a surrogate model or emulator for cold dark matter density fields conditional on input cosmological parameters, and as a parameter inference model that solves the inverse problem of constraining the cosmological parameters of an input field. The model is able to emulate fields with summary statistics consistent with those of the simulated target distribution. We then leverage the approximate likelihood of the diffusion generative model to derive tight constraints on cosmology by using the Hamiltonian Monte Carlo method to sample the posterior on cosmological parameters for a given test image. Finally, we demonstrate that this parameter inference approach is more robust to small perturbations of noise to the field than baseline parameter inference networks.
Concepts
The Big Picture
Two numbers govern the universe’s large-scale structure: Ω_m (the fraction of the universe made of matter) and σ_8 (how unevenly that matter is clumped). These leave fingerprints on the dark matter density field, a map of where matter clusters across cosmic scales. Reading those fingerprints back out from an observation means working from effect to cause, and that’s hard.
Traditional approaches compress the field into a power spectrum, a curve capturing structure at each scale, losing information in the process. Richer methods that work directly with the full map require computing how likely a given observation is, which gets expensive fast. Nayantara Mudur, Carolina Cuesta-Lazaro, and Douglas Finkbeiner built a single diffusion model that does both: generate realistic dark matter fields for any cosmology, and work backward from an observed field to tightly constrain Ω_m and σ_8.
Key Insight: Using the approximate likelihood built into a trained diffusion model, the researchers run Hamiltonian Monte Carlo to sample the full posterior over cosmological parameters, getting tight, noise-resistant constraints without bespoke summary statistics.
How It Works
The model is a denoising diffusion probabilistic model (DDPM), the same mathematical machinery behind image generators like Stable Diffusion. It learns to progressively blur an image into noise, then reverse that process to reconstruct the original. Because it is conditional, it takes Ω_m and σ_8 as inputs, learning what dark matter fields look like across a wide range of cosmologies.
Training data comes from the CAMELS Multifield Dataset, drawn from IllustrisTNG simulations. Each simulation tracks 256³ dark matter particles over a 25 Mpc/h volume. The team trained on 10,500 two-dimensional field slices spanning 700 cosmologies, using a U-Net with circular convolutions to respect the periodic boundary conditions of cosmological simulations.
Once trained, the model serves two purposes:
- Forward emulator: Given Ω_m and σ_8, generate realistic dark matter density fields whose power spectra and higher-order statistics match true simulation outputs across the full parameter range.
- Inference engine: Given an observed field, constrain the parameters that produced it.
The inference step uses Hamiltonian Monte Carlo (HMC), a sampling method borrowed from physics. Picture a ball rolling across a hilly surface where valleys mark likely parameter combinations. HMC uses gradient information to explore efficiently, and those gradients come from the diffusion model itself.
By tracking how reconstruction error changes as Ω_m and σ_8 vary, the team extracts a differentiable approximate likelihood. HMC then walks through this likelihood surface, building a full posterior distribution.
What comes out is not a single best-fit cosmology but a map of which parameter combinations are consistent with the data, correlations and degeneracies included. The method also holds up under noise. When the researchers added Gaussian noise to test fields and compared Diffusion-HMC against a standard discriminative CNN, the CNN’s performance degraded noticeably. The diffusion model, trained to denoise fields at every corruption level, retained structural knowledge of the signal even in noisy conditions.
Why It Matters
This is simulation-based inference (SBI): the likelihood is too expensive to compute analytically, so you need other ways to do statistics. The idea that a diffusion model’s approximate likelihood can drive HMC is new. Generative models built for synthesis can be repurposed for analysis, running the same learned probability structure in reverse.
Full-field inference matters for cosmology because non-Gaussian features in the dark matter distribution carry real information about fundamental physics, especially at small scales where gravity has scrambled initial conditions. Doing that with a single trainable model, without hand-crafting bespoke statistics, is what makes this approach worth paying attention to.
Noise performance matters too. Upcoming surveys from DESI, the Vera Rubin Observatory’s LSST, and the Nancy Grace Roman Space Telescope will generate data at unprecedented scale but with complex systematic effects. Methods that degrade gracefully under noise are far more useful than those that need pristine conditions.
The current model is conditioned only on Ω_m and σ_8. Extending it to the full CAMELS parameter space is an obvious next step. And since HMC relies on an approximate likelihood, pinning down where that approximation breaks will be essential before applying the method to real survey data.
Bottom Line: A single diffusion model trained on dark matter simulations can both generate realistic cosmic fields and constrain cosmological parameters. HMC-based inference proves more noise-resistant than discriminative neural network baselines, opening a practical route to field-level cosmological analysis with upcoming surveys.
IAIFI Research Highlights
Score-based generative modeling from machine learning meets Hamiltonian Monte Carlo from statistical physics. A model built for synthesis doubles as a precision measurement tool for cosmological inference.
Diffusion models here pull double duty as generative emulators and approximate likelihood engines. The same idea applies to scientific inverse problems well beyond cosmology.
Full-field inference on dark matter density maps extracts cosmological constraints beyond power-spectrum summaries, probing non-Gaussian structure that carries information about matter clustering and fundamental cosmological parameters.
Future work could extend the framework to real observed fields and the full CAMELS parameter space; the paper is available at [arXiv:2405.05255](https://arxiv.org/abs/2405.05255).
Original Paper Details
Diffusion-HMC: Parameter Inference with Diffusion-model-driven Hamiltonian Monte Carlo
2405.05255
Nayantara Mudur, Carolina Cuesta-Lazaro, Douglas P. Finkbeiner
Diffusion generative models have excelled at diverse image generation and reconstruction tasks across fields. A less explored avenue is their application to discriminative tasks involving regression or classification problems. The cornerstone of modern cosmology is the ability to generate predictions for observed astrophysical fields from theory and constrain physical models from observations using these predictions. This work uses a single diffusion generative model to address these interlinked objectives -- as a surrogate model or emulator for cold dark matter density fields conditional on input cosmological parameters, and as a parameter inference model that solves the inverse problem of constraining the cosmological parameters of an input field. The model is able to emulate fields with summary statistics consistent with those of the simulated target distribution. We then leverage the approximate likelihood of the diffusion generative model to derive tight constraints on cosmology by using the Hamiltonian Monte Carlo method to sample the posterior on cosmological parameters for a given test image. Finally, we demonstrate that this parameter inference approach is more robust to small perturbations of noise to the field than baseline parameter inference networks.