# Topology of Wilson-loop spectrum and periodic evolution of surface-state
Fermi arc

Yi-Dong Wu<EMAIL_ADDRESS>Key Laboratory for Microstructural Material
Physics of Hebei Province, School of Science, Yanshan University, Qinhuangdao
066004, China

###### Abstract

Wilson-loop has been widely used to characterize the topological property of
topological insulators, high order topological insulators and topological
semimetals. Both bulk topological invariants and nontrivial boundary
properties can be deduced from the topology of Wilson-loop spectrum. However,
no attempt has been made to observe it. In this letter we demonstrate the
topology of Wilson-loop spectrum can be observed by using the existing
technology for observing surface-state Fermi arc. We predict that the surface-
state Fermi arc sweeps the whole surface Brillouin zone in a continuous
periodic process and thus generates a surface that is topologically equivalent
to the Wilson-loop spectrum. So by observing the evolution of surface-state
Fermi arc in the periodic process we can get the topology of Wilson-loop
spectrum.

Bulk-boundary correspondence is a defining feature of topological state of
matters. The bulk bands of topological matters are characterized by
topological invariants and there are nontrivial boundary states corresponding
to the nontrivial bulk bands[1, 2, 3]. Wilson-loop serves as a bridge between
the bulk property and the boundary property: the bulk topological invariant
and the boundary state property can both be derived from the Wilson-loop[4, 5,
6, 7].

The phases of eigenvalues the Wilson-loop along one direction are proportional
to the centers of the maximumly localized Wannier functions[8]. In 2D or 3D
sytems the Wannier functions are hybrid[6], that is, they are localized in one
direction and extended in other directions. By studying the Wannier centers as
function of $\mathbf{k}_{||}$ we can get the topological information of the
bulk bands[4, 5, 6, 9], where $\mathbf{k}_{||}$ is the wave vector in the
direction that the Wannier functions are extended. In this letter the Wannier
centers as functions of $\mathbf{k}_{||}$ are referred to as Wilson-loop
spectrums.

On the other hand, spectrum of Wilson-loop can be continuously deformed to
boundary state spectrum[7]. So the bulk-boundary correspondence can be
established by studying the topology of Wilson-loop spectrum.

Despite of its great versatility in characterizing topological property of
matters, up to now the Wilson-loop exists only in theory. That is, the Wilson-
loops are only used as auxiliary tools either for computing the topological
invariant[4, 5, 6] or for characterizing the topology of boundary states[7].
No attempt has been made to measure the topology of Wilson-loop.

For all the applications of Wilson-loop, the knowledge of the topology of
Wilson-loop spectrum is enough[4, 5, 6, 7]. So we can obtain the topology of
Wilson-loop by measuring something that is topologically equivalent to the
Wilson-loop spectrum. In ref.[7] it is shown that the boundary-state spectrum
is a continuous deformation of Wilson-loop spectrum if the bulk system is
gapped. Thus it is natural to obtain the topology of Wilson-loop spectrum by
measuring the boundary-state spectrum. However, as we show in supplementary
material[10], the boundary-state spectrum is only a partial deformation of the
Wilson-loop spectrum in general and this partial information from the
boundary-state spectrum is sometime misleading. In the topological semimetals,
where the bulk system is not gapped, the surface state spectrum is not a
partial deformation of Wilson-loop spectrum at all.

Another defect of obtaining topological information of Wilson-loop from
boundary state spectrum is that boundary state spectrum can not reflect the
periodic feature of Wilson-loop spectrum. Wilson-loop spectrum is defined
modulo a lattice constant. So it is periodic in the direction that the Wannier
functions are localized. When Wilson-loop spectrum deforms to boundary state
spectrum this periodic feature disappears. So we need a systematic way to
measure the topology of Wilson-loop spectrum.

Figure 1: (a) and (c) Brillouin zone of Weyl semimetal and Dirac semimetal.
The spheres in the cylinders are Weyl points and Dirac points. (b) and (d)
Behavior of the Wilson-loop spectrums of topological semimetals near the Weyl
and Dirac points. The intersections of the Wilson-loop spectrums and the
cylinders are highlighted by yellow curves.

Our work is inspired by the recent interests in surface-state Fermi arc of
topological matters, especially, of topological Weyl and Dirac semimetals[11,
12, 13, 14, 15, 16, 17]. The surface-state Fermi arcs have been observed in
both electronic[11, 12] and artificial periodic structures, such as phononic
crystals and photonic crystals[13, 14, 15, 16, 17].

The spectrum of surface states is sensitive to boundary conditions[18]. So the
surface-state Fermi arc is also boundary-dependent[19]. This boundary-
dependence of Fermi arc has been used to realize the topological negative
refraction of surface acoustic waves in a Weyl phononic crystal[19]. When the
boundary of the system varies continuously with a parameter, the Fermi arc
changes with the parameter and thus generates a surface.

If the system is half infinite, when one layer of unit cells are removed from
the system in one continuous process, the process can be considered as
periodic. In supplemental material[10] we show the evolution of Fermi arc in
such periodic process parameterized by $\phi_{d}$ generates a surface that is
topologically equivalent to the Wilson-loop spectrum. We call this surface the
“boundary Fermi surface” in following discussions.

The topological equivalence between Wilson-loop spectrum and “boundary Fermi
surface” is based on the fact that the bulk system is gapped. However, in Weyl
or Dirac semimetals the bulk gap closes at the Weyl or Dirac points. Here we
show the topological equivalence is not affected by the presence of Weyl or
Dirac points.

To establish the topological equivalence we divided the Brillouin zone of the
topological semimetal into two parts. One part consists of the open cylinders
that include the Weyl or Dirac points as is show in Fig.1 (a) and (c). The
rest of Brillouin zone constitutes the second part. Clearly in the second part
the bulk energy spectrum is gapped. So topological equivalence can be
established in this part by using the arguments in supplementary material[10].
Then we let the radius of the cylinders approaches zero. In this process the
topological equivalence does not change because it does not depends on the
size of bulk band gap. In this way we prove the topological equivalence in the
whole Brillouin zone.

One interesting phenomenon is that the Wilson-loop spectrum approaches to
vertical lines at projection points the Weyl or Dirac points. For example, if
the cylinder encloses a Weyl point with monopole charge $q$, the integral of
the Berry flux over the surface of the cylinder is $q2\pi$. If we consider the
surface of one cylinder as a 2D system, the system will have a Chern number
$q$. It is well know the Wilson-loop spectrum of a system with nonzero Chern
number is gapless. So when the radiuses of the cylinders approach to zero the
Wilon-loop spectrum will approach to vertical lines as is shown in Fig.1 (b)
and (d). To the best of our knowledge, no proposal have been made to use this
topological feature to characterize the topological semimetals.

Figure 2:  (a)Wilson-loop spectrum along the $y$ direction of the Weyl
semitmetal of Model (1). (b)“Boundary Fermi surface” of the same weyl
smeimetal. It is generated by the Fermi arcs of the surface states on the $XZ$
surface at Fermi energy $E_{F}=0$. (c) and (d) Surface state spectrum of the
this model with two different boundary conditions. $\phi_{d}=0$ in (c) and
$\phi_{d}=0.25$ in (d). The red surfaces are the surface state spectrums. The
yellow curves are Fermi arc of surface states at $E_{F}=0$. The semi-
transparent surfaces correspond to the bottom of upper bulk bands and top of
the lower bulk bands. They form Dirac cones at the neighbourhood of the
project points of the Wyel points. (e) and (f) Wilson-loop spectrum and
“boundary Fermi surface” of the continuous model (2).
$\hbar=m=V_{1}=V_{2}/2=1$ in the calculation. In (f) the Fermi energy equals
to the energy at the Weyl points.

To illustrate our conclusion we use the simplest tight-binding model for Weyl
semimetal

$\begin{split}H=&\sin k_{x}\sigma_{x}+\sin k_{y}\sigma_{y}\\\ &+(M_{0}+\cos
k_{x}+\cos k_{y}+\cos k_{z})\sigma_{z},\end{split}$ (1)

where $\sigma_{i}$, $i=x,y,z$, are the Pauli matrices. There are two Weyl
points at $\mathbf{k}=(\pi,\pi,\pi\pm\pi/2)$ when $M_{0}=2$. In Fig.2 (a) and
(b) we show the topological equivalence between Wilson-loop spectrum and
“boundary Fermi surface” of the Weyl semimetal. Especially, we show the
“boundary Fermi surface” approaches to vertical lines at the projection points
of the Weyl points. In comparison, the surface state spectrums in Fig.2 (c)
and (d) fail to reflect this topological feature of the Wilson-loop spectrum.
It is easily seen from Fig.2 (a), (c) and (d) the surface state spectrums of
the Weyl semimetal are not partial deformations of the Wilson-loop spectrum.

Figure 3:  (a) Wilson-loop spectrum along the $z$ direction of the 2D TI.
$M=2$, $t=\Delta_{1}=1$ and $\Delta_{2}=0$ in the calculations. (b) “Boundary
Fermi surface” of the 2D TI. (c) Wilson-loop spectrum along the $z$ direction
of a 2D trivial insulator which can exhibit gapless edge states with some
boundary conditions. (d) “Boundary Fermi surface” of the trivial insulator.
(e) and (f), Edge state spectrums of the trivial insulator with different
boundary conditions. In (e)((f)) $\phi_{d}=0$($\phi_{d}=0.75$) and the edge
state spectrum is gapless(gapped).

It is technically challenging to continuously remove one layer of unit cells
from the electronic system without breaking the translational symmetry in the
parallel direction. However, in classical phononic crystals and photonic
crystals, the unit cells are of macro scale and the boundary of the system can
be easily controlled.

We can remove one layer of unit cells from these artificial structures by
cutting the unit cells slice by slice. The surface-state Fermi arc of the
system is measured when each slice is cut off. Such a process can be
considered as an approximation of the continuous process we propose. The
thickness(in unit of lattice constant) of the layer that is cut off can be
chosen as the parameter $\phi_{d}$. So we can use the exiting experimental
setup for measuring surface-state Fermi arc to detect the topology of Wilson-
loop spectrum of the artificial structures. The only extra effort is to cut
the unit cells, which is particularly easy for phononic crystals because the
unit cells are fabricated by 3D printing[15, 16, 17].

Figure 4:  (a) Wilson-loop spectrum along the $z$ direction of the 3D TI. The
insets indicated the gaplessness of the Wilson-loop spectrum. The parameters
take the same value as in the 2D TI. (b) “Boundary Fermi surface” of the same
TI. The insets show it reflects the gaplessness of Wilson-loop spectrum. (c)
and (d) Surface state spectrum of the same bulk TI with two different
boundaries. $\phi_{d}=0$ in (c) and $\phi_{d}=0.75$ in (d). The red surfaces
are the surface state spectrums. The yellow cures are Fermi arcs of surface
states at $E_{F}=0$. In (c) the Fermi arc reduces to a point. The semi-
transparent surfaces also correspond to the bottom of upper bulk bands and top
of the lower bulk bands. They are gapped in 3D TI.

Since we predict the topology of Wilson-loop spectrum can be observed in
phononic crystal and since the rigid boundary conditions are used when Fermi
arc is measured in phononic crystal[15, 16, 17], it is more convincing to use
a continuous model and rigid boundary condition to show the topological
equivalence between Wilson-loop spectrum and “boundary Fermi surface”. We use
a 1D quantum model with two parameters to simulate a 3D Wyel semimetal. The
Hamiltonian of the semimetal can be expressed as

$\begin{split}H=&-\frac{\hbar^{2}}{2m}\frac{d^{2}}{dx^{2}}+V_{1}\cos(2\pi
x/a)\\\ &+V_{2}\sin^{2}(k_{z})\cos(2\pi x/a-k_{y}),\end{split}$ (2)

where a is lattice constant, $V_{1}$ and $V_{2}$ are constants. $k_{y}$ and
$k_{z}$ are two parameters to map the system to a 3D system. When
$|V_{2}|>|V_{1}|>0$ it represents a Weyl semimetal. The rigid boundary
condition can be realized by choosing $\psi(\phi_{d}a)=0$. If the system is on
the right side of $x=\phi_{d}a$, one layer of unit cells will be continuously
removed from the system when $\phi_{d}$ continuously increases by one. In
Fig.2 (e) and (f) we show the topological equivalence between Wilson-loop
spectrum and “boundary Fermi surface” of this model.

The topology of Wilson-loop spectrums of 2D and 3D topological insulators(TIs)
can be similarly measured. The models for 2D and 3D TIs used in this letter
are all adapted from the model for chiral high order topological
insulator(HOTI) in ref.20

$\begin{split}H(\mathbf{k})=&(M+t\sum_{i}\cos
k_{i})\tau_{z}\sigma_{0}+\Delta_{1}\sum_{i}\sin k_{i}\tau_{x}\sigma_{i}\\\
&+\Delta_{2}(\cos k_{x}-\cos k_{y})\tau_{y}\sigma_{0},\end{split}$ (3)

where $\sigma_{i}$ and $\tau_{i}$,$i=x,y,z$, are Pauli matrices acting on spin
and orbital degree of freedoms, respectively. When $\Delta_{2}=0$ the time-
reversal symmetry breaking term disappears and it is a model for strong 3D TI.
The model for 2D TI can be obtained by setting $k_{x}=\pi$ with
$\Delta_{2}=0$.

In Fig.3 and Fig.4 we show the topological equivalence between “boundary Fermi
surfaces” and Wilson-loop spectrums of the TIs. In the 2D systems the
“surface-state Fermi arcs” are points in 1D Brillouin zone. The evolution of
these points generate curves that are topologically equivalent to the Wilson-
loop spectrum.

If the bulk bands are gapped, the spectrums of edge or surface states are
partial deformation of Wilson-loop spectrums as is shown in Fig.3 (e) and (f)
and Fig.4 (c) and (d). However, this partial information from boundary-state
spectrum can be misleading as is indicated in Fig.2 (c) and (e). In Fig.2 (c)
we show the Wilson-loop spectrum of a 2D trivial insulator. The model for this
trivial insulator is constructed by setting $k_{x}=\pi$ with $\Delta_{2}=1$ in
(3). We can see the Wilson-loop spectrum of this model is gapped. However, in
Fig.3 (e) and (f) we show that the edge state spectrums can be both gapless
and gapped depending on the choice of boundary conditions. In contrast, the
“boundary Fermi surface” in Fig.3 (d) contains all the topological information
of the Wilson-loop spectrum.

The topology of Wilson-loop spectrum of HOTIs[21, 22, 20, 23, 24, 25, 26] can
also be measured by observing the evolution of surface-state Fermi arc.
However, the Wilson-loop spectrums of HOTIs are gapped as that of trivial
topological insulators. So Wilson-loop spectrum can not be used to
characterize the topology of HOTIs. The HOTIs distinguish themselves from
trivial insulators by the fact that the Wannier bands carry nontrivial
topological invariant. The topological invariant of Wannier band can be
obtained from the topology of nested Wilson-loop spectrum. The topology of
nested Wilson-loop spectrum can also be observed if we measure the boundary
wave functions in the periodic process measuring the “boundary Fermi
surfaces”[27].

## References

  * Kane and Mele [2005] C. L. Kane and E. J. Mele, Phys. Rev. Lett. 95, 146802 (2005).
  * Fu  _et al._ [2007] L. Fu, C. L. Kane, and E. J. Mele, Phys. Rev. Lett. 98, 106803 (2007).
  * Qi  _et al._ [2008] X.-L. Qi, T. L. Hughes, and S.-C. Zhang, Phys. Rev. B 78, 195424 (2008).
  * Yu  _et al._ [2011] R. Yu, X. L. Qi, A. Bernevig, Z. Fang, and X. Dai, Phys. Rev. B 84, 075119 (2011).
  * Soluyanov and Vanderbilt [2011a] A. A. Soluyanov and D. Vanderbilt, Phys. Rev. B 83, 235401 (2011a).
  * Soluyanov and Vanderbilt [2011b] A. A. Soluyanov and D. Vanderbilt, Phys. Rev. B 83, 035108 (2011b).
  * Fidkowski  _et al._ [2011] L. Fidkowski, T. S. Jackson, and I. Klich, Phys. Rev. Lett. 107, 036601 (2011).
  * Marzari and Vanderbilt [1997] N. Marzari and D. Vanderbilt, Phys. Rev. B 56, 12847 (1997).
  * Alexandradinata  _et al._ [2014] A. Alexandradinata, X. Dai, and B. A. Bernevig, Phys. Rev. B 89, 155114 (2014).
  * [10] See supplementary material .
  * Xu  _et al._ [2015a] S.-Y. Xu, C. Liu, S. K. Kushwaha, R. Sankar, J. W. Krizan, I. Belopolski, M. Neupane, G. Bian, N. Alidoust, T.-R. Chang, H.-T. Jeng, C.-Y. Huang, W.-F. Tsai, H. Lin, P. P. Shibayev, F.-C. Chou, R. J. Cava, and M. Z. Hasan, Science 347, 294 (2015a), https://science.sciencemag.org/content/347/6219/294.full.pdf .
  * Xu  _et al._ [2015b] S.-Y. Xu, I. Belopolski, N. Alidoust, M. Neupane, G. Bian, C. Zhang, R. Sankar, G. Chang, Z. Yuan, C.-C. Lee, S.-M. Huang, H. Zheng, J. Ma, D. S. Sanchez, B. Wang, A. Bansil, F. Chou, P. P. Shibayev, H. Lin, S. Jia, and M. Z. Hasan, Science 349, 613 (2015b), https://science.sciencemag.org/content/349/6248/613.full.pdf .
  * Yang  _et al._ [2017] B. Yang, Q. Guo, B. Tremain, L. E. Barr, W. Gao, H. Liu, B. Beri, Y. Xiang, D. Fan, A. P. Hibbins, _et al._ , Nature Communications 8, 97 (2017).
  * Noh  _et al._ [2017] J. Noh, H. Sheng, D. Leykam, Y. D. Chong, K. P. Chen, and M. C. Rechtsman, Nature Physics 13 (2017).
  * Li  _et al._ [2018] F. Li, X. Huang, J. Lu, J. Ma, and Z. Liu, Nature Physics 14, 30 (2018).
  * Ge  _et al._ [2018] H. Ge, X. Ni, Y. Tian, S. K. Gupta, M.-H. Lu, X. Lin, W.-D. Huang, C. T. Chan, and Y.-F. Chen, Phys. Rev. Applied 10, 014017 (2018).
  * Xie  _et al._ [2019] B. Xie, H. Liu, H. Cheng, Z. Liu, S. Chen, and J. Tian, Phys. Rev. Lett. 122, 104302 (2019).
  * He  _et al._ [2010] C. He, M.-H. Lu, W.-W. Wan, X.-F. Li, and Y.-F. Chen, Solid State Communications 150, 1976 (2010).
  * He  _et al._ [2018] H. He, C. Qiu, L. Ye, X. Cai, X. Fan, M. Ke, F. Zhang, and Z. Liu, Nature 560, 61 (2018).
  * Schindler  _et al._ [2018] F. Schindler, A. M. Cook, M. G. Vergniory, Z. Wang, S. S. P. Parkin, B. A. Bernevig, and T. Neupert, Science Advances 4 (2018).
  * Benalcazar  _et al._ [2017a] W. A. Benalcazar, B. A. Bernevig, and T. L. Hughes, Science 357, 61 (2017a), https://science.sciencemag.org/content/357/6346/61.full.pdf .
  * Benalcazar  _et al._ [2017b] W. A. Benalcazar, B. A. Bernevig, and T. L. Hughes, Phys. Rev. B 96, 245115 (2017b).
  * Imhof  _et al._ [2018] S. Imhof, C. Berger, F. Bayer, J. Brehm, L. W. Molenkamp, T. Kiessling, F. Schindler, C. H. Lee, M. Greiter, T. Neupert, _et al._ , Nature Physics 14, 925 (2018).
  * Serragarcia  _et al._ [2018] M. Serragarcia, V. Peri, R. Susstrunk, O. R. Bilal, T. Larsen, L. G. Villanueva, and S. D. Huber, Nature 555, 342 (2018).
  * Zhang  _et al._ [2019] X. Zhang, H. Wang, Z. Lin, Y. Tian, B. Xie, M. Lu, Y. Chen, and J. Jiang, Nature Physics 15, 582 (2019).
  * Luo and Zhang [2019] X.-W. Luo and C. Zhang, Phys. Rev. Lett. 123, 073601 (2019).
  * [27] To be disscussed later .


# Diffusion-HMC: Parameter Inference with Diffusion Model driven Hamiltonian
Monte Carlo

Nayantara Mudur  Department of Physics, Harvard University, 17 Oxford Street,
Cambridge, MA 02138, USA  Harvard-Smithsonian Center for Astrophysics, 60
Garden Street, Cambridge, MA 02138, USA  Carolina Cuesta-Lazaro  Harvard-
Smithsonian Center for Astrophysics, 60 Garden Street, Cambridge, MA 02138,
USA  The NSF AI Institute for Artificial Intelligence and Fundamental
Interactions Massachusetts Institute of Technology, Cambridge, MA 02139, USA
Department of Physics, Massachusetts Institute of Technology, Cambridge, MA
02139, USA  Douglas P. Finkbeiner  Department of Physics, Harvard University,
17 Oxford Street, Cambridge, MA 02138, USA  Harvard-Smithsonian Center for
Astrophysics, 60 Garden Street, Cambridge, MA 02138, USA

###### 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 the addition of noise than baseline parameter inference networks.

##  I Introduction

Ongoing and upcoming missions, such as the Dark Energy Spectroscopic
Instrument DESI, the Vera C. Rubin Observatory’s Legacy Survey of Space and
Time and the Nancy Grace Roman Space Telescope will map the cosmos at
unprecedented resolution and volume. This has created a proportionate demand
for simulations that can generate predictions from theory. Cosmological
simulations, however, are expensive to run, and can only be generated for a
limited set of initial conditions and points in parameter space.

The canonical summary statistic used for parameter inference is the two point
correlation function, or the power spectrum $Pk$ at large, linear scales where
perturbation theory holds. At smaller scales, however, gravitational collapse
induces non-Gaussianity in the fields. This means the information content of
cosmological fields is not fully captured by the power spectrum at the large
scales alone. For example, [11] derived 27% stronger constraints of
$\sigma_{8}$ from the two point correlation function of galaxy clustering data
by analyzing non-linear modes at smaller scales, ($k>$0.25 $h\rm{Mpc}^{-1}$).
A multitude of other statistics – such as the marked power spectrum, the
bispectrum, the wavelet scattering transform, and void probability functions –
have also been devised [44, 28, 32, 12] in an effort to capture higher order
correlations in non-Gaussian fields. Previous work [14, 23, 28, 45, 38] has
addressed the prohibitive cost of simulations by creating emulators or
surrogate models that learn to interpolate between predictions of a specific
summary statistic between training points using formalisms such as Gaussian
processes. More recently, simulation-based inference at the field level has
also been used, where a likelihood model parameterized by a neural network is
learned on the fields [9, 8].

Generative models are a class of machine learning approaches that enable one
to simulate the ability to draw samples from a complicated target probability
density and include variational autoencoders, normalizing flows and generative
adversarial networks (GANs). Diffusion generative models [41] involve a
forward diffusion (noising) process that transform samples from the target
distribution to those from the standard normal. In the denoising diffusion
probabilistic model (DDPM)[16], the noising process consists of a variance
schedule ${\beta_{t}}$ over a fixed number of time steps, $T$, that determines
the incremental noise added to the image. Since the diffusion process can be
formulated as a stochastic differential equation (SDE) the DDPM variance
schedule corresponds to the discretization of this SDE. In the generative
direction, a neural network or score model is used to parameterize the reverse
transformation.

Diffusion models are alternatively referred to as a score-based generative
model since parameterizing the reverse diffusion process is equivalent to to
learning the ‘score’ or $\vec{\nabla}\log p_{t}(x)$ of the data [1, 41]. Since
in high dimensions, the target distribution invariably lies on a thin
manifold, the incremental addition of the random noise, blurs the distribution
and makes the score progressively easier to learn. Diffusion models are the
underlying mathematical framework that have given rise to the photorealistic
image generation successes of DALL.E and Stable Diffusion [34] and have been
shown to mitigate mode collapse, a phenomenon, often encountered with GANs a
generative model fails to generate multiple modes in a distribution. In
scientific applications, they have used for problems involving protein folding
and ligand prediction and medical imaging reconstruction [7, 40] . They have
been applied in astrophysics to reconstruction problems involving dust [15,
22], cosmological simulations and initial conditions reconstruction [20, 36,
27, 8] and strong lensing problems [18, 33].

In this work, we apply diffusion generative models to emulating cold dark
matter density fields conditional on cosmological parameters and demonstrate
that the trained model can also be used to derive tight and robust constraints
on cosmological parameters. In Section III, we examine the ability of the
model to appropriately capture the statistics of the distribution of fields
corresponding to different parameters, and further compare the effect of
modulating a single parameter at a time on the statistics of the resulting
fields in the true and the generated set. We then quantify the ability of the
model to capture the full range of cosmic variance for a single parameter, as
a means to assess the extent of mode collapse. In Section IV, we examine how
the diffusion model’s approximate likelihood can be used to solve the inverse
problem of constraining the cosmological parameters of a given input field. We
then use Hamiltonian Monte Carlo [10, 24, 3] to draw samples from the
estimated posterior on the cosmological parameters given an input field, and
compare our estimates with a power spectrum baseline. A novel contribution of
this work is our use of an HMC to sample a posterior consisting of an
approximation to the diffusion model conditional likelihood to solve a
downstream inference task. Finally, we demonstrate that the Diffusion-HMC
based parameter inference estimates are more robust to the addition of
uncorrelated noise relative to the estimates from a discriminative neural
network directly trained to estimate parameters.

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.

##  II Datasets, Architecture and Training

##### Datasets

We work with cold Dark Matter density fields at $z=0$ from the IllustrisTNG
[25, 30] suite from the CAMELS Multifield Dataset (CMD) [49, 47]. The dataset
contains of 1000 simulations for 1000 different cosmologies with 15 two-
dimensional fields per simulation. The ‘cosmology’ is parameterized by a
parameter vector with 2 cosmological ($\Omega_{m}$ and $\sigma_{8}$) and 4
astrophysical parameters. Each simulation tracked the evolution of $256^{3}$
dark matter particles and $256^{3}$ fluid elements and took around 6000 CPU
hours to generate (see [49, 47] for more details). The fields span $25$
Mpc.$h^{-1}$ on each side. We train on 70% of the parameters in the LH set,
i.e. 700 parameters or 10,500 fields, and condition the model only on the
cosmological parameters $\Omega_{m}$ and $\sigma_{8}$. Example generated dark
matter fields are shown in Figure 2.

##### Diffusion Model Setup

We follow the denoising diffusion probabilistic model (DDPM) formalism, in
which a target image $x_{0}$ is transformed to a sample from
$x_{T}\sim\mathcal{N}(0,\mathbb{I})$ over the course of $T=1000$ timesteps.
The forward diffusion process follows an incremental noise schedule
{$\beta_{t}$}, and the noise is added in a variance-preserving way.

$\displaystyle\textrm{For
t}\in[0,T-1],q(x_{t+1}|x_{t})=\mathcal{N}(\sqrt{1-\beta_{t}}x_{t},\beta_{t}\mathbb{I})$
$\displaystyle\textrm{ and
}q(x_{t+1}|x_{0})=\mathcal{N}(\sqrt{\bar{\alpha_{t}}}x_{0},(1-\bar{\alpha_{t}})\mathbb{I})$
$\displaystyle\text{ where
}\bar{\alpha_{t}}=\prod_{t^{\prime}=0}^{t}1-\beta_{t^{\prime}}$ (1)

The score / noise-predictor model is a U-Net [35] similar to that used in
[16]. There are 4 downsampling layers, with each layer consisting of 2 ResNet
blocks [53], group-normalization [52], and attention [46, 39]. We use circular
convolutions in the downsampling layers since the input fields have periodic
boundary conditions. Each parameter is normalized to lie between [0, 1] with
respect to its range, $\Omega_{m}\in[0.1,0.5],\sigma_{8}\in[0.6,1.0]$. A
multilayer perceptron (MLP)transforms the cosmology vector into a space with
the same dimension as the time embedding, and each ResNet block additionally
has an MLP conditional on cosmology. The variance schedule $\beta_{t}$ is non-
linear with smaller steps at smaller $t$ and larger steps for larger values of
$t$ (see Figure 8). We train the model to generate the log of the fields, and
randomly rotate and flip the image to account for these invariances. During
training, for each batch of images $\mathbf{x}_{0}$, a batch of timesteps is
sampled uniformly along with a noise pattern
$\epsilon\sim\mathcal{N}(0,\mathbf{I})$. The loss function minimized is
$||\mathbf{\epsilon}-\mathbf{\epsilon}_{\phi}(\sqrt{\bar{\alpha}_{t}}\mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}}\mathbf{\epsilon},t,\theta))||^{2}$
for each set of $\\{\mathbf{x}_{0},t,\theta,\epsilon\\}$, where $\theta$ is
the parameter vector. In practice, we minimize the Huber loss, which behaves
as an L1 loss for larger values and L2 for smaller values of the loss. We used
the Weights and Biases framework [4] to track experiments. The model has 31.2
million parameters.

##### Training

We first downsample the images by a factor of 4 and train the conditional
diffusion model on these 64x64 images for 60000 iterations. Since the UNet is
formulated in terms of relative downsampling, the same architecture can be
applied to images with different resolution. To train the model to emulate
256x256 fields, we initialize the checkpoint with the weights of the
64x64-trained model after 60000 iterations and trained for over 400000
iterations. We found that initializing the 256x256 model with the weights of
the 64x64 model and then training the model led to faster convergence.

The noise prediction loss does not fully capture sample quality and
convergence, and we need an alternative metric to assess the quality of the
generated samples [43]. We sampled 500 fields (with 50 fields for 10 different
validation parameters) for the checkpoints after 200k, 220k, 240k, 260k, 280k,
300k, 320k and 340k iterations. We computed the reduced chi-squared statistic
(Equation 10) of the power spectrum of each generated field, $s$, with respect
to the reference distribution comprised of the 15 true fields for that
parameter. We then compute the mean and standard error of these values across
all parameters and sampled fields.

While the diffusion generative model is trained to generate the fields in log
space, the power spectra we compute here and in Figures 1\- 3 correspond to
the overdensity power spectra of the ‘linear’ ($10^{\rm{GeneratedFields}}$)
fields. The checkpoint corresponding to the 260kth iteration had the lowest
value, corresponding to $2.29\pm 0.49$. To put this number in perspective, we
can examine the effect of cosmic variance on this metric using a leave-one-out
cross-validation approach, by computing the reduced chi-squared statistic of
each sample of a true field, using the 14 other true fields corresponding to
the same parameter as the reference distribution. The mean of this value
across the 10 parameters is $1.70\pm 0.36$. We use the 260k checkpoint for our
analysis. We plot these values, along with the reduced chi-squared statistic
of the power spectra of the log fields and the $p$ values of the mean
intensity in Appendix .3.1.

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.

##  III Summary Statistics

In this section we examine the consistency of our generated fields relative to
the true fields using three sets of simulations under the IllustrisTNG CAMELS
suite. The latinhypercube, LH, suite varies the largest range of cosmological
and astrophysical parameters but has a limited number (15) of fields for each
parameter, with all parameters varying randomly. The one-parameter, 1P, suite
consists of 15 fields with the same seed, and one dimensional variations of
each parameter modulated systematically, while the others are fixed to the
fiducial value. The Cosmic Variance, CV, set consists of 405 fields for the
fiducial parameter value. Sampling a batch of 50 $256\times 256$ fields from
our model takes 310 seconds (6.2s/field).

Varying cosmology in a Latin Hypercube (LH):  We examine the consistency of
the summary statistics of the distribution of true and generated fields for a
given validation parameter from the LH set in Figure 1. We have 15 true fields
and 50 generated fields for each parameter. We derived the boundaries of the
envelope using the estimates of the $16^{th}$ and $84^{th}$ percentiles, while
the solid line demarcates the mean of the distribution of power spectra in 35
log spaced $k$ bins. The lower panel depicts the density histograms of the log
fields. The envelopes are again derived using the percentiles, while the solid
lines indicate the means of the histograms for the true and generated fields.
The distribution of the power spectra and the density histograms of the true
and the generated fields are in good agreement with each other.

Varying cosmological parameters one at a time (1P): In Figure 2, we generated
‘1P’ sets and examined whether the effect of modulating a single parameter,
while keeping the others constant is the same as is observed in the 1P CAMELS
suite. We sample 15 fields corresponding to 15 different seeds for each of the
parameters. The fields corresponding to the same seed across parameters have
the same position and orientation of their seeded structures as visible in
Figure 2. For each seed, we then compute the ratio of the power spectrum of a
field at a different parameter to the power spectrum of the fiducial parameter
value, and compute the average and the percentile-based standard deviation
envelopes across all seeds as depicted in the right column of Figure 2. The
dashed (solid) line and envelope correspond to the generated (true) ratios.
The ratios for the generated ‘1P’ set are in good agreement with the those of
the true ‘1P’ set, since the dashed lines are within the solid envelopes.

Reproducing cosmic variance (CV): The CV set has 405 ($27\times 15$) fields
for the fiducial parameter value of [0.3, 0.8] and varying initial conditions,
designed to quantify the effect of cosmic variance. The CV set allows us to
quantify the consistency between the second moments of the true and the
generated distributions. In particular, it allows us to test the ability of
the model to generate a diverse set of samples for the same cosmological
parameters such that it reproduces the true underlying distribution at fixed
cosmology.

We generate 405 samples from our trained diffusion model, compute their power
spectra, and examine the standard deviation, and the correlation matrices of
different $k$ modes of the power spectrum in Figures 3 and 4. The correlation
between the modes of the power spectra is largely consistent between the true
and the generated samples, although the generated samples appear to have a
slight excess correlation around 5 $h\rm{Mpc}^{-1}$.

The standard deviations of the distribution are also consistent although the
standard deviations of the generated power spectra are slightly underpredicted
relative to the true power spectra at the largest length scales
($k<5$$h\rm{Mpc}^{-1}$). To construct the standard deviation estimator, we use
the jackknife-approach to generate a distribution of 450 estimates of the
standard deviation for each set, where the $\rm{i^{th}}$ estimate corresponds
to the standard deviation computed using all samples excluding that of the
$\rm{i^{th}}$ field. We then compute the mean and the standard deviation for
these sets in order to capture the mean and the standard deviation on the
estimate of the standard deviation in the upper right panel of Figure 3. We
use the ratios of the jackknife-estimated means and errors in the lower right
panel.

The ability to capture the full diversity of $k$ modes for a single parameter
may be in tension with the ability to distinguish between and appropriately
modulate the power spectrum for different parameters. This tension is enhanced
for our assessment of mode collapse in spectral space, compared to canonical
machine learning datasets involving discrete classes, since the cosmological
parameters that we condition the diffusion generative model on lie on a
continuum. Thus, for a given conditioning parameter, a generated field with
too much or too little power on certain scales relative to the mean of the
statistic is more likely to wander into the typical set of a field with a
different cosmological parameter since modulating the parameters also
modulates the power spectrum (as in Figure 2). In the context of natural
images, one often deals with categorical descriptors, where the boundaries
between whether an object qualifies as one class or another are typically more
clearly delineated, and the ability to capture the full diversity of cats is
unlikely to cause the model to ‘wander’ into islands of image space
corresponding to a dog or an airplane. Thus the slightly lower standard
deviation can be partly attributed to the possibility that the model chooses
to compromise on the diversity of samples for a single parameter, in order to
be able to accurately generate fields that look different for different
parameters.

Moreover, the high dimensionality of the initial conditions relative to our
training set size of 10,000 fields poses a challenge. Increasing the size of
the training dataset may help mitigate this issue.

We compute the covariance between the modes of the power spectrum. With the
full covariance we can now statistically quantify the consistency of the
generated samples relative to that of the true samples using the
multidimensional reduced chi-squared statistic. We use 350 samples of the true
fields to set up the reference distribution used to compute the covariance. We
compute the inverse covariance and adjust for the Hartlap factor (Equation 8)
[13]. We then compute the multidimensional reduced chi-squared statistic of
the entire sample of the true and generated fields, and compute the means of
the 405 chi-squared statistics for each distribution following Equation 9. For
the true fields, the mean of the chi-squared distribution is 33.4 while that
of the generated is 36.6. Since our power spectra consist of 35 log-spaced
bins, this is consistent with the expected mean for a chi-squared distribution
with 35 degrees of freedom, i.e. 35.

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.

##  IV Parameter Inference

A trained diffusion model can be used to estimate the variational lower bound
(VLB) of the log likelihood [19, 8]. In the case of a conditional diffusion
model, this likelihood estimate is also conditional, i.e.

$\displaystyle\mathbb{E}_{q}[-\log
p_{\phi}(x_{0}|\theta)]\leq\mathbb{E}_{q}[{\color[rgb]{0,0,0}-\log
p_{\phi}(x_{0}|x_{1},\theta)}+$ $\displaystyle{\color[rgb]{0,0,0}\sum_{t\geq
1}\rm{D_{KL}}[q(x_{t}|x_{t+1},x_{0})||p_{\phi}(x_{t}|x_{t+1},\theta)]}+$
$\displaystyle{\color[rgb]{0,0,0}\rm{D_{KL}}[q(x_{T}|x_{0})||p(x_{T})]}]$ (2)
$\displaystyle
L_{\rm{VLB}}={\color[rgb]{0,0,0}L_{0}}+{\color[rgb]{0,0,0}L_{1}...L_{T-1}}+{\color[rgb]{0,0,0}L_{T}}$
(3) $\displaystyle\hat{L}_{\rm{VLB}}=\sum_{t<T_{\rm{MAX}}}L_{t}$ (4)

where $\phi$ denotes the diffusion model architecture, $\theta$ is the
conditioning cosmology (in our case, a vector with $\Omega_{m}$ and
$\sigma_{8}$), $p_{\phi}$ are the reverse (learned) distributions, and $q$ are
the forward (analytical) distributions. During training, the diffusion model’s
noise prediction loss terms are equivalent to terms of the reweighted VLB
[19]. Since the predicted noise is conditional on cosmology, the terms of the
variational lower bound thus encode dependencies on the cosmological
parameters. The contrast between the VLB evaluated at one parameter
$\theta_{1}$ relative to another parameter $\theta_{2}$ for a fixed field
$x_{0}$ can thus be used to find the region of parameter space that maximizes
the conditional likelihood for the field.

###  IV.1 Examining the influence of different timesteps

Since using all the terms $L_{t}$ that contribute to the VLB is
computationally expensive, we first investigate how sensitive each of the
contributing terms $L_{t}$ is to changes in cosmological parameters over a
grid in Figure 5a). For an input field $x_{0}$ and a single seed, we evaluate
each of the $L_{t}(x_{0}|\vec{\theta}_{\textsc{Eval}})$ terms over a $50\times
50$ grid in $[\Omega_{m},\sigma_{8}]$, centered on the value of the true field
$\vec{\theta}_{\textsc{True}}$ and extending to $\pm 0.06$ about the true
parameter. To disentangle the individual contribution of each term, we
subtract the minimum value of each $L_{t}(x_{0}|\vec{\theta}_{\textsc{Eval}})$
on the grid, and multiply it by 2 to yield $-2\Delta ln\mathcal{\hat{L}}_{t}$.
We plot the contour corresponding to 2.30, or the one-sigma contour for a chi-
squared distribution with two degrees of freedom. All contours are minimized
in the vicinity of the true parameter, and the curvature of the contours
decreases with increasing time. Increasing time involves increasing the amount
of noise added to the input image, which can explain the increased uncertainty
in the true value of the parameter. Thus dropping the terms corresponding to
the higher timesteps in Equation 2 is unlikely to result in weaker constraints
on cosmology.

Figure 5: Investigating the contribution of the timesteps used in the VLB sum.
Left: One sigma contours of each $-2\Delta lnL_{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.

###  IV.2 HMC-based Parameter Inference

To compute the parameter estimates for fields, we draw samples from the
posterior on the parameter using the Hamiltonian Monte Carlo (HMC) method. The
Hamiltonian Monte Carlo, or Hybrid Monte Carlo method [10, 24, 3] draws
samples from a probability distribution $\pi(\theta)$ via the introduction of
an auxiliary momentum variable $\mathbf{p}$ and solves the equations of
Hamiltonian dynamics in order to update the momentum and position ($\theta$).
HMC enables a more efficient exploration of high dimensional probability
distributions. In our case, using an HMC also helps us circumvent the problem
of having to redefine the extents and granularity of the parameter grid
depending on how confident the constraints for a given $T_{\rm{MAX}}$ are. The
Hamiltonian governing the dynamics in the HMC chain is:

$\displaystyle H=U+K\text{ where }U=-\log\pi(\theta)\text{ and
}K=\frac{\bm{p}^{-1}\bm{M}^{-1}\bm{p}}{2}$ (5)
$\displaystyle\log\pi(\theta)=\log p_{\phi}(x_{0}|\theta)+\log
p_{\textsc{Prior}}(\theta)$ $\displaystyle\simeq-\hat{L}_{\rm{VLB}}+\log
p_{\textsc{Prior}}=-\sum_{t<T_{\rm{MAX}}}L_{t}+\log p_{\textsc{Prior}}$ (6)

We used the Hamiltorch [6] package for the HMC. We explain more details about
the HMC implementation in Appendix .1. The prior is chosen to be a flat prior
over $\Omega_{m}\in[0.1,0.5]$ and $\sigma_{8}\in[0.6,1.0]$. The initial
parameter value is always the fiducial value of $[0.3,0.8]$ and we designate
the first 100 samples as burn-in samples that are discarded.

We now examine the effect of truncating terms in Equation 4 using the HMC
based parameter inference. Truncating terms allows us to perform inference
faster and can allow us to explore the tradeoff between dropping terms for
speed and higher precision with more timesteps. In Figure 5b), for a single
field, we plot the mean predictions and the $15.9-84.1$ (1 sigma) percentiles
computed using 200 samples with the approximate $-\log p_{\phi}(x_{0}|\theta)$
using Equation 4 as a function of $T_{\rm{MAX}}$ on the x axis for the same
field. For this field and parameter, using more terms asymptotically removes
the bias on $\Omega_{m}$ while increasing the bias on $\sigma_{8}$. However,
using the first 20 timesteps only changes the mean prediction for $\Omega_{m}$
and $\sigma_{8}$ by $-0.26\%$ and $-1.27\%$ of the prediction using all 1000
timesteps respectively.

Figure 6: Top: Comparison of the parameter estimates for a single field with
the power spectrum based estimator and with the diffusion model likelihood.
The star demarcates the true parameter corresponding to the input fields. The
shaded contours demarcate the $68.27\%,95.45\%$ and $99.73\%$ confidence
intervals. Lower Panel: Predicted parameter and ground truth parameter for 10
different fields. The error bars correspond to the $15.9-84.1$ percentiles for
the marginal probability distributions for each parameter for each input
field.

We now turn our attention to the performance of our parameter inference
approach relative to a power spectrum baseline. For subsequent HMC-based
parameter inference in this section, we use Equation 4 with $T_{\rm{MAX}}=20$
to approximate the conditional negative log likelihood. Drawing 500 samples
with $T_{\rm{MAX}}=20$ takes $\sim 32$ minutes for a single field.

To assess the constraining power of the power spectrum, we use the neural
posterior estimation approach and train a normalizing flow to represent the
posterior of $p_{\rm{NF}}(\theta|Pk\cdot\rm{Mean})$ using the Lampe [37]
package. Since computing the overdensity power spectrum involves dividing by
the mean of the field, we further concatenate the mean of the log field as an
additional feature, since the prediction for $\Omega_{m}$ for a small box size
of $25$ Mpc.$h^{-1}$ is very sensitive to the mean of the fields. For a single
field, we compare the posteriors obtained by drawing 10000 samples from the
power spectrum based estimator and 400 samples from the diffusion model+HMC in
Figure 6. Other details about our implementation are in the Appendix .2.

The diffusion model has significantly narrower constraints for the parameters
relative to the power spectrum baseline. Note that, as shown in Figure 2, the
cosmological parameters $\Omega_{m}$ and $\sigma_{8}$ are strongly correlated
at the level of the power spectrum on the small scales probed by our
simulations, since they are both modulating the amplitude of the power
spectrum. This result is consistent with [50] who also found that the galaxy
power spectrum in conjunction with a multilayer perceptron yielded weak
constraints on $\Omega_{m}$ and $\sigma_{8}$.

In the second row, we plot the predicted cosmological parameters relative to
the truth for 10 different fields across different parameters in the
validation set. The dots annotate the mean of the 400 samples and the error
bars correspond to the $15.9-84.1$ percentiles of the samples. The mean of the
samples is close to the true value of $\vec{\theta}_{\textsc{True}}$ over a
broad range of parameters. The $\Omega_{m}$ predictions are biased by $-0.006$
$(2\%$ of the true prediction) on average over the range of parameters while
the mean absolute bias for $\sigma_{8}$ is 0.01 ($1.26\%$ of the true
prediction). The $\sigma_{8}$ prediction uncertainties (mean z score
$\bar{|z|}=1.14$) are better calibrated relative to the uncertainties for
$\Omega_{m}$ ($\bar{|z|}=2.75$), but overall, we find the error-bars to be
slightly under-predicted. We attribute the mis-calibration to the truncation
of terms. It is possible that including more timesteps, averaging over more
seeds and examining alternate choices of the variance schedule could yield
better calibrated uncertainties, we leave this exploration for future work.

Figure 7: Mean and $15.8-84.1$ percentile predictions using the Diffusion-HMC
estimates (green) and the mean and sigma predictions obtained from the
parameter inference network in [47], denoted by InfNet. The mean absolute
percentage bias (of the true prediction) for each parameter is also indicated.
The Diffusion-HMC constraints are more robust to the perturbations to the
input image.

####  IV.2.1 Robustness

Although our constraints are much tighter than those derived from the two
point correlation function, robustness to noise and observational systematics
is an important consideration guiding the use of parameter inference methods
on survey data. Since the terms of the VLB (and the noise prediction loss
terms) include terms where the original image has been noised, we hypothesize
that the parameter estimates learned by the model are naturally more robust to
perturbations involving the addition of scaled white noise to the field.

To test this, we examine the extent to which our parameter estimates change
relative to the predictions of the parameter inference network in [48]. The
neural network in [48] is trained to predict the mean and the standard
deviation of the parameters given an input dark matter field. In the leftmost
column of Figure 7, we compare the sample predictions obtained from the
diffusion model (green) relative to those obtained from the parameter
inference network (red) on 5 fields ($x_{0}$). In the next two columns, we add
increasing levels of $\mathcal{N}(0,\sigma)$ noise with
$\sigma\in[0.01,0.02]$, and examine the estimates for both approaches. We
perform this experiment in a noise-agnostic setting, i.e. we assume that we do
not know the level of noise a priori and do not modify the parameter inference
approach for the diffusion model. We use the same $T_{\rm{MAX}}=20$ and use
200 samples for the estimates.

##  V Conclusion

In this work, we used a diffusion generative model to emulate dark matter
density fields conditional on cosmological parameters. The model learns to
reproduce the modulation of summary statistics, such as the power spectra, for
changes in cosmological parameters. We additionally assess the extent of mode
collapse through the lens of cosmic variance by examining the diversity of the
power spectra at the fiducial cosmology.

We then directed our attention to the inverse problem of parameter inference
and disentangled the contribution of each term $L_{t}$ in the expression for
the Variational Lower Bound to constraints on cosmological parameters. Our
findings reveal that the strength of the constraints decrease as $t$
increases. [21] explored this question in the context of using the diffusion
model’s conditional VLB terms for classification and found that intermediate
timesteps had the highest accuracy when only a single timestep is used. The
difference in which terms contain the most information about the conditioning
attribute (cosmology / class) is interesting, and could be partly explained by
the difference in their formulation and weighting of the VLB terms and in how
changes in the conditioning attribute affect an image in the two settings.
Modulating cosmology modifies global attributes of the fields, such as their
intensity and power spectra, as seen in Figure 2, while the information that
distinguishes different breeds of pets from each other tends to be relatively
localized in the bounding box containing the animal.

If it is indeed always the case that the timesteps nearest to the image
manifold contain most cosmological information, one could in future, swap out
our discrete time architecture with continuous time diffusion models [41],
where $t$ is a continuous variable with $t\in[0,1]$ and prioritize steps that
lie near the image manifold or $t=0$.

These insights motivated us to truncate the diffusion model conditional VLB
based approximation for $p_{\phi}(x_{0}|\theta)$ by subsampling the terms to
only use the first $T_{\rm{MAX}}$ terms (Equation 4). This approximation
allows us to backpropagate its gradient and plug this estimate for
$p_{\phi}(x_{0}|\theta)$ into an HMC and sample the posterior on
$p_{\phi}(\theta|x_{0})$. The Diffusion-HMC approach yields tight constraints
on cosmological parameters, competitive with a bespoke parameter inference
network [48] trained only to infer cosmology given a field.

In our experiments, we use only a single seed in each HMC step and a
$T_{\rm{MAX}}=20$ to speed up inference. In our case, the use of an HMC
allowed us to circumvent the requirement of choosing a grid with the
appropriate resolution required to resolve the constraints, and eliminates the
dependence on the number of points in the grid. However, the Diffusion-VLB
estimates could also be used in other settings, that do not entail the use of
an HMC. While our approach scaled with the number of timesteps we use in the
sum, in constrast, [21] scaled with the number of classes in the
classification dataset.

Lastly, we demonstrate that the diffusion model likelihood confers the
Diffusion-HMC constraints with greater robustness to the addition of noise to
the input image relative to the behavior of a parameter inference network.
This echoes the behavior of diffusion models in other discriminative tasks
such as [21, 5, 31], where diffusion model based classifiers have been shown
to possess higher robustness to adversarial examples or perturbations. This is
a pertinent finding since [17] showed that the powerful constraints derived
from neural network based parameter inference may often come at the cost of
their susceptibility to slight perturbations that the canonical two-point
correlation based analyses are impervious to.

The simulation volume of the dataset we work with is still much smaller than
the scales mapped by astrophysical surveys. Future work could focus on scaling
up to more survey-realistic scenarios involving larger simulation volumes and
directly observed tracers. [8] also showed that diffusion generative models
that work with point clouds can allow one to emulate and perform cosmological
parameter inference with point cloud data. Our exploration into robustness
could inform applications to real data, and future investigation could focus
on conferring and quantifying robustness against other survey-related and
observational noise effects. Alternative formulations of the generative
process [2, 51] could also be relevant to this exercise.

In this work, we demonstrated that a diffusion model can be trained not just
to emulate fields, but that it’s likelihood estimate can be adapted to work
with the Hamiltonian Monte Carlo framework to derive tight constraints on
cosmological parameters. This makes a step toward advancing the use of
diffusion model based priors for a range of downstream tasks from image
generation and restoration to inference problems.

##  VI Code

The code for this project is available at Diffusion-HMC . We acknowledge [48],
Improved-Diffusion [26], The Annotated Diffusion and DDPM for use of their
code snippets.

Packages: Hamiltorch [6], Lampe [37], GNU Parallel [42].

##  VII Acknowledgements

This work was supported by the National Science Foundation under Cooperative
Agreement PHY2019786 (The NSF AI Institute for Artificial Intelligence and
Fundamental Interactions). We are grateful to Yueying Ni, Francisco
Villaescusa-Navarro, Core Francisco Park, Andrew K. Saydjari, Justina R. Yang,
Yilun Du, Ana Sofia Uzsoy, Shuchin Aeron and many others for insightful
conversations.

## References

  * Anderson [1982] Anderson, B. D. 1982, Stochastic Processes and their Applications, 12, 313
  * Bansal et al. [2024] Bansal, A., Borgnia, E., Chu, H.-M., et al. 2024, Advances in Neural Information Processing Systems, 36
  * Betancourt [2017] Betancourt, M. 2017, arXiv preprint arXiv:1701.02434
  * Biewald [2020] Biewald, L. 2020, Experiment Tracking with Weights and Biases.  https://www.wandb.com/
  * Chen et al. [2024] Chen, H., Dong, Y., Shao, S., et al. 2024, arXiv preprint arXiv:2402.02316
  * Cobb et al. [2019] Cobb, A. D., Baydin, A. G., Markham, A., & Roberts, S. J. 2019, arXiv preprint arXiv:1910.06243
  * Corso et al. [2022] Corso, G., Stärk, H., Jing, B., Barzilay, R., & Jaakkola, T. 2022, arXiv preprint arXiv:2210.01776
  * Cuesta-Lazaro & Mishra-Sharma [2023] Cuesta-Lazaro, C., & Mishra-Sharma, S. 2023, arXiv preprint arXiv:2311.17141
  * Dai & Seljak [2023] Dai, B., & Seljak, U. 2023, arXiv preprint arXiv:2306.04689
  * Duane et al. [1987] Duane, S., Kennedy, A. D., Pendleton, B. J., & Roweth, D. 1987, Physics letters B, 195, 216
  * Hahn et al. [2023] Hahn, C., Eickenberg, M., Ho, S., et al. 2023, Proceedings of the National Academy of Sciences, 120, e2218810120
  * Hamaus et al. [2016] Hamaus, N., Pisani, A., Sutter, P. M., et al. 2016, Physical Review Letters, 117, 091302
  * Hartlap et al. [2007] Hartlap, J., Simon, P., & Schneider, P. 2007, Astronomy & Astrophysics, 464, 399
  * Heitmann et al. [2009] Heitmann, K., Higdon, D., White, M., et al. 2009, Astrophys. J., 705, 156, doi: 10.1088/0004-637X/705/1/156
  * Heurtel-Depeiges et al. [2023] Heurtel-Depeiges, D., Burkhart, B., Ohana, R., & Blancard, B. R.-S. 2023, arXiv preprint arXiv:2310.16285
  * Ho et al. [2020] Ho, J., Jain, A., & Abbeel, P. 2020, Advances in Neural Information Processing Systems, 33, 6840
  * Horowitz & Melchior [2022] Horowitz, B., & Melchior, P. 2022, arXiv preprint arXiv:2211.14788
  * Jagvaral et al. [2022] Jagvaral, Y., Mandelbaum, R., & Lanusse, F. 2022, arXiv preprint arXiv:2212.05592
  * Kingma & Gao [2023] Kingma, D. P., & Gao, R. 2023, arXiv preprint arXiv:2303.00848
  * Legin et al. [2024] Legin, R., Ho, M., Lemos, P., et al. 2024, Monthly Notices of the Royal Astronomical Society: Letters, 527, L173
  * Li et al. [2023] Li, A. C., Prabhudesai, M., Duggal, S., Brown, E., & Pathak, D. 2023, in Proceedings of the IEEE/CVF International Conference on Computer Vision, 2206–2217
  * Mudur & Finkbeiner [2022] Mudur, N., & Finkbeiner, D. P. 2022, arXiv preprint arXiv:2211.12444
  * Mustafa et al. [2019] Mustafa, M., Bard, D., Bhimji, W., et al. 2019, Computational Astrophysics and Cosmology, 6, 1
  * Neal et al. [2011] Neal, R. M., et al. 2011, Handbook of markov chain monte carlo, 2, 2
  * Nelson et al. [2019] Nelson, D., Springel, V., Pillepich, A., et al. 2019, Computational Astrophysics and Cosmology, 6, 1
  * Nichol & Dhariwal [2021] Nichol, A. Q., & Dhariwal, P. 2021, in International Conference on Machine Learning, PMLR, 8162–8171
  * Ono et al. [2024] Ono, V., Park, C. F., Mudur, N., et al. 2024, arXiv preprint arXiv:2403.10648
  * Paillas et al. [2023] Paillas, E., Cuesta-Lazaro, C., Percival, W. J., et al. 2023, Cosmological constraints from density-split clustering in the BOSS CMASS galaxy sample.  https://arxiv.org/abs/2309.16541
  * Papamakarios et al. [2017] Papamakarios, G., Pavlakou, T., & Murray, I. 2017, Advances in neural information processing systems, 30
  * Pillepich et al. [2018] Pillepich, A., Springel, V., Nelson, D., et al. 2018, Monthly Notices of the Royal Astronomical Society, 473, 4077
  * Prabhudesai et al. [2023] Prabhudesai, M., Ke, T.-W., Li, A. C., Pathak, D., & Fragkiadaki, K. 2023, in Thirty-seventh Conference on Neural Information Processing Systems
  * Régaldo-Saint Blancard et al. [2022] Régaldo-Saint Blancard, B., Allys, E., Auclair, C., et al. 2022, arXiv e-prints, arXiv
  * Remy et al. [2022] Remy, B., Lanusse, F., Jeffrey, N., et al. 2022, arXiv preprint arXiv:2201.05561
  * Rombach et al. [2021] Rombach, R., Blattmann, A., Lorenz, D., Esser, P., & Ommer, B. 2021, High-Resolution Image Synthesis with Latent Diffusion Models.  https://arxiv.org/abs/2112.10752
  * Ronneberger et al. [2015] Ronneberger, O., Fischer, P., & Brox, T. 2015, in International Conference on Medical image computing and computer-assisted intervention, Springer, 234–241
  * Rouhiainen et al. [2023] Rouhiainen, A., Gira, M., Münchmeyer, M., Lee, K., & Shiu, G. 2023, arXiv preprint arXiv:2311.05217
  * Rozet et al. [2021] Rozet, F., Delaunoy, A., Miller, B., et al. 2021, LAMPE: Likelihood-free Amortized Posterior Estimation, doi: 10.5281/zenodo.8405782
  * Sharma et al. [2024] Sharma, D., Dai, B., Villaescusa-Navarro, F., & Seljak, U. 2024, arXiv preprint arXiv:2401.15891
  * Shen et al. [2021] Shen, Z., Zhang, M., Zhao, H., Yi, S., & Li, H. 2021, in Proceedings of the IEEE/CVF winter conference on applications of computer vision, 3531–3539
  * Song et al. [2021] Song, Y., Shen, L., Xing, L., & Ermon, S. 2021, arXiv preprint arXiv:2111.08005
  * Song et al. [2020] Song, Y., Sohl-Dickstein, J., Kingma, D. P., et al. 2020, arXiv preprint arXiv:2011.13456
  * Tange [2018] Tange, O. 2018, GNU Parallel 2018 (Ole Tange), doi: 10.5281/zenodo.1146014
  * Theis et al. [2015] Theis, L., Oord, A. v. d., & Bethge, M. 2015, arXiv preprint arXiv:1511.01844
  * Valogiannis & Dvorkin [2022] Valogiannis, G., & Dvorkin, C. 2022, Physical Review D, 106, 103509
  * Valogiannis et al. [2023] Valogiannis, G., Yuan, S., & Dvorkin, C. 2023, arXiv preprint arXiv:2310.16116
  * Vaswani et al. [2017] Vaswani, A., Shazeer, N., Parmar, N., et al. 2017, Advances in neural information processing systems, 30
  * Villaescusa-Navarro et al. [2021a] Villaescusa-Navarro, F., Anglés-Alcázar, D., Genel, S., et al. 2021a, The Astrophysical Journal, 915, 71
  * Villaescusa-Navarro et al. [2021b] —. 2021b, arXiv preprint arXiv:2109.09747
  * Villaescusa-Navarro et al. [2022] Villaescusa-Navarro, F., Genel, S., Angles-Alcazar, D., et al. 2022, The Astrophysical Journal Supplement Series, 259, 61
  * Villanueva-Domingo & Villaescusa-Navarro [2022] Villanueva-Domingo, P., & Villaescusa-Navarro, F. 2022, The Astrophysical Journal, 937, 115
  * Wildberger et al. [2024] Wildberger, J., Dax, M., Buchholz, S., et al. 2024, Advances in Neural Information Processing Systems, 36
  * Wu & He [2018] Wu, Y., & He, K. 2018, in Proceedings of the European conference on computer vision (ECCV), 3–19
  * Zagoruyko & Komodakis [2016] Zagoruyko, S., & Komodakis, N. 2016, arXiv preprint arXiv:1605.07146

## Appendix

Figure 8: The scale / standard deviation of the cumulative noise in Equation 1
added to the image over different timesteps.

###  .1 Parameter Inference

In Equation 5, we set $\bm{M}^{-1}$ to be diagonal with [1, 5]. Setting the
inverse mass matrix in an HMC to be close to the covariance of the expected
posterior distribution helps the chain explore the space better. In this case,
we choose a step size of $5\times 10^{-4}$, because of the steep gradient of
the posterior distribution with respect to $\Omega_{m}$. For the parameter on
field where the true parameter is 0.101, i.e. on the prior range, we need to
further reduce the step size to $1\times 10^{-4}$ in order for the parameters
to be accepted. We modified the Hamiltorch package in order to generate
samples. For $T_{\rm{MAX}}\geq 20$ we compute the contributions to the VLB
loss in batches of 10 timesteps and accumulate the gradient contribution for
each batch. This enables us to compute the gradients with 1000 timesteps. The
choice of mass matrix accelerates the chain’s convergence to the correct
region of parameter space for $\sigma_{8}$. We approximate -$\log
p_{\phi}(x_{0}|\theta)$ with the truncated variational lower bound in Equation
4. While we do not compute the expectation over multiple seeds within a single
evaluation of Equation 5, for speed, every evaluation uses a different seed
and noise pattern. The prior is chosen to be a flat prior over
$\Omega_{m}\in[0.1,0.5]$ and $\sigma_{8}\in[0.6,1.0]$.

$\displaystyle L_{vlb}=L_{0}+L_{1}...L_{T-1}+L_{T}=$ (7)
$\displaystyle=\mathbb{E}[-\log
p_{\phi}(x_{0}|\theta)]\leq\mathbb{E}_{q}[\rm{D_{KL}}[q(x_{T}|x_{0})||p(x_{T})]+\sum_{t\geq
1}D_{KL}[q(x_{t}|x_{t+1},x_{0})||p_{\phi}(x_{t}|x_{t+1},\theta)]-\log
p_{\phi}(x_{0}|x_{1},\theta)]$

Notation:

  * •

$x_{0}$: Normalized input field

  * •

$\phi$: Noise Model (Neural Network)

  * •

$\theta$: Conditioning cosmology i.e. a vector with $\Omega_{m}$ and
$\sigma_{8}$

Computing VLB terms:

$\displaystyle L_{0}=-\log
p_{\phi}(x_{0}|x_{1},\theta)=-ln\mathcal{N}(x_{0}|\mu_{0},\beta_{0})=\sum_{p}\frac{(x_{0}-\mu_{0})^{2}}{2\beta_{0}}+0.5ln|2\pi\beta_{0}|$
$\displaystyle\textrm{For
t}\in[1,T-1],L_{t}=D_{KL}[q(x_{t}|x_{t+1},x_{0})||p_{\phi}(x_{t}|x_{t+1},\theta)]$
$\displaystyle
q(x_{t}|x_{t+1},x_{0})=\mathcal{N}(\tilde{\mu_{t}}(x_{t+1}x_{0}),\tilde{\beta}_{t})$
$\displaystyle\tilde{\mu_{t}}(x_{t+1}x_{0})=\frac{\sqrt{\bar{\alpha}_{t-1}\beta_{t}}}{1-\bar{\alpha}_{t}}x_{0}+\frac{\sqrt{\alpha_{t}(1-\bar{\alpha}_{t-1})}}{1-\bar{\alpha}_{t}}x_{t+1}\text{
and
}\tilde{\beta}_{t}=\beta_{t}\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_{t}}$
$\displaystyle
p_{\phi}(x_{t}|x_{t+1},\theta)=\mathcal{N}(\tilde{\mu_{t}}(x_{t+1},\hat{x}_{0,\phi}),\tilde{\beta}_{t})$
$\displaystyle\hat{x}_{0,\phi}=\frac{x_{t+1}-\sqrt{1-\bar{\alpha}_{t}\epsilon_{\phi}(x_{t+1},t,\theta)}}{\sqrt{\bar{\alpha_{t}}}}$

####  .1.1 HMC Convergence

For a single field, we examine the convergence of parameters, when the chain
starts from different initial parameters in Figure 9. The chains are
indistinguishable beyond around 50 samples. We thus choose a burn-in of 100
samples. The $\hat{R}$ for both parameters computed using the samples in
$[100-300]$ is 0.997. A $\hat{R}$ of greater than 1.1 usually indicates that
the chains have not converged and still retain some memory of their
initialization. While the $\hat{R}$ is theoretically expected to be around 1
or slightly greater, some numerical variation about this expected value can
result in values that are slightly less than 1. The $\hat{R}$ is a measure of
the variance between chains divided by the variance within chains.

Figure 9: HMC chains for 300 samples, for the same field, starting at two
different initial parameters: [0.11, 0.62] (beige) and [0.3, 0.8] (teal). The
chains are well mixed by the cutoff we designate as our burn-in (100 samples),
denoted by the dashed line.

###  .2 Parameter Estimation Baselines

Power Spectrum NPE Baseline: We use a Masked Autoregressive Flow[29] to
implement the normalizing flow that predicts the 2 dimensional parameter
vector given the 129 dimensional feature vector for a single field (128 bins
for the power spectrum+1 for the mean of the log fields). The power spectrum
is the log of the overdensity power spectrum of the (unlogged) fields. In
Figure 6, the power spectrum sample contours are smoothed by convolving with a
Gaussian kernel with a scale of 0.8 and the Diffusion-HMC samples are smoothed
by a kernel of 0.2. The ellipses in the inset figure are computed using the
covariance of the 400 diffusion model samples and finding the ellipses
corresponding to the 68.3, 99.4 and 99.7$\%$ confidence intervals, using the
eigendecomposition of the covariance.

###  .3 Summary Statistics

####  .3.1 Reduced Chi-squared Statistics

For the CV fields, where we have 450 samples of the true and generated fields
for a single parameter, we compute the reduced chi-squared statistic using an
estimate of the covariance between different $k$ bins. The number of $k$ bins
here is 35.

$\displaystyle\vec{\mu}=\langle Pk_{Ref}\rangle\quad\quad\quad
C=\rm{Cov}[Pk_{Ref}]\quad\quad\quad\hat{C}^{-1}=C^{-1}\frac{N-p-2}{N-1}$ (8)
$\displaystyle\chi^{2}_{r}(Pk_{Test})=\sum_{k}(Pk_{Test}-\vec{\mu})\cdot(\hat{C}^{-1}(Pk_{Test}-\vec{\mu})^{T})^{T}$
(9)

For the LH fields during model selection, since we just have 15 fields in the
true dataset, we cannot reliably estimate a covariance. We use the following
formula instead. The number of $k$ bins here is 128.

$\displaystyle\chi_{r}^{2}(s)=\frac{1}{|k|}\sum_{k}\frac{(P(k)_{s}-<P(k)_{\textsc{True}}>)^{2}}{\sigma[P(k)_{\textsc{True}}]^{2}}$
(10)

####  .3.2 Across Checkpoints

For the eight checkpoints we generated 500 fields, with 50 fields for each of
the 10 validation parameters. We then examined how different the reduced chi-
squared statistic of the power spectra of the log fields, the linear fields
and the $p$ values of the means of the distributions of true and generated
fields. A value of less than 0.05 would indicate that the two distributions of
the means are statistically different. The $p$ values are above 0.05 for all
eight checkpoints. Since these comparisons are limited by the number of
samples in the true set for each parameter (15), we additionally plot the
reduced chi-squared statistic derived by using each true fields as the test
and the other 14 fields as the reference. While there is some oscillation
across checkpoints for each of these three statistics, the variation appears
random.

Figure 10: Mean and standard error on the mean for the reduced $\chi_{r}^{2}$
statistic of the power spectra of the 50 linear (left) and log (center) fields
for each parameter relative to the 15 true fields’ power spectra for that
parameter, across 10 different parameters for each of 8 checkpoints. The
dashed line demarcates the mean reduced $\chi_{r}^{2}$ of the true fields
using the other 14 true fields as the reference distribution (leave-one-out).
Right: Mean and standard error on the mean of the $p$ values of the
distribution of the means of the 50 generated log fields relative to the
distribution of the means of the 15 true fields for the same parameter. The
$p$ values are above 0.05 for all of the 8 checkpoints.