Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Diffusion-HMC: Parameter Inference with Diffusion-model-driven Hamiltonian Monte Carlo

Nayantara Mudur, Carolina Cuesta-Lazaro, Douglas P. Finkbeiner

TL;DR

Diffusion-HMC tackles cosmological parameter inference by training a conditional diffusion model to emulate cold dark matter density fields as a function of $(\Omega_m, \sigma_8)$ and by using the model’s conditional variational lower bound as an approximate likelihood for Hamiltonian Monte Carlo sampling of the posterior. The approach yields field-level emulation that reproduces key statistics, and, when combined with HMC, produces tighter, more robust constraints than a power-spectrum baseline, with demonstrated resilience to uncorrelated noise. The authors explore the information content across diffusion timesteps, show potential speed-accuracy tradeoffs with timesteps, and discuss robustness strategies, including truncation and priors, while releasing code for broader use. This work advances diffusion-model priors as a versatile tool for both generating cosmological fields and performing principled, robust inference on cosmological parameters at the field level.

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.

Diffusion-HMC: Parameter Inference with Diffusion-model-driven Hamiltonian Monte Carlo

TL;DR

Diffusion-HMC tackles cosmological parameter inference by training a conditional diffusion model to emulate cold dark matter density fields as a function of and by using the model’s conditional variational lower bound as an approximate likelihood for Hamiltonian Monte Carlo sampling of the posterior. The approach yields field-level emulation that reproduces key statistics, and, when combined with HMC, produces tighter, more robust constraints than a power-spectrum baseline, with demonstrated resilience to uncorrelated noise. The authors explore the information content across diffusion timesteps, show potential speed-accuracy tradeoffs with timesteps, and discuss robustness strategies, including truncation and priors, while releasing code for broader use. This work advances diffusion-model priors as a versatile tool for both generating cosmological fields and performing principled, robust inference on cosmological parameters at the field level.

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.
Paper Structure (22 sections, 7 equations, 12 figures)

This paper contains 22 sections, 7 equations, 12 figures.

Table of Contents

  1. Introduction
  2. Datasets, Architecture and Training
  3. Datasets
  4. Diffusion Model Setup
  5. Training
  6. Summary Statistics
  7. Parameter Inference
  8. Examining the influence of different timesteps
  9. HMC-based Parameter Inference
  10. Robustness
  11. Conclusion
  12. Code
  13. Acknowledgements
  14. Parameter Inference
  15. HMC Convergence
  16. ...and 7 more sections

Figures (12)

  • Figure 1: Generated fields at different cosmologies.Upper Row: Power spectrum of the unlogged fields for five different validation parameters. The lines depict the mean power spectrum and the envelope indicates the $16^{th}$ and $84^{th}$ percentiles of the distribution. True simulations are shown in black, generated in blue. Lower Row: Mean and standard deviation envelopes for the density histograms of the log fields.
  • Figure 2: Generated '1P' fields.Left column: Generated fields corresponding to the extreme values of each parameter for a single seed, with the other value held fixed at the fiducial value (0.3 for $\Omega_m$ and 0.8 for $\sigma_8$). Middle column: Power spectra of the generated fields for the same seed, for different values of each parameter, holding the other fixed. Right column: Mean and standard deviation for the ratio of the power spectra at the modified parameter value to the power spectra for the field at the fiducial parameter value (black) for 15 slices from the CAMELS dataset (solid) and 15 seeds for the generated fields from the diffusion model (dashed). The effect of modulating a parameter on the generated fields' power spectra is consistent with that of the true fields.
  • Figure 3: Generated 'CV' fields.Left: Power spectra of the 405 true and generated fields, with the mean, and $16-84$th percentiles. Right: Correlation matrix of the power spectra of the true and generated fields. The lower triangular matrix corresponds to the correlation matrix of the true fields while the upper triangular matrix corresponds to the correlation matrix of the generated fields.
  • Figure 4: Generated 'CV' fields.Upper: Ratio of the standard deviations of the generated fields to that of the true fields. Lower: Standard deviations of the power spectra in each $k$ bin for the generated and true fields.
  • Figure 5: Investigating the contribution of the timesteps used in the VLB sum. Left: One sigma contours of each $-2\Delta ln L_t$'s individual contribution for different timesteps. The contours are all centered near the true parameter (black star) and become wider as t increases. Right: Mean, and 1 sigma predictions for a single field as a function of the number of timesteps used in the VLB sum optimized by the HMC. Reducing the number of timesteps used to compute the VLB does not significantly affect the diffusion model predictions.
  • ...and 7 more figures