# Scaling of pre-equilibrium dilepton production in QCD Kinetic Theory

Oscar Garcia-Montero<EMAIL_ADDRESS>Fakultät für Physik,
Universität Bielefeld, D-33615 Bielefeld, Germany  Philip Plaschke
<EMAIL_ADDRESS>Fakultät für Physik, Universität Bielefeld,
D-33615 Bielefeld, Germany  Sören Schlichting<EMAIL_ADDRESS>bielefeld.de Fakultät für Physik, Universität Bielefeld, D-33615 Bielefeld,
Germany

###### Abstract

We use QCD kinetic theory to compute dilepton production coming from the pre-
equilibrium phase of the Quark-Gluon Plasma created in high-energy heavy-ion
collisions. We demonstrate that the dilepton spectrum exhibits a simple
scaling in terms of the specific shear viscosity $\eta/s$ and entropy density
$dS/d\zeta\sim{\scriptstyle\left(T\tau^{1/3}\right)_{\infty}^{3/2}}$, which
can be derived from dimensional analysis in the presence of a pre-equilibrium
attractor. Based on this scaling we perform event-by-event calculations of in-
medium dilepton production and determine the invariant mass range where the
pre-equilibrium yield is the leading contribution.

## Introduction.

Developing a first principles understanding of the non-equilibrium QCD
processes that lead to the formation of a close to equilibrium Quark-Gluon
Plasma (QGP), is one of the most outstanding challenges in modern heavy-ion
physics. By now it has been established [1, 2] that – in various different
microscopic theories – the onset of hydrodynamic behavior typically occurs on
the time scale of a single equilibrium relaxation time
$\sim(4\pi\eta/s)/T_{\rm eq}$, as determined by the specific shear viscosity
$\eta/s$ and (local) temperature of the system. Over the course of this
evolution, the non-equilibrium QGP also undergoes a rapid information loss,
whereby (some) details of the initial conditions become irrelevant, such that
the pre-equilibrium evolution is governed by pre-equilibrium attractors (see
Ref. [3, 4, 5, 6, 7, 8] and references therein) even before a viscous fluid
dynamic description becomes applicable.

Despite significant theoretical progress, experimental measurements to probe
the early-time pre-equilibrium dynamics of the QGP at present and future
facilities at RHIC and the LHC remain elusive. Electromagnetic probes, namely
photons $(\gamma)$ and dileptons ($l^{+}l^{-}$), provide promising candidates,
as they are produced throughout the entire space-time evolution of the
collisions, and due to the lack of re-interactions with the QGP medium, escape
the reaction unscathed. However, experimental measurement that isolate this
signal are extremely challenging, as in addition to the “direct” production,
which can be further sub-divided into prompt contributions (Drell-Yan [9]),
thermal photons/dileptons and early-time pre-equilibrium radiation of the QGP,
photons and dileptons are also copiously produced from late stage hadronic
decays (see e.g. [10, 11] and references therein).

Evidently, an important property of dileptons is the invariant mass $M$ of the
lepton pair, as it can be used to discriminate between different production
mechanism [12, 13]. Generally speaking, lepton pairs with larger invariant
masses tend to be produced at earlier times, such that e.g. the Drell-Yan
process is the main source of high-$M$ dileptons. The intermediate mass region
is produced by in-medium dileptons coming from the pre-equilibrium and thermal
phase of the QGP. Recent estimates of the pre-equilibrium dilepton production
have shown that in an invariant mass range of $3\;\text{GeV}<M<4\;\text{GeV}$
the pre-equilibrium contribution can indeed exceed the thermal QGP and Drell-
Yan contributions [13, 14], and first ideas to develop a phenomenology of the
pre-equilibrium phase based on dilepton measurements have been developed [12,
15].

In this letter we show that the pre-equilibrium dilepton production exhibits a
simple scaling behavior (c.f. Eq. 5) in terms of the specific shear viscosity
$\eta/s$ and local entropy density $dS/d^{2}\mathbf{x}_{T}d\zeta$ of the QGP,
which can also be extended to other probes of the early pre-equilibrium phase
[16]. By performing microscopic calculations in QCD kinetic theory [17, 18,
19, 20], which simultaneously provides a consistent theoretical description of
the thermalization of the QGP and the production of electromagnetic probes, we
explicitly verify this scaling and compute the associated scaling function for
the invariant mass spectrum $dN_{l^{+}l^{-}}/d^{2}\mathbf{x}_{T}MdMdy_{Q}$ of
pre-equilibrium dileptons. Since the scaling allows for a particularly
efficient implementation of pre-equilibrium dilepton production, we further
perform state-of-the art phenomenological calculations of (direct) dilepton
production in heavy-ion collisions, which demonstrate the importance of event-
by-event fluctuations in determining the dilepton yield in the intermediate
mass range $3\;\text{GeV}<M<4\;\text{GeV}$, where the pre-equilibrium
production dominates.

## Pre-equilibrium dilepton production

Dilepton production in leading order QCD kinetic theory is governed by quark
anti-quark annihilation, with the corresponding production rate determined by

$\displaystyle\frac{dN_{l^{+}l^{-}}}{d^{4}Xd^{4}Q}=\frac{8\pi}{3}N_{c}\sum_{f}q_{f}^{2}\alpha_{EM}^{2}\int\frac{d^{3}p_{1}}{(2\pi)^{3}p_{1}}\frac{d^{3}p_{2}}{(2\pi)^{3}p_{2}}$
$\displaystyle\qquad\times
f_{q_{f}}(x,P_{1})f_{\overline{q}_{f}}(x,P_{2})\delta^{(4)}(Q-P_{1}-P_{2})\,,$
(1)

where $Q=P_{l^{+}}+P_{l^{-}}$ is the four momentum of the dilepton pair. Over
the short course of the pre-equilibrium phase, the space-time evolution of the
non-equilibrium QGP can be locally described in terms of a Bjorken flow, such
that the space-time differential is given by $d^{4}X=\tau d\tau d\zeta
d^{2}\mathbf{x}_{T}$ where $\tau=\sqrt{t^{2}-z^{2}}$ is the proper time,
$\zeta=\text{arctanh}(z/t)$ is the space-time rapidity and $\mathbf{x}_{T}$
denotes the transverse coordinates. While in thermal equilibrium, the
production rate can be integrated to yield the well known leading order result
(see also [21, 22] for NLO corrections)

$\displaystyle\frac{dN_{l^{+}l^{-}}}{d^{4}Xd^{4}Q}=\frac{N_{C}\alpha_{\text{EM}}^{2}}{12\pi^{4}}\sum_{f}q_{f}^{2}\frac{F(q)}{\exp(q^{0}/T)-1}\,,$
(2)

with

$F(q)=\frac{2T}{q}\ln\bqty{\cosh(\frac{q^{0}+q}{4T})\left/\cosh(\frac{q^{0}-q}{4T})\right.}$
(3)

the non-equilibrium production of dilepton pairs is governed by the phase-
space distributions of quarks and anti-quarks $f_{q_{f}}(x,P_{1})$ and
$f_{\overline{q}_{f}}(x,P_{2})$ in Eq. (Pre-equilibrium dilepton production).

During the pre-equilibrium phase the QGP is believed to be highly anisotropic
in momentum space, as it is subject to a rapid longitudinal expansion, and
dominated by gluon degrees a freedom, as these are more copiously produced in
the initial (semi-) hard scatterings [2, 23]. Notably, these effects result in
a momentum-space anisotropy [24, 25], a non-trivial polarization [12] and an
overall suppression of the dilepton yield, as discussed in previous works [14,
13, 12], where a phenomenological parametrization of the phase-space
distributions $f_{q/\bar{q}}$ was employed to compute the pre-equilibrium
dilepton production. Conversely, in this work we will compute the pre-
equilibrium dilepton production directly from a microscopic calculation in QCD
kinetic theory. However, before we present results of our microscopic
calculations, it proves particularly insightful to perform a dimensional
scaling analysis of the pre-equilibrium dilepton yield.

## Scaling of pre-equilibrium dilepton production

Generally, the pre-equilibrium production of electromagnetic probes will
proceed differently at different positions $\mathbf{x}_{T}$ in the transverse
plane and depend on the timescale $\tau_{\rm eq}(\mathbf{x}_{T})$, and the
temperature $T_{\rm eq}(\mathbf{x}_{T})$ at which the system locally
approaches a viscous hydrodynamic description. By dimensional analysis, it is
natural to assume, that the local production rate
$\frac{dN_{l^{+}l^{-}}}{d^{4}Xd^{4}Q}$ at each transverse position
$\mathbf{x}_{T}$ only depends on the dimensionless ratios $\tau/\tau_{\rm
eq}(\mathbf{x}_{T})$ and $Q/T_{\rm eq}(\mathbf{x}_{T})$,

$\displaystyle\frac{dN_{l^{+}l^{-}}}{d^{4}Xd^{4}Q}=\frac{dN_{l^{+}l^{-}}}{d^{4}Xd^{4}Q}\left(\frac{\tau}{\tau_{\rm
eq}(\mathbf{x}_{T})},\frac{Q}{T_{\rm eq}(\mathbf{x}_{T})}\right)$ (4)

such that upon integrating over the production time $\tau$, the space-time
time rapidity $\zeta$, and the transverse momentum $q_{T}$ of the dilepton
pair, the local yield
$\frac{dN_{l^{+}l^{-}}}{d^{2}\mathbf{x}_{T}MdMdy_{Q}}=\int d\tau\tau\int
d\zeta\int d^{2}q_{T}\frac{dN_{l^{+}l^{-}}}{d^{4}Xd^{4}Q}$ per unit rapidity
$y_{Q}$ and invariant mass $M=\sqrt{Q^{2}}$ is determined by

$\displaystyle\frac{dN_{l^{+}l^{-}}}{d^{2}\mathbf{x}_{T}MdMdy_{Q}}=\tau_{\rm
eq}^{2}(\mathbf{x}_{T})T_{\rm eq}^{2}(\mathbf{x}_{T})\leavevmode\nobreak\
\mathcal{N}_{{l^{+}l^{-}}}\left(\frac{M}{T_{\rm eq}}\right)$ (5)

where $N_{l^{+}l^{-}}$ is a scaling function, which only depends on the ratio
$M/T_{\rm eq}$, and is determined by

$\displaystyle\mathcal{N}_{{l^{+}l^{-}}}\left(\frac{M}{T_{\rm eq}}\right)$
$\displaystyle=\int d\frac{\tau}{\tau_{\rm eq}}\frac{\tau}{\tau_{\rm eq}}\int
d\zeta\int d^{2}\frac{q_{T}}{T_{\rm eq}}$ (6)
$\displaystyle\quad\frac{dN_{l^{+}l^{-}}}{d^{4}Xd^{4}Q}\left(\frac{\tau}{\tau_{\rm
eq}},\frac{Q}{T_{\rm eq}}\right)$ (7)

While the calculation of the scaling function
$\mathcal{N}_{l^{+}l^{-}}\left(\frac{M}{T_{\rm eq}}\right)$ is highly non-
trivial, it is comparatively straightforward to isolate the dependence on the
local entropy density and coupling strength of the system. We first note that
different studies at weak [26, 27] and strong coupling [4, 5, 8]) indicate,
that equilibration of the out-of-equilibrium QGP is controlled by the
equilibrium relaxation rate $\Gamma_{\rm eq}=T_{\rm eq}/(4\pi\eta/s)$, such
that hydrodynamics becomes applicable when the dimensionless scaling variable
$\tilde{w}=\Gamma_{\rm eq}\tau_{\rm eq}$ is of order unity, such that

$\displaystyle\tau_{\rm eq}T_{\rm eq}\approx 4\pi\eta/s\,.$ (8)

Since the equilibrium temperature $T_{\rm eq}$ can be related to the local
entropy density of the system $\frac{dS}{d^{2}\mathbf{x}_{T}d\zeta}$, which in
turn is directly related to the multiplicity of charged particles produced in
the collisions, one further obtains the relation

$\displaystyle\frac{2\pi^{2}}{45}\nu_{\rm eff}\tau_{\rm eq}T_{\rm
eq}^{3}\approx\frac{dS}{d^{2}\mathbf{x}_{T}d\zeta}$ (9)

where $\nu_{\rm eff}\approx 32$ denotes the effective number of bosonic
degrees of freedom [28, 29]. Denoting the dependence on the local entropy
density in terms of the quantity
$(T\tau^{\frac{1}{3}})^{3/2}_{\infty}=\sqrt{\frac{dS}{d^{2}\mathbf{x}_{T}d\zeta}/\left(\frac{2\pi^{2}}{45}\nu_{\rm
eff}\right)}$, by virtue of Eqns. (8) and (9) the equilibration time and
equilibration temperature can then be determined as

$\displaystyle\tau_{\rm eq}$
$\displaystyle\approx(T\tau^{\frac{1}{3}})^{-3/2}_{\infty}\left(4\pi\eta/s\right)^{3/2}$
(10) $\displaystyle T_{\rm eq}$
$\displaystyle\approx(T\tau^{\frac{1}{3}})^{3/2}_{\infty}\leavevmode\nobreak\
\left(4\pi\eta/s\right)^{-1/2}$ (11)

and the pre-equilibrium dilepton yield can be expressed as

$\displaystyle\frac{dN_{l^{+}l^{-}}}{d^{2}\mathbf{x}_{T}MdMdy_{Q}}=(4\pi\eta/s)^{2}\leavevmode\nobreak\
\mathcal{N}_{l^{+}l^{-}}\left(\frac{\sqrt{4\pi\eta/s}\leavevmode\nobreak\
M}{(T\tau^{1/3})^{3/2}_{\infty}}\right)$ (12)

Strikingly, the above scaling can be derived on rather general grounds, and
thus easily extended to a variety of other quantities, such as e.g. the
transverse momentum ($q_{T}$) spectrum, which takes the form

$\displaystyle\frac{dN_{l^{+}l^{-}}}{d^{2}\mathbf{x}_{T}d^{2}\mathbf{q}_{T}MdMdy_{Q}}=(4\pi\eta/s)^{3}\leavevmode\nobreak\
(T\tau^{\frac{1}{3}})^{-3}_{\infty}$ (13)
$\displaystyle\quad\mathcal{N}_{l^{+}l^{-}}\left(\frac{\sqrt{4\pi\eta/s}\leavevmode\nobreak\
M}{(T\tau^{\frac{1}{3}})^{3/2}_{\infty}},\frac{\sqrt{4\pi\eta/s}\leavevmode\nobreak\
q_{T}}{(T\tau^{\frac{1}{3}})^{3/2}_{\infty}}\right)$ (14)

Most importantly, such scaling can also be expected to be present for other
processes, such as e.g. the pre-equilibrium production of photons [16], heavy
flavor quarks [30] or pre-equilibrium jet quenching [31], which will greatly
facilitate their phenomenological treatment in the future.

## Dilepton production from non-equilibrium QGP evolution.

(a) Dilepton spectra

(b) Universal scaling curve

Figure 1:  Spectrum of dileptons produced during the pre-equilibrium stage as
a function of (a) $M/(T\tau^{1/3})_{\infty}^{3/2}$ and (b)
$\sqrt{\eta/s}M/(T\tau^{1/3})_{\infty}^{3/2}$. Different colors correspond to
different coupling strengths $\lambda=5,10,20$ and different panels correspond
to different final integration times $\tilde{w}=0.5,1.0$. Spectra in the right
panels are divided by $(\eta/s)^{2}$ in order to obtain the universal scaling
curve.

In order to verify the predicted scaling we compute the pre-equilibrium
spectrum of dileptons for different values of the coupling $\lambda=4\pi
N_{c}\alpha_{s}$ using a QCD kinetic theory description, where we incorporate
all leading order processes of quarks and gluons into the Boltzmann equation
to compute the respective distribution functions during the pre-equilibrium
evolution of the QGP [18, 16, 27]. Within this framework infrared divergencies
are regulated by an isotropic screening [32, 33, 34, 26, 27] and in-medium
effects associated with the Landau-Pomeranchuk-Migdal (LPM) effect are treated
via an effective vertex resummation. Starting from a gluon dominated initial
state as in [16], we numerically solve the Boltzmann equation to determine the
evolution of the phase-space distributions $f_{q/\overline{q}}$, which then
enables us to compute the dilepton rate at each time step.

$\lambda$ | 5 | 10 | 20
---|---|---|---
$\alpha_{s}$ | 0.133 | 0.265 | 0.531
$\eta/s$ | 2.716 | 0.994 | 0.385
$(\tau^{1/3}T)_{\infty}/Q_{s}^{2/3}$ | 0.799 | 0.612 | 0.474
Table 1: Strong coupling $\alpha_{s}$, shear viscosity to entropy density
ratio $\eta/s$ and asymptotic temperature
$(\tau^{1/3}T)_{\rm\infty}/Q_{s}^{2}$ for different coupling strengths
$\lambda$.

Numerical results for the pre-equilibrium spectrum of dileptons as a function
of invariant mass are presented in Fig. 1(a), where we show the integrated
dilepton yield up to $\tilde{w}=0.5$ (top), during the thermalization process
of the QGP and $1.0$ (bottom), when visc. hydrodynamics becomes applicable.
When comparing the results for different coupling strength $\lambda=5,10,20$,
which correspond to different values of the specific shear viscosity $\eta/s$,
we express the spectra in terms of the invariant mass
$M/{(\tau^{1/3}T)^{3/2}_{\infty}}$ (in Fig. 1(a)) and the re-scaled invariant
mass variable ${(\eta/s)^{1/2}}M/{(\tau^{1/3}T)^{3/2}_{\infty}}$ (in Fig.
1(b)). The specific shear viscosity as well as the asymptotic constant
$(\tau^{1/3}T)_{\infty}$ are obtained by comparison to the first order viscous
hydrodynamic expansions of the longitudinal pressure resp. the effective
temperature (see [33, 35, 16] for more details). The results for
$\lambda=5,\,10,\,20$ are compactly summarized in Table 1.

When comparing the results for different coupling strength as a function of
$M/{(\tau^{1/3}T)^{3/2}_{\infty}}$, we see that stronger couplings tend to
produce more large-$M$ dileptons. Setting
$(\tau^{1/3}T)^{3/2}_{\infty}=0.402\;\text{GeV}$ as in [16], to mimic the
conditions in central Pb+Pb collisions at LHC energies, this occurs for
invariant masses $M\gtrsim 2\;\text{GeV}$. Below $M\lesssim 2\;\text{GeV}$
this changes and more weakly coupled systems produce more dileptons during the
pre-equilibrium phase, i.e. for $\tilde{w}<1$. We further see that high-$M$
dileptons are predominantly produced at the earliest times of the evolution.
Indeed, the curves for $\tilde{w}=0.5$ are almost not modified during the
evolution up to $\tilde{w}=1.0$. However, in the mean time the QGP produces
many low-mass dileptons, such that the spectrum increases significantly in the
low-$M$ regime.

Besides the effective kinetic theory (EKT) results, we also show results for
an effective parameterization (EP) of the distribution functions already used
in previous works for pre-equilibrium dilepton production [14, 13, 12], where
the authors used anisotropic equilibrium distribution functions of the
Romatschke-Strickland form, together with an overall quark suppression factor.
Compared to the EKT results, the spectra computed from the effective
parameterization are suppressed at both times, although the suppression is
stronger for earlier times. This is expected as the description of pre-
equilibrium dynamics in terms of anisotropic equilibrium distributions is
worse for earlier time dynamics. Over the course of the evolution, the spectra
for the effective parametrization get closer to the EKT results and for
smaller couplings we find almost the same spectra for the low and intermediate
mass range.

Next, in order to verify the universal scaling, we re-plot the same spectra in
Fig. 1(b) but this time with the predicted scaling from Eq. 12. By
construction, the spectra obtained using the effective parametrization
perfectly collapse to each other, while the EKT results also get much closer
to each other. Since the scaling function $N_{l^{+}l^{-}}$ is derived on basis
of an attractor solution, the scaling for the EKT results works better and
better for later evolution. Especially for $\tilde{w}\gtrsim 1$, which is in
accordance with the hydrodynamization time found in earlier studies in the
context of QCD kinetic theory [36, 7, 33, 37], the scaling works very well,
whereas at very early times the spectra are sensitive to the initial
conditions for the kinetic evolution.

## Phenomenology of the pre-equilibrium dilepton.

One remarkable consequence of the scaling formula is the practical usage for
phenomenological calculations and comparisons to experimental data. In order
to demonstrate this, we match the EKT evolution to the initial conditions of
the subsequent hydrodynamic evolution of the QGP. By use of the scaling
formula in Eq. 12, the only other ingredients to compute the spectrum of
dileptons produced during the pre-equilibrium are the (local) temperature
$T(\tau_{\rm hydro},\mathbf{x}_{T})\equiv T_{\rm hydro}(\mathbf{x}_{T})$ at
initialization time $\tau_{\rm hydro}$ and the specific shear-viscosity
$\eta/s$. Based on this information, the local values of the scaling variable
$\tilde{w}(\mathbf{x}_{T})$ can then be determined as

$\tilde{w}(\mathbf{x}_{T})=\frac{\tau_{\rm
hydro}T_{\text{hydro}}(\mathbf{x}_{T})}{4\pi\eta/s}\,,$ (15)

and similarly, the energy scale $(\tau^{1/3}T)_{\infty}(\mathbf{x}_{T})$ can
be obtained according to [33, 35]

$(\tau^{1/3}T)_{\infty}(\mathbf{x}_{T})=\tau_{\rm hydro}^{1/3}\,\left(T_{\rm
hydro}(\mathbf{x}_{T})+\frac{2}{3}\frac{\eta/s}{\tau_{\rm hydro}}\right)\,$
(16)

accounting for first order viscous corrections. Once those two quantities are
determined at each point in the transverse plane, the total pre-equilibrium
dilepton yield is then given by

$\displaystyle\frac{dN_{l^{+}l^{-}}}{MdMdy_{Q}}=(\eta/s)^{2}\\!\\!\int\\!d^{2}\mathbf{x}_{T}\leavevmode\nobreak\
\mathcal{N}_{l^{+}l^{-}}\pqty{\tilde{w}(\mathbf{x}_{T}),\textstyle\frac{\sqrt{\eta/s}\leavevmode\nobreak\
M}{\left(T\tau^{1/3}\right)_{\infty}^{3/2}(\mathbf{x}_{T})}}\,.$ (17)

We compute the event-by-event dilepton spectra at midrapidity for Pb-Pb
collisions at $\sqrt{s_{NN}}=5.02$ TeV. To account for a background evolution
well fit to data, we created the events using the hybrid model _Trajectum_
[38], and used one of the consistent parameter sets extracted using the
Bayesian analysis in Ref. [39]. We focus on the $0-5\%$ class, and obtain the
pre-equilibrium dilepton yield by integrating event-by-event over the initial
profiles for the hydrodynamical evolution as in Eq. 17. By matching the
initial profiles to the subsequent hydrodynamic evolution at
$\tau_{hydro}=1.0$ fm as in [16], the thermal QGP contribution is then found
by folding the thermal dilepton production rate with the background
hydrodynamic evolution from $\tau_{hydro}=1.0$ fm onwards111The QGP rate is
folded with the hydro evolution for temperatures higher than $T_{c}=154$ MeV.
Below this temperature, hadronic degrees of freedom dominate, and the rate can
be computed from effective theories [38]. Nevertheless, the hadronic rates are
too soft to contribute significantly to the invariant mass range relevant to
this work.. Clearly, one of the main advantages of using this description is
that it guarantees for a smooth matching between pre-equilibrium and
hydrodynamical dilepton production, such that the overall spectrum of in-
medium dileptons is virtually independent of the matching time $\tau_{\rm
hydro}$.

We present our main phenomenological result in Fig. 2, where we compare the
in-medium dilepton production, to the Drell-Yan (DY) contribution, computed to
NLO precision with the EPPS nPDFs [40] by using the DYTurbo software [41],
where as in previous works [13] the centrality dependence of the nPDFs was
neglected and $T_{AA}$ scaling was assumed. While the top panel of Fig. 2
shows the total dilepton yield $dN_{l^{+}l^{-}}/dMdy_{Q}$, the lower panel
shows a break-up of the total yield into DY and thermal QGP contributions, as
well as in the the pre-equilibrium contribution from ranges of evolution time
$\tilde{w}$, and clearly illustrates how larger invariant masses correspond to
earlier production times.

While at higher invariant mass ranges the DY contribution is well constrained
and therefore presents small uncertainties, at low $M$s, it presents large
error bands, due to large scale uncertainties. Despite these large theoretical
uncertainties, we find that the pre-equilibrium dilepton contribution obtained
from realistic event-by-event EKT+Hydro simulations dominates over both the DY
and thermal QGP background for invariant masses $\sim 3\;\text{GeV}$.

Beyond the results obtained from event-by-event (EByE) simulations, we have
also computed the in-medium dilepton spectra for a coarse grained (CG)
description, where similar to [42, 43] we compute the emission for a single
average event, as well as comparing to previous obtained results for a simple
Bjorken scaling scenario [13], where transverse dynamics of the QGP is
completely neglected. While the Bjorken scaling underestimates the cooling of
the QGP, and thus overestimates the dilepton yield by almost a factor of
three, the coarse grained averaged description provides a very good
description for low invariant masses $\lesssim 2\;\text{GeV}$. Nevertheless,
the coarse grained description also but underestimates the high $M$ tails of
in-medium dilepton production, which are more likely produced from local hot
spots, and we therefore conclude from this analysis, that an event-by-event
description is desirable to provide accurate predictions for high and
intermediate mass dileptons. This difference between the EbE and CG
descriptions becomes more important for less central collisions.

Figure 2:  (Upper panel) Comparison of the in-medium dilepton spectrum for
$0-5\%$ centrality Pb-Pb collisions at $\sqrt{s_{NN}}=5.02$ TeV to their
corresponding Drell-Yan contribution. The in-medium contributions are shown
for three different cases (transversely homogeneous Bjorken expansion (data
taken from [13]), realistic EbE ICs, and smooth initial conditions, the
average of the latter). (Lower panel) Relative contribution of the different
contributions to the total yield (in-medium + Drell-Yan). The EbE pre-
equilibrium has been split into different $\tilde{w}$ ranges for better
comparison.

## Summary and Outlook.

We computed the pre-equilibrium spectrum of dileptons in QCD kinetic theory
and showed that it exhibits a simple universal scaling in terms of the
specific shear viscosity $\eta/s$ and entropy density
$(T\tau^{{1}/{3}})_{\infty}$. Such scaling functions are powerful tools as
they allow for event-by-event calculations of the corresponding spectra by
matching $\eta/s$ and $(T\tau^{{1}/{3}})_{\infty}$ to realistic values of a
heavy-ion collision. By performing realistic event-by-event simulations of
pre-equilibrium, thermal QGP and Drell-Yan dilepton production for
$\sqrt{s_{NN}}=5.02$ TeV Pb-Pb collisions, we find that the pre-equilibrium
production dominates over the irreducible background for invariant masses
$\sim 3\;\text{GeV}$. However, this result only serve as an illustrative
example and in order to make full use of our results for phenomenological
applications, we have included the pre-equilibrium scaling functions into the
initial state framework KøMPøST. This version including photons and dileptons
is called ShinyKøMPøST and is publicly available under [44].

Since the pre-equilibrium scaling in Eq. (12) can be derived on rather general
grounds, it can also be extended to fully differential dilepton observables,
which is highly beneficial for the development of a Monte-Carlo generator to
facilitate the detection of pre-equilibrium dileptons in present and future
heavy-ion experiments. Since earlier studies showed that the polarization of
dileptons is a key observable to distinguish the pre-equilibrium emission from
other sources [12], clearly a next step will be to derive analogous scaling
functions for polarization observables. However, the use of pre-equilibrium
scaling functions is by no means limited to electromagnetic probes, an
extension to relevant topics such as heavy-flavor production [45, 46, 47, 30],
heavy-quark diffusion [48] or jet energy loss [49] would be interesting to see
and deserves further theoretical attention.

## Acknowledgments.

We thank Maurice Coquet, Travis Dore, Xiaojian Du, Stephan Ochsenfeld, Jean-
Yves Ollitrault, Michael Winn for valuable discussions. The authors want to
thank Aleksas Mazeliauskas for collaboration on related work. OGM, PP and SS
acknowledge support by the Deutsche Forschungsgemeinschaft (DFG, German
Research Foundation) through the CRC-TR 211 ‘Strong-interaction matter under
extreme conditions’-project number 315477589 – TRR 211. OGM and SS acknowledge
also support by the German Bundesministerium für Bildung und Forschung (BMBF)
through Grant No. 05P21PBCAA. The authors acknowledge computing time provided
by the Paderborn Center for Parallel Computing (PC2) and the National Energy
Research Scientific Computing Center, a DOE Office of Science User Facility
supported by the Office of Science of the U.S. Department of Energy under
Contract No. DE-AC02-05CH11231.

## References

  * Schlichting and Teaney [2019] S. Schlichting and D. Teaney, The First fm/c of Heavy-Ion Collisions, Ann. Rev. Nucl. Part. Sci. 69, 447 (2019), arXiv:1908.02113 [nucl-th] .
  * Berges  _et al._ [2021] J. Berges, M. P. Heller, A. Mazeliauskas, and R. Venugopalan, QCD thermalization: Ab initio approaches and interdisciplinary connections, Rev. Mod. Phys. 93, 035003 (2021), arXiv:2005.12299 [hep-th] .
  * Heller and Spalinski [2015] M. P. Heller and M. Spalinski, Hydrodynamics Beyond the Gradient Expansion: Resurgence and Resummation, Phys. Rev. Lett. 115, 072501 (2015), arXiv:1503.07514 [hep-th] .
  * Florkowski  _et al._ [2018] W. Florkowski, M. P. Heller, and M. Spalinski, New theories of relativistic hydrodynamics in the LHC era, Rept. Prog. Phys. 81, 046001 (2018), arXiv:1707.02282 [hep-ph] .
  * Romatschke and Romatschke [2019] P. Romatschke and U. Romatschke,  _Relativistic Fluid Dynamics In and Out of Equilibrium_, Cambridge Monographs on Mathematical Physics (Cambridge University Press, 2019) arXiv:1712.05815 [nucl-th] .
  * Kurkela  _et al._ [2020] A. Kurkela, W. van der Schee, U. A. Wiedemann, and B. Wu, Early- and Late-Time Behavior of Attractors in Heavy-Ion Collisions, Phys. Rev. Lett. 124, 102301 (2020), arXiv:1907.08101 [hep-ph] .
  * Du and Schlichting [2021a] X. Du and S. Schlichting, Equilibration of the Quark-Gluon Plasma at Finite Net-Baryon Density in QCD Kinetic Theory, Phys. Rev. Lett. 127, 122301 (2021a), arXiv:2012.09068 [hep-ph] .
  * Soloviev [2022] A. Soloviev, Hydrodynamic attractors in heavy ion collisions: a review, Eur. Phys. J. C 82, 319 (2022), arXiv:2109.15081 [hep-th] .
  * Drell and Yan [1970] S. D. Drell and T.-M. Yan, Massive Lepton Pair Production in Hadron-Hadron Collisions at High-Energies, Phys. Rev. Lett. 25, 316 (1970), [Erratum: Phys.Rev.Lett. 25, 902 (1970)].
  * Rapp [2013] R. Rapp, Dilepton Spectroscopy of QCD Matter at Collider Energies, Adv. High Energy Phys. 2013, 148253 (2013), arXiv:1304.2309 [hep-ph] .
  * Vujanovic  _et al._ [2016] G. Vujanovic, J.-F. Paquet, G. S. Denicol, M. Luzum, S. Jeon, and C. Gale, Electromagnetic radiation as a probe of the initial state and of viscous dynamics in relativistic nuclear collisions, Phys. Rev. C 94, 014904 (2016), arXiv:1602.01455 [nucl-th] .
  * Coquet  _et al._ [2023a] M. Coquet, X. Du, J.-Y. Ollitrault, S. Schlichting, and M. Winn, Dilepton polarization as a signature of plasma anisotropy, (2023a), arXiv:2309.00555 [nucl-th] .
  * Coquet  _et al._ [2021] M. Coquet, X. Du, J.-Y. Ollitrault, S. Schlichting, and M. Winn, Intermediate mass dileptons as pre-equilibrium probes in heavy ion collisions, Phys. Lett. B 821, 136626 (2021), arXiv:2104.07622 [nucl-th] .
  * Coquet  _et al._ [2023b] M. Coquet, X. Du, J.-Y. Ollitrault, S. Schlichting, and M. Winn, Transverse mass scaling of dilepton radiation off a quark-gluon plasma, Nucl. Phys. A 1030, 122579 (2023b), arXiv:2112.13876 [nucl-th] .
  * Seck  _et al._ [2023] F. Seck, B. Friman, T. Galatyuk, H. van Hees, E. Speranza, R. Rapp, and J. Wambach, Polarization of Thermal Dilepton Radiation, (2023), arXiv:2309.03189 [nucl-th] .
  * Garcia-Montero  _et al._ [2023] O. Garcia-Montero, A. Mazeliauskas, P. Plaschke, and S. Schlichting, Pre-equilibrium photons from the early stages of heavy-ion collisions, (2023), arXiv:2308.09747 [hep-ph] .
  * Arnold  _et al._ [2002] P. B. Arnold, G. D. Moore, and L. G. Yaffe, Photon and gluon emission in relativistic plasmas, JHEP 06, 030, arXiv:hep-ph/0204343 .
  * Arnold  _et al._ [2003] P. B. Arnold, G. D. Moore, and L. G. Yaffe, Effective kinetic theory for high temperature gauge theories, JHEP 01, 030, arXiv:hep-ph/0209353 .
  * Arnold  _et al._ [2001a] P. B. Arnold, G. D. Moore, and L. G. Yaffe, Photon emission from quark gluon plasma: Complete leading order results, JHEP 12, 009, arXiv:hep-ph/0111107 .
  * Arnold  _et al._ [2001b] P. B. Arnold, G. D. Moore, and L. G. Yaffe, Photon emission from ultrarelativistic plasmas, JHEP 11, 057, arXiv:hep-ph/0109064 .
  * Baier  _et al._ [1988] R. Baier, B. Pire, and D. Schiff, Dilepton production at finite temperature: Perturbative treatment at order $\alpha_{s}$, Phys. Rev. D 38, 2814 (1988).
  * Churchill  _et al._ [2023] J. Churchill, L. Du, C. Gale, G. Jackson, and S. Jeon, Dilepton production at NLO and intermediate invariant-mass observables, (2023), arXiv:2311.06675 [nucl-th] .
  * Iancu and Venugopalan [2003] E. Iancu and R. Venugopalan, The Color glass condensate and high-energy scattering in QCD, in  _Quark-gluon plasma 4_, edited by R. C. Hwa and X.-N. Wang (2003) pp. 249–3363, arXiv:hep-ph/0303204 .
  * Kasmaei and Strickland [2019] B. S. Kasmaei and M. Strickland, Dilepton production and elliptic flow from an anisotropic quark-gluon plasma, Phys. Rev. D 99, 034015 (2019), arXiv:1811.07486 [hep-ph] .
  * Kasmaei and Strickland [2020] B. S. Kasmaei and M. Strickland, Photon production and elliptic flow from a momentum-anisotropic quark-gluon plasma, Phys. Rev. D 102, 014037 (2020), arXiv:1911.03370 [hep-ph] .
  * Kurkela and Mazeliauskas [2019] A. Kurkela and A. Mazeliauskas, Chemical equilibration in weakly coupled QCD, Phys. Rev. D 99, 054018 (2019), arXiv:1811.03068 [hep-ph] .
  * Du and Schlichting [2021b] X. Du and S. Schlichting, Equilibration of weakly coupled QCD plasmas, Phys. Rev. D 104, 054011 (2021b), arXiv:2012.09079 [hep-ph] .
  * Borsanyi  _et al._ [2014] S. Borsanyi, Z. Fodor, C. Hoelbling, S. D. Katz, S. Krieg, and K. K. Szabo, Full result for the QCD equation of state with 2+1 flavors, Phys. Lett. B 730, 99 (2014), arXiv:1309.5258 [hep-lat] .
  * Bazavov  _et al._ [2014] A. Bazavov _et al._ (HotQCD), Equation of state in ( 2+1 )-flavor QCD, Phys. Rev. D 90, 094503 (2014), arXiv:1407.6387 [hep-lat] .
  * Du [2024] X. Du, Heavy quark drag and diffusion coefficients in the prehydrodynamic QCD plasma, Phys. Rev. C 109, 014901 (2024), arXiv:2306.02530 [hep-ph] .
  * Andres  _et al._ [2023] C. Andres, L. Apolinário, F. Dominguez, M. G. Martinez, and C. A. Salgado, Medium-induced radiation with vacuum propagation in the pre-hydrodynamics phase, JHEP 03, 189, arXiv:2211.10161 [hep-ph] .
  * Kurkela and Zhu [2015] A. Kurkela and Y. Zhu, Isotropization and hydrodynamization in weakly coupled heavy-ion collisions, Phys. Rev. Lett. 115, 182301 (2015), arXiv:1506.06647 [hep-ph] .
  * Kurkela  _et al._ [2019a] A. Kurkela, A. Mazeliauskas, J.-F. Paquet, S. Schlichting, and D. Teaney, Effective kinetic description of event-by-event pre-equilibrium dynamics in high-energy heavy-ion collisions, Phys. Rev. C 99, 034910 (2019a), arXiv:1805.00961 [hep-ph] .
  * Kurkela and Lu [2014] A. Kurkela and E. Lu, Approach to Equilibrium in Weakly Coupled Non-Abelian Plasmas, Phys. Rev. Lett. 113, 182301 (2014), arXiv:1405.6318 [hep-ph] .
  * Du  _et al._ [2022] X. Du, M. P. Heller, S. Schlichting, and V. Svensson, Exponential approach to the hydrodynamic attractor in Yang-Mills kinetic theory, Phys. Rev. D 106, 014016 (2022), arXiv:2203.16549 [hep-ph] .
  * Giacalone  _et al._ [2019] G. Giacalone, A. Mazeliauskas, and S. Schlichting, Hydrodynamic attractors, initial state energy and particle production in relativistic nuclear collisions, Phys. Rev. Lett. 123, 262301 (2019), arXiv:1908.02866 [hep-ph] .
  * Kurkela  _et al._ [2019b] A. Kurkela, A. Mazeliauskas, J.-F. Paquet, S. Schlichting, and D. Teaney, Matching the Nonequilibrium Initial Stage of Heavy Ion Collisions to Hydrodynamics with QCD Kinetic Theory, Phys. Rev. Lett. 122, 122302 (2019b), arXiv:1805.01604 [hep-ph] .
  * Nijs  _et al._ [2021] G. Nijs, W. van der Schee, U. Gürsoy, and R. Snellings, Transverse Momentum Differential Global Analysis of Heavy-Ion Collisions, Phys. Rev. Lett. 126, 202301 (2021), arXiv:2010.15130 [nucl-th] .
  * Giacalone  _et al._ [2023] G. Giacalone, G. Nijs, and W. van der Schee, Determination of the Neutron Skin of Pb208 from Ultrarelativistic Nuclear Collisions, Phys. Rev. Lett. 131, 202302 (2023), arXiv:2305.00015 [nucl-th] .
  * Eskola  _et al._ [2017] K. J. Eskola, P. Paakkinen, H. Paukkunen, and C. A. Salgado, EPPS16: Nuclear parton distributions with LHC data, Eur. Phys. J. C 77, 163 (2017), arXiv:1612.05741 [hep-ph] .
  * Camarda  _et al._ [2020] S. Camarda _et al._ , DYTurbo: Fast predictions for Drell-Yan processes, Eur. Phys. J. C 80, 251 (2020), [Erratum: Eur.Phys.J.C 80, 440 (2020)], arXiv:1910.07049 [hep-ph] .
  * Endres  _et al._ [2015] S. Endres, H. van Hees, J. Weil, and M. Bleicher, Coarse-graining approach for dilepton production at energies available at the CERN Super Proton Synchrotron, Phys. Rev. C 91, 054911 (2015), arXiv:1412.1965 [nucl-th] .
  * Endres  _et al._ [2017] S. Endres, H. van Hees, and M. Bleicher, Probing the QCD phase diagram with dileptons - a study using coarse-grained transport dynamics, Acta Phys. Polon. Supp. 10, 549 (2017), arXiv:1609.09140 [nucl-th] .
  * [44] Kømpøst, https://github.com/KMPST/KoMPoST [Accessed: 2023-08-17)].
  * Uphoff  _et al._ [2010] J. Uphoff, O. Fochler, Z. Xu, and C. Greiner, Heavy quark production at RHIC and LHC within a partonic transport model, Phys. Rev. C 82, 044906 (2010), arXiv:1003.4200 [hep-ph] .
  * Zhang  _et al._ [2008] B.-W. Zhang, C.-M. Ko, and W. Liu, Thermal charm production in a quark-gluon plasma in Pb-Pb collisions at s**(1/2)(NN) = 5.5-TeV, Phys. Rev. C 77, 024901 (2008), arXiv:0709.1684 [nucl-th] .
  * Zhou  _et al._ [2016] K. Zhou, Z. Chen, C. Greiner, and P. Zhuang, Thermal Charm and Charmonium Production in Quark Gluon Plasma, Phys. Lett. B 758, 434 (2016), arXiv:1602.01667 [hep-ph] .
  * Boguslavski  _et al._ [2020] K. Boguslavski, A. Kurkela, T. Lappi, and J. Peuron, Heavy quark diffusion in an overoccupied gluon plasma, JHEP 09, 077, arXiv:2005.02418 [hep-ph] .
  * Zhou  _et al._ [2024] F. Zhou, J. Brewer, and A. Mazeliauskas, Minijet quenching in non-equilibrium quark-gluon plasma, (2024), arXiv:2402.09298 [hep-ph] .


# Diffusion probabilistic models enhance variational autoencoder for crystal
structure generative modeling

Teerachote Pakornchote  Natthaphon Choomphon-anomakhun  Sorrjit Arrerut
Extreme Conditions Physics Research Laboratory (ECPRL) and Center of
Excellence in Physics of Energy Materials (CE:PEM), Department of Physics,
Faculty of Science, Chulalongkorn University, Bangkok, 10330, Thailand
Chayanon Atthapak  Sakarn Khamkaeo  Extreme Conditions Physics Research
Laboratory (ECPRL) and Center of Excellence in Physics of Energy Materials
(CE:PEM), Department of Physics, Faculty of Science, Chulalongkorn University,
Bangkok, 10330, Thailand  Thailand Center of Excellence in Physics , Ministry
of Higher Education, Science, Research and Innovation, 328 Si Ayutthaya Road,
Bangkok 10400, Thailand  Thiparat Chotibut  Chula Intelligent and Complex
Systems, Department of Physics, Faculty of Science, Chulalongkorn University,
Bangkok, Thailand, 10330  Thiti Bovornratanaraks<EMAIL_ADDRESS>Extreme
Conditions Physics Research Laboratory (ECPRL) and Center of Excellence in
Physics of Energy Materials (CE:PEM), Department of Physics, Faculty of
Science, Chulalongkorn University, Bangkok, 10330, Thailand  Thailand Center
of Excellence in Physics , Ministry of Higher Education, Science, Research and
Innovation, 328 Si Ayutthaya Road, Bangkok 10400, Thailand

###### Abstract

The crystal diffusion variational autoencoder (CDVAE) is a machine learning
model that leverages score matching to generate realistic crystal structures
that preserve crystal symmetry. In this study, we leverage novel diffusion
probabilistic (DP) models to denoise atomic coordinates rather than adopting
the standard score matching approach in CDVAE. Our proposed DP-CDVAE model can
reconstruct and generate crystal structures whose qualities are statistically
comparable to those of the original CDVAE. Furthermore, notably, when
comparing the carbon structures generated by the DP-CDVAE model with relaxed
structures obtained from density functional theory calculations, we find that
the DP-CDVAE generated structures are remarkably closer to their respective
ground states. The energy differences between these structures and the true
ground states are, on average, 68.1 meV/atom lower than those generated by the
original CDVAE. This significant improvement in the energy accuracy highlights
the effectiveness of the DP-CDVAE model in generating crystal structures that
better represent their ground-state configurations.

Keywords

denoising diffusion probabilistic models, variational autoencoder,
crystal structures, ground state, density functional theory

Open data

The dataset used in this work has been made available at
https://github.com/trachote/dp-cdvae

††preprint: APS/123-QED

##  I Introduction

Advances in computational materials science has enabled the accurate
prediction of novel materials possessing exceptional properties. Remarkably,
these computational advancements have facilitated the successful experimental
synthesis of materials that exhibit the anticipated properties. Some predicted
materials, such as near-room-temperature superconductors, have been
successfully synthesized under high-pressure conditions, with their
superconducting temperatures in accordance with density functional theory
(DFT) calculations [1, 2]. To achieve accurate predictions, a priori knowledge
of plausible molecular and crystal structures play a vital role in both
theoretical and experimental studies. Several algorithms, such as evolutionary
algorithms, swarm particle optimization, random sampling method, and etc.,
have been employed for structure prediction [3, 4, 5]. These algorithms rely
on identifying local minima on the potential energy landscape obtained from
DFT calculations [6, 7]. In the case of crystal structures, where atoms are
arranged in three-dimensional space with periodic boundaries, additional
criteria are necessary to enforce crystal symmetry constraints [5].

Recent approach for structure prediction employs denoising diffusion models to
perform probabilistic inference. These models sample molecular and crystal
structures from a probability distribution of atomic coordinates and types [8,
9, 10, 11], bypassing the computationally intensive DFT calculation to
tediously determine the potential energy landscape. By leveraging sufficiently
large datasets containing various compounds, this method enables the
generation of diverse compositions and combinations of elements
simultaneously. Furthermore, the models allow for the control of desired
physical properties of the generated structures through conditional
probability sampling [12, 13, 14]. These machine learning-based algorithms
also hold promise for solving inverse problem efficiently, resolving
structures from experimental characterizations, e.g., x-ray absorption
spectroscopy and other techniques, a challenging problem in materials science
[15, 16, 17].

There are two primary types of denoising diffusion models: score matching
approach and denoising diffusion probabilistic models (DDPM) [18, 19, 20].
These two models can denoise (reverse) a normal distribution such that the
distribution gradually transforms into the data distribution of interest. The
score matching approach estimates the score function of the perturbed data
directing the normal distribution toward the data distribution and employing
large step sizes of variance. In contrast, DDPM gradually denoises the random
noise through a joint distribution of data perturbed at different scales of
variance.

Since atomic positions in crystal structures are periodic and can be invariant
under some rotation groups depending on their crystal symmetry, the core
neural networks should favourably possess roto-translational equivariance [21,
22, 23]. Xie et al. [8] has proposed a model for crystal prediction by a
combination between variational autoencoder (VAE) [24] and the denoising
diffusion model, called crystal diffusion VAE (CDVAE). The model employs the
score matching approach with (annealed) Langevin dynamics to generate new
crystal structures from random coordinates [18]. The neural networks for an
encoder and the diffusion model are roto-translationally equivariant. As a
result, CDVAE can generate crystal structures with realistic bond lengths and
respect crystal symmetry.

Because of the periodic boundary condition imposed on the unit cell, gradually
injecting sufficiently strong noises (in the forward process) to the
fractional coordinates can lead to the uniform distribution of atomic
positions at late times, the consequence of ergodicity in statistical
mechanical sense. Rather than beginning with a Gaussian distribution and
denoising it as in the original CDVAE formulation, Jiao et al. [25] perturbed
and sampled atomic positions beginning with a wrapped normal distribution
which satisfies the periodic boundary condition. With this approach, the
reconstruction performance has been significantly improved. Other circular
(periodic) distributions, e.g., the wrapped normal and von Mises
distributions, are not natural for DDPM framework since there is no known
analytical method to explicitly incorporate such distributions into the
framework. There, one needs to resort to an additional sampling procedure to
construct the DDPM [26].

In this work, we introduce a crystal generation framework called diffusion
probabilistic CDVAE (DP-CDVAE). Similar to the original CDVAE, our model
consists of two parts: the VAE part and the diffusion part. The purpose of the
VAE part is to predict the lattice parameters and the number of atoms in the
unit cell of crystal structures. On the other hand, the diffusion part
utilizes the diffusion probabilistic approach to denoise fractional
coordinates and predict atomic coordinates. By employing the DDPM instead of
the score matching approach, the DP-CDVAE model shows reconstruction and
generation task performances that are statistically comparable to those
obtained from original CDVAE. Importantly, we demonstrate the significantly
higher ground-state generation performance of DP-CDVAE, through the distance
comparison between generated structures and those optimized using the DFT
method. We also analyze the changes in energy and volume after relaxation to
gain further insights into models’ capabilities.

Figure 1: The schematic summarizing the architecture for training the DP-CDVAE
model. Multiple sub-networks are trained to minimize the total loss function
of Eq.(S1). The encoder ($G_{\phi}(\mathbf{L}\mathbf{r}_{f},Z,N_{a})$)
compresses input pristine crystal structures into the latent feature
($\mathbf{z}$). The predicted lattice parameters ($\mathbf{L}_{\mathbf{z}}$),
the predicted number of atoms ($N_{\mathbf{z}}$), and
$\mathbf{A}_{\mathbf{z}}$ are decoded from $\mathbf{z}$. Here,
$\mathbf{A}_{\mathbf{z}}$ enables the sampling of atomic types ($Z_{t}$), and
all the decoded features enable the reconstruction of crystal structures. The
input fractional coordinates $\mathbf{r}_{f}$ undergo perturbation (dash-
dotted line) at time step $t$ and then are transformed by $\bm{\pi}(\cdot)$ to
satisfy the periodic boundary condition (dotted line), serving as the
coordinates for the reconstructed crystal structures. These reconstructed
structures, $\left(\mathbf{L_{z}r}_{f_{t}},Z_{t},\mathbf{z},t\right)$, are
subsequently fed into the diffusion network
($D_{\theta}(\mathbf{L}_{\mathbf{z}}\mathbf{r}_{f_{t}},\mathbf{f}_{t})$),
where $\mathbf{f}_{t}$ is a node feature composing of $Z_{t}$, $\mathbf{z}$,
and $t$. The diffusion network predicts the noise added to the fractional
coordinates ($\bm{\epsilon}_{\theta}$) as well as the one-hot vector of atomic
types ($\mathbf{A}_{\theta}$), see Sec. IV.3. Dashed-line boxes represent the
unit cells of the crystal structures.

##  II Results

The performances of DP-CDVAE models are herein presented. There are four DP-
CDVAE models, differed by the choice of the encoder (see Fig. S1). DimeNet++
has been employed for the main encoder for every DP-CDVAE models [27]. We then
modify the encoder of DP-CDVAE to encode the crystal structure by two
additional neural networks: a multilayer perceptron that takes the number of
atoms in the unit cell ($N_{a}$) as an input, and a graph isomorphism network
(GINE) [28]. Their latent features are combined with the latent features from
DimeNet++ through another multilayer perceptron. The $N_{a}$ is encoded such
that the model can decode the $N_{a}$ accurately, and GINE encoder is inspired
by GeoDiff [10] whose model is a combination of SchNet [29] and GINE which
yields better performance.

Three datasets, Perov-5 [30, 31], Carbon-24 [32], and MP-20 [33], were
selected to evaluate the performance of the model. The Perov-5 dataset
consists of perovskite materials with cubic structures, but with variations in
the combinations of elements within the structures. The Carbon-24 dataset
comprises carbon materials, where the data consists of carbon element with
various crystal systems obtained from ab initio random structure searching
algorithm at pressure of 10 GPa [32]. The MP-20 dataset encompasses a wide
range of compounds and structure types.

###  II.1 Reconstruction performance

The reconstruction performance is determined by the similarity between
reconstructed and ground-truth structures. The similarity can be evaluated
using Niggli’s algorithm implemented in StructureMatcher method from pymatgen
library [34]. The reconstructed and ground-truth structures are similar if
they pass the criteria of StructureMatcher which are stol=0.5, angle_tol=10,
ltol=0.3. Match rate is the percentage of those structures passed the
criteria. If the reconstructed and ground-truth structures are similar under
the criteria, the root-mean-square distance between their atomic positions is
computed and then normalized by $\sqrt[3]{V/N_{a}}$, where $V$ is the unit-
cell volume, and $N_{a}$ is the number of atoms in the unit cell. An average
of the distances of every pair of structures
($\langle\delta_{\text{rms}}\rangle$), computed from Niggli’s algorithm, is
used as the performance metric.

Table 1 presents the reconstruction performance of different models for three
different datasets: Perov-5, Carbon-24, and MP-20. For the Perov-5 dataset,
the DP-CDVAE model achieves a match rate of 90.04%, indicating its ability to
reconstruct a significant portion of the ground-truth structures. This
performance is slightly lower than the CDVAE model but still demonstrates the
effectiveness of our model. In terms of $\langle\delta_{\text{rms}}\rangle$,
the DP-CDVAE model achieves a value of 0.0212, comparable to the FTCP model
[35], but slightly higher than the CDVAE model. Similarly, for the Carbon-24
and MP-20 datasets, the DP-CDVAE model performs well in terms of both match
rate and $\langle\delta_{\text{rms}}\rangle$. It achieves match rates of
45.57% and 32.42% for Carbon-24 and MP-20, respectively. The corresponding
$\langle\delta_{\text{rms}}\rangle$ values for Carbon-24 and MP-20 are 0.1513
and 0.0383, respectively, comparable to the CDVAE model.

Regarding the DP-CDVAE+$N_{a}$ model, the additional encoding of $N_{a}$ into
the model leads to improved match rates for all datasets, with an increase of
2–5%. This enhancement can be attributed to the accurate prediction of
$N_{a}$. However, in terms of $\langle\delta_{\text{rms}}\rangle$, only the
Perov-5 dataset shows an improvement, with a value of 0.0149. On the other
hand, for the Carbon-24 and MP-20 datasets, the
$\langle\delta_{\text{rms}}\rangle$ values are higher compared to the DP-CDVAE
model.

For the DP-CDVAE+GINE and DP-CDVAE+$N_{a}$+GINE models, the additional
encoding of GINE into the models leads to a substantial drop in match rates
compared to the DP-CDVAE model, particularly for the Perov-5 and Carbon-24
datasets. In contrast, there is a moderate increase in the match rates for the
MP-20 dataset. The $\langle\delta_{\text{rms}}\rangle$ values for the Perov-5
and Carbon-24 datasets are comparable to those of the DP-CDVAE model. However,
for the MP-20 dataset, the $\langle\delta_{\text{rms}}\rangle$ is noticeably
higher in the models with GINE encoder compared to the DP-CDVAE model.

Overall, while the reconstruction performance of the DP-CDVAE model may be
slightly lower than the CDVAE model in terms of match rate, it still
demonstrates competitive performance with relatively low
$\langle\delta_{\text{rms}}\rangle$. The match rate can be enhanced by
additionally encoding the $N_{a}$, but the performance is traded off by the
increase in $\langle\delta_{\text{rms}}\rangle$.

Table 1: Reconstruction performance. Models | Match rate (%) $\uparrow$ |  $\langle\delta_{\text{rms}}\rangle$ $\downarrow$
---|---|---
| Perov-5 | Carbon-24 | MP-20 | Perov-5 | Carbon-24 | MP-20
FTCP [8] | 99.34 | 62.28 | 69.89 | 0.0259 | 0.2563 | 0.1593
CDVAE [8] | 97.52 | 55.22 | 45.43 | 0.0156 | 0.1251 | 0.0356
DP-CDVAE | 90.04 | 45.57 | 32.42 | 0.0212 | 0.1513 | 0.0383
DP-CDVAE+$N_{a}$ | 91.86 | 50.99 | 36.17 | 0.0149 | 0.1612 | 0.0560
DP-CDVAE+GINE | 80.50 | 49.02 | 34.08 | 0.0214 | 0.1599 | 0.0455
DP-CDVAE+$N_{a}$+GINE | 88.30 | 38.28 | 37.44 | 0.0180 | 0.1921 | 0.0525

###  II.2 Generation performance

We follow the CDVAE model that used three metrics to determine the generation
performance of the models [8]. The first metric is Validity percentage. The
structures that are valid for structure and composition tests must satisfy two
criteria: distances of every pair of atoms are larger than 0.5 Å and the total
charge in the unit cell is neutral. The second metric is called coverage
(COV), which utilizes structure and composition fingerprints to evaluate the
similarity between the generated and ground-truth structures. COV-R (Recall)
represents the percentage of ground-truth structures covered by the generated
structures. COV-P (Precision) represents the percentage of generated
structures that are similar to the ground-truth structures, indicating the
quality of the generation. The third metric is the Wasserstein distance
between property distributions of generated and ground-truth structures. Three
property statistics are density ($\rho$), which is total atomic mass per
volume (unit g/cm3), formation energy ($E_{form}$, unit eV/atom), and the
number of elements in the unit cell (# elem.). A separated and pre-trained
neural network is employed to predict $E$ of the structures where the detail
of the pre-training can be found in Ref. [8]. The first and second metrics are
computed over 10,240 generated structures, and 1000 structures are randomly
chosen from the generated structures that pass the validity tests to compute
the third metric. The ground-truth structures used to evaluate the generation
performance are from the test set.

In Table 2, the DP-CDVAE model achieves a validity rate of 100% for the
Perov-5 dataset and close to 100% for the Carbon-24 and MP-20 datasets in
terms of structure. The validity rate for composition is comparable to that of
the CDVAE model. The DP-CDVAE model also demonstrates comparable COV-R values
to the CDVAE model across all three datasets. Furthermore, the DP-CDVAE models
with $N_{a}$ and/or GINE encoders exhibit similar Validity and COV-R metrics
to those of the DP-CDVAE model. However, for COV-P, all DP-CDVAE models yield
lower values compared to CDVAE.

On the other hand, our models show significant improvements in property
statistics. In the case of the MP-20 dataset, the DP-CDVAE models,
particularly those with the GINE encoder, yield substantially smaller
Wasserstein distances for $\rho$, $E_{form}$, and the number of elements
compared to other models. For the Carbon-24 dataset, our models also exhibit a
smaller Wasserstein distance for $\rho$ compared to the CDVAE model.

###  II.3 Ground-state performance

Another objective of the structure generator is to generate novel structures
that also are close to the ground state. To verify that, the generated
structures are relaxed using the DFT calculation where the relaxed structures
exhibit balanced internal stresses with external pressures and reside in local
energy minima. These relaxed structures are then compared with the generated
structures to evaluate their similarity. In this study, we have chosen a set
of 100 generated structures from each of CDVAE, CDVAE+Fourier, and DP-CDVAE
models for relaxation where CDVAE+Fourier model is CDVAE model with Fourier
embedding features of the perturbed coordinates. However, relaxation
procedures for multi-element compounds can be computationally intensive. To
address this, we have specifically selected materials composed solely of
carbon atoms, using the model trained on Carbon-24 dataset. This selection
ensures a convergence of the self-consistent field in DFT calculation.
Moreover, in the relaxation, we consider the ground state of the relaxed
structures at a temperature of 0 K and a pressure of 10 GPa since the carbon
structures in the training set are stable at 10 GPa [32].

We here introduce a ground-state performance presented in Table 3. The
StructureMatcher with the same criteria as in the reconstruction performance
is used to evaluate the similarity between the generated and relaxed
structures. The relaxed structure was used as a based structure to determine
if the generated structure can be matched. Four metrics used to determine the
similarity are 1) match rate, 2) $\langle\delta_{\text{rms}}\rangle$, 3)
$\Delta V_{\text{rms}}$ and 4) $\Delta E_{\text{rms}}$. The $\Delta
V_{\text{rms}}$ and $\Delta E_{\text{rms}}$ represent the root mean square
differences in volume and energy, respectively, between the generated
structures and the relaxed structures in the dataset.

In Table 3, the DP-CDVAE model achieves the highest match rate and the lowest
$\langle\delta_{\text{rms}}\rangle$ and $\Delta E_{\text{rms}}$. Although the
CDVAE+Fourier model achieves the lowest $\Delta V_{\text{rms}}$, the DP-CDVAE
model demonstrates the $\Delta V_{\text{rms}}$ that is comparable to that of
the CDVAE+Fourier model.

Table 2: Generation performance Datasets | Models | Validity (%) $\uparrow$ | COV (%) $\uparrow$ | Property statistics $\downarrow$
---|---|---|---|---
Struc. | Comp. | R. | P. | $\rho$ | $E_{form}$ | # elem.
Perov-5 | G-SchNet [8] | 99.92 | 98.79 | 0.18 | 0.23 | 1.625 | 4.746 | 0.0368
P-G-SchNet [8] | 79.63 | 99.13 | 0.37 | 0.25 | 0.2755 | 1.388 | 0.4552
CDVAE [8] | 100 | 98.59 | 99.45 | 98.46 | 0.1258 | 0.0264 | 0.0628
DP-CDVAE | 100 | 98.07 | 99.52 | 98.39 | 0.1807 | 0.0713 | 0.0767
DP-CDVAE+$N_{a}$ | 99.99 | 97.34 | 99.55 | 97.22 | 0.1027 | 0.0287 | 0.0437
DP-CDVAE+GINE | 100 | 96.11 | 98.94 | 95.63 | 0.2114 | 0.0832 | 0.0498
DP-CDVAE+$N_{a}$+GINE | 100 | 97.09 | 99.52 | 96.73 | 0.1368 | 0.0425 | 0.0210
Carbon-24 | G-SchNet [8] | 99.94 | – | 0.00 | 0.00 | 0.9427 | 1.320 | –
P-G-SchNet [8] | 48.39 | – | 0.00 | 0.00 | 1.533 | 134.7 | –
CDVAE [8] | 100 | – | 99.80 | 83.08 | 0.1407 | 0.2850 | –
DP-CDVAE | 99.92 | – | 99.56 | 77.98 | 0.1109 | 0.2596 | –
DP-CDVAE+$N_{a}$ | 99.73 | – | 99.61 | 72.29 | 0.1080 | 0.3030 | –
DP-CDVAE+GINE | 99.50 | – | 100 | 68.13 | 0.0977 | 0.3623 | –
DP-CDVAE+$N_{a}$+GINE | 98.61 | – | 99.21 | 65.13 | 0.1267 | 0.4136 | –
MP-20 | G-SchNet [8] | 99.65 | 75.96 | 38.33 | 99.57 | 3.034 | 42.09 | 0.6411
P-G-SchNet [8] | 77.51 | 76.40 | 41.93 | 99.74 | 4.04 | 2.448 | 0.6234
CDVAE [8] | 100 | 86.70 | 99.15 | 99.49 | 0.6875 | 0.2778 | 1.432
DP-CDVAE | 99.59 | 85.44 | 98.93 | 98.96 | 0.4037 | 0.1547 | 0.9179
DP-CDVAE+$N_{a}$ | 99.81 | 84.95 | 99.36 | 99.33 | 0.4889 | 0.1800 | 1.053
DP-CDVAE+GINE | 99.82 | 81.92 | 99.48 | 99.00 | 0.2785 | 0.0603 | 0.5679
DP-CDVAE+$N_{a}$+GINE | 99.90 | 83.89 | 95.51 | 99.27 | 0.1790 | 0.0522 | 0.6909
Table 3: Ground-state performance. Model | Match rate (%) $\uparrow$ |  $\langle\delta_{\text{rms}}\rangle$ $\downarrow$ |  $\Delta V_{\text{rms}}$ (Å3/atom) $\downarrow$ |  $\Delta E_{\text{rms}}$ (meV/atom) $\downarrow$
---|---|---|---|---
CDVAE | 63 | 0.0321 | 0.0227 | 468.8
CDVAE+Fourier | 62 | 0.0216 | 0.0157 | 494.4
DP-CDVAE | 64 | 0.0141 | 0.0158 | 400.7

##  III Discussion

The DP-CDVAE models significantly enhance the generation performance,
particularly in terms of property statistics, while maintaining comparable
COVs to those of CDVAE. Specifically, for Carbon-24 and MP-20 datasets, the
density distributions between the generated and ground-truth structures from
DP-CDVAE models exhibit substantially smaller Wasserstein distance compared
those of the CDVAE model (see Table 2). The $\Delta V_{\text{rms}}$ of the DP-
CDVAE model presented in Table 3 is significantly lower than that of the
original CDVAE. This is corresponding to smaller Wasserstein distance of
$\rho$ shown in Table 2. The DP-CDVAE model also demonstrates significantly
smaller $\langle\delta_{\text{rms}}\rangle$ than the original CDVAE. These
suggest that our lattice generation closely approximates the relaxed lattice,
while also achieving atomic positions that closely resemble the ground-state
configuration. Additionally, the distribution of the number of elements in the
unit cells is relatively similar to that of the data in the test set,
particularly in the results from the models with GINE encoder. This could be
attributed to the capability of GINE to search for graph isomorphism [36].

Moreover, $\Delta E$ is the energy difference between the generated structures
and their corresponding relaxed structures. The ground-state energy represents
a local minimum that the generated structure is relaxed towards. A value of
$\Delta E$ close to zero indicates that the generated structure is in close
proximity to the ground state. In Table 3, it can be observed that our model
achieves the $\Delta E_{\text{rms}}$ value of 400.7 meV/atom which is about
68.1 meV/atom lower than the $\Delta E_{\text{rms}}$ of CDVAE. The mode of
$\Delta E$ of our model is 64 – 128 meV/atom, which is lower than its root-
mean-square value (see Fig. S2). Nevertheless, both the $\Delta
E_{\text{rms}}$ and the mode of $\Delta E$ exhibit relatively high values. In
many cases, the formation energy of synthesized compounds is reported to be
above the convex hull less than 36 meV/atom [37, 38, 39]. To obviate the need
for time-consuming DFT relaxation, it is essential for the generated
structures to be even closer to the ground state. Therefore, achieving lower
$\Delta E_{\text{rms}}$ values remains a milestone for future work.

##  IV Methods

###  IV.1 Diffusion probabilistic model

In the diffusion probabilistic model, the data distribution is gradually
perturbed by noise in the forward process until it becomes a normal
distribution at late times. In this study, the distribution of the fractional
coordinate ($\mathbf{r}_{f}$) is considered since their values of every
crystal structures distribute over the same range ,i.e.,
$\mathbf{r}_{f}\in[0,1)^{3}$. The Markov process is assumed for the forward
diffusion such that the joint distribution is a product of the conditional
distributions conditioned on the knowledge of the fractional coordinate at the
previous time step:

$\displaystyle q(\mathbf{r}_{1:T}|\mathbf{r}_{0})$
$\displaystyle=\prod_{t=1}^{T}q(\mathbf{r}_{t}|\mathbf{r}_{t-1}),$ (1)
$\displaystyle q(\mathbf{r}_{t}|\mathbf{r}_{t-1})$
$\displaystyle=\mathcal{N}(\mathbf{r}_{t};\sqrt{\alpha_{t}}\mathbf{r}_{t-1},(1-\alpha_{t})\mathbf{I}),$

where $\mathbf{r}_{0}\sim q(\mathbf{r}_{f})$ which is the data distribution of
the fractional coordinate, $t$ is the discretized diffusion time step, $T$ is
the final diffusion time, $\alpha_{t}$ is a noise schedule with a sigmoid
scheduler [40], and the conditional $q(\cdot|\cdot)$ is a Gaussian kernel due
to the Markov diffusion process assumption. Then $\mathbf{r}_{t}$ can be
expressed in the Langevin’s form through the reparameterization trick as

$\mathbf{r}_{t}=\sqrt{\bar{\alpha}_{t}}\mathbf{r}_{0}+\sqrt{1-\bar{\alpha}_{t}}\bm{\epsilon},$
(2)

where $\bm{\epsilon}\sim\mathcal{N}(0,\mathbf{I})$, and
$\bar{\alpha}_{t}=\prod_{i=1}^{t}\alpha_{i}$. This update rule does not
necessitate $\mathbf{r}_{t}$ to remain in $[0,1)^{3}$; however, we can impose
the periodic boundary condition for the fractional coordinate so that

$\mathbf{r}_{f_{t}}=\bm{\pi}(\mathbf{r}_{t})\coloneqq\mathbf{r}_{t}-\lfloor\mathbf{r}_{t}\rfloor.$
(3)

Then, $\mathbf{r}_{f_{t}}\in[0,1)^{3}$.

In the reverse diffusion process, if the consecutive discretized time step is
small compared to the diffusion timescale, the reverse coordinate trajectories
can be approximately sampled also from the product of Gaussian diffusion
kernels as

$\displaystyle p_{\theta}(\mathbf{r}_{0:T})$
$\displaystyle=p(\mathbf{r}_{T})\prod_{t=1}^{T}p_{\theta}(\mathbf{r}_{t-1}|\mathbf{r}_{t}),$
(4) $\displaystyle p_{\theta}(\mathbf{r}_{t-1}|\mathbf{r}_{t})$
$\displaystyle=\mathcal{N}(\mathbf{r}_{t-1};\bm{\mu}_{\theta},\sigma^{2}_{t}\mathbf{I}),$

where

$\displaystyle\bm{\mu}_{\theta}$
$\displaystyle=\frac{1}{\sqrt{\alpha_{t}}}\Big{(}\mathbf{r}_{t}-\frac{1-\alpha_{t}}{\sqrt{1-\bar{\alpha}_{t}}}\bm{\epsilon}_{\theta}\Big{)},$
(5) $\displaystyle\sigma_{t}^{2}$
$\displaystyle=\frac{(1-\bar{\alpha}_{t-1})(1-\alpha_{t})}{1-\bar{\alpha}_{t}}.$

The reverse conditional distribution can be trained by minimizing the
Kullback–Leibler divergence between
$p_{\theta}(\mathbf{r}_{t-1}|\mathbf{r}_{t})$ and
$q(\mathbf{r}_{t-1}|\mathbf{r}_{t},\mathbf{r}_{0})$, the posterior of the
corresponding forward process [20]. We use GemNetT for the diffusion network
to train the parametrized noise $\bm{\epsilon}_{\theta}$ [41]. Then, the
coordinate in the earlier time can be sampled from $\mathbf{r}_{t-1}\sim
p_{\theta}(\mathbf{r}_{t-1}|\mathbf{r}_{t})$, whose corresponding reverse
Langevin’s dynamics reads

$\mathbf{r}_{t-1}=\frac{1}{\sqrt{\alpha_{t}}}\Big{(}\mathbf{r}_{t}-\frac{1-\alpha_{t}}{\sqrt{1-\bar{\alpha}_{t}}}\bm{\epsilon}_{\theta}\Big{)}+\sigma_{t}\bm{\epsilon}^{\scalebox{0.8}{\scriptsize{\hbox{\prime}}}},$
(6)

where
$\bm{\epsilon}^{\scalebox{0.8}{\scriptsize$\prime$}}\sim\mathcal{N}(0,\mathbf{I})$.
Crucially, we empirically found that the final reconstruction performance is
considerably improved when we impose the periodic boundary condition on the
fractional coordinate at every time step such that $\mathbf{r}_{t-1}\sim
p_{\theta}(\mathbf{r}_{t-1}|\mathbf{r}_{f_{t}})$ and $\alpha_{t}$ in the first
term of Eq. (6) is replaced by $\bar{\alpha}_{t}$. Namely, in our modified
reverse process, the coordinate is sampled from

$\displaystyle\mathbf{r}_{t-1}$
$\displaystyle=\frac{1}{\sqrt{\bar{\alpha_{t}}}}\Big{(}\mathbf{r}_{f_{t}}-\sqrt{1-\bar{\alpha}_{t}}\bm{\epsilon}_{\theta}\Big{)}+\sigma_{t}\bm{\epsilon}^{\scalebox{0.8}{\scriptsize$\prime$}},$
(7) $\displaystyle\mathbf{r}_{f_{t}}$
$\displaystyle=\bm{\pi}(\mathbf{r}_{t}).$

An illustration of denoising atomic coordinates with Eq. (7) is demonstrated
in Fig. 2. The model performance using Eq. (6) is shown in Table S1, whereas
the performance using Eq. (7) is shown in Table 1.

Figure 2:  The schematic depicting the reverse diffusion process of the DP-
CDVAE model. Initially, atomic coordinates are sampled from a normal
distribution and subsequently mapped into the unit cell (dashed-line box)
using the periodic boundary-imposing function $\bm{\pi}(\cdot)$. White circles
outside the unit cell depict the atomic coordinates prior to the periodic
boundary condition is imposed, while colored circles represent atoms that are
inside the unit cell of interest. The action of $\bm{\pi}(\cdot)$ on the atoms
outside the unit cell is represented by an arrow that translates the white
circles into colored circles in the unit cell. Left to right show the reverse
direction of the arrow of time, depicting the reverse diffusion process.

###  IV.2 Graph neural networks

Graph neural networks architecture facilitate machine learning of crystal
graphs $\mathcal{G}=(\mathcal{V},\mathcal{E})$, graph representations of
crystal structures. $\mathcal{V}$ and $\mathcal{E}$ are sets of nodes and
edges, respectively, defined as

$\displaystyle\mathcal{V}=\\{(\mathbf{f}_{n},\mathbf{r}_{c_{n}})\;|\;\mathbf{f}_{n}\in\mathbb{R}^{M},\;\mathbf{r}_{c_{n}}=\mathbf{L}\mathbf{r}_{f_{n}}\in\mathbb{R}^{3}\\},$
$\displaystyle\mathcal{E}=\\{\Delta\mathbf{r}_{c_{mn}}^{(\mathbf{T})}\;|\;\Delta\mathbf{r}_{c_{mn}}^{(\mathbf{T})}=\mathbf{r}_{c_{m}}-\mathbf{r}_{c_{n}}+\mathbf{T};\
\mathbf{r}_{c_{m}},\mathbf{r}_{c_{n}}\in\mathbb{R}^{3}\\},$

where $n$ and $m$ are indices of atoms in a crystal structure,
$\mathbf{f}_{n}$ is a vector of $M$ features of an atom in the unit cell,
$\mathbf{T}$ is a translation vector, and $\mathbf{L}$ is a lattice matrix
that converts a fractional coordinate $\mathbf{r}_{f_{n}}$ into its atomic
Cartesian coordinate $\mathbf{r}_{c_{n}}$. The atomic features, fractional
coordinates, and atomic Cartesian coordinates of the crystal structure are
vectorized (concatenated) as
$\mathbf{f}=(\mathbf{f}_{1},\dots,\mathbf{f}_{N_{a}})\in\mathbb{R}^{N_{a}\times
M}$,
$\mathbf{r}_{f}=(\mathbf{r}_{f_{1}},\dots,\mathbf{r}_{f_{N_{a}}})\in\mathbb{R}^{N_{a}\times
3}$, and
$\mathbf{r}_{c}=(\mathbf{r}_{c_{1}},\dots,\mathbf{r}_{c_{N_{a}}})\in\mathbb{R}^{N_{a}\times
3}$. Three graph neural networks implemented in this work are DimeNet++ [27],
GINE [28], and GemNetT [41]. DimeNet++ and GINE are employed for encoders, and
GemNetT is used for a diffusion network. DimeNet++ and GemNetT, whose based
architecture concerns geometry of the graphs, are rotationally equivariant.
GemNetT has been devised by incorporating the polar angles between four atoms
into DimeNet++. This development grants GemNetT a higher degree of expressive
power compared to DimeNet++ [42]. Furthermore, GINE has been developed to
distinguish a graph isomorphism, but not graph geometry nor the distance
between nodes, which is important for our study. Thus we supplement the edge
attributes into GINE with the distances between atoms, i.e.
$\mathcal{E}=\\{||\Delta\mathbf{r}_{c_{mn}}^{(\mathbf{T})}||\\}$.

###  IV.3 DP-CDVAE’s architecture

The forward process of DP-CDVAE model is illustrated in Fig. 1. The model is a
combination of two generative models, which are VAE and diffusion
probabilistic model. The pristine crystal structures consist of the fractional
coordinate ($\mathbf{r}_{f}$), the lattice matrix ($\mathbf{L}$), ground-truth
atomic type ($Z$), and the number of atoms in a unit cell ($N_{a}$). For
crystal graphs of the encoders, their node features are $\mathbf{f}=Z$. The
number of atoms in a unit cell $N_{a}$ is encoded through multilayer
perceptron before concatenated with the latent features from other graph
encoders. They are encoded to train $\bm{\mu}_{\phi}$ and $logvar_{\phi}$
where $\phi$ is a learnable parameter of the encoders. The latent variables
($\mathbf{z}$) can be obtained by

$\mathbf{z}=\bm{\mu}_{\phi}+e^{logvar_{\phi}}\bm{\epsilon}^{\scalebox{0.8}{\scriptsize{\hbox{\prime}}\scriptsize{\hbox{\prime}}}},$
(8)

where
$\bm{\epsilon}^{\scalebox{0.8}{\scriptsize$\prime$\scriptsize$\prime$}}\sim\mathcal{N}(0,\mathbf{I})$.
Then, $\mathbf{z}$ will be decoded to compute the lattice lengths and angles,
which then yield the lattice matrix ($\mathbf{L}_{\mathbf{z}}$), $N_{a}$, and
$\mathbf{A}_{\mathbf{z}}$. In the original CDVAE, $\mathbf{A}_{\mathbf{z}}$ is
the probability vector indicating the fraction of each atomic type in the
compound and is used to perturb $Z$ by

$Z_{t}\sim\mathcal{M}(\text{softmax}(\mathbf{A}+\sigma_{t}^{\scalebox{0.8}{\scriptsize{\hbox{\prime}}}}\mathbf{A}_{\mathbf{z}}))$
(9)

where $\mathcal{M}$ is a multinomial distribution, $\mathbf{A}$ is a one-hot
vector of ground-truth atomic type $Z$, and
$\sigma_{t}^{\scalebox{0.8}{\scriptsize$\prime$}}$ is the variance for
perturbing atomic types at time $t$, which is distinct from $\sigma_{t}$ used
for perturbing the atomic coordinates. Similar to the original CDVAE,
$\sigma_{t}^{\scalebox{0.8}{\scriptsize$\prime$}}$ is selected from the range
of [0.01, 5].

For the diffusion network, the input structures are constructed from
$\mathbf{r}_{f_{t}}$, $Z_{t}$, and $\mathbf{L}_{\mathbf{z}}$ where the
(Cartesian) atomic coordinates at time $t$ are computed by
$\mathbf{r}_{c_{t}}=\mathbf{L}_{\mathbf{z}}\mathbf{r}_{f_{t}}$. These are then
utilized by the crystal graphs for the diffusion network, whose node features
are $\mathbf{f}_{t}=(Z_{t},\mathbf{F}_{t},\mathbf{z},t)$ where
$\mathbf{F}_{t}$ is a Fourier embedding feature of $\mathbf{r}_{t}$ (see
section S3). As proposed by Ho et al. [20], we use the simple loss to train
the model such that

$\mathcal{L}_{simple}=\|\bm{\epsilon}-\bm{\epsilon}_{\theta}(\mathbf{r}_{c_{t}},\mathbf{f}_{t})\|^{2}.$
(10)

Since the diffusion model is trained to predict both $\bm{\epsilon}$ and
$\mathbf{A}$, so the loss of the diffusion network is

$\mathcal{L}_{diff}=\mathcal{L}_{simple}+\lambda\mathcal{L}_{CE}(\mathbf{A},\mathbf{A}_{\theta}(\mathbf{r}_{c_{t}},\mathbf{f}_{t})),$
(11)

where $\mathcal{L}_{CE}$ is the cross entropy loss, $\lambda$ is a loss
scaling factor, and $t\in\\{1,...,T\\}$ where $T=1000$. In this work, $t$ is
randomly chosen for each crystal graph and randomly reinitialized for each
epoch in the training process. The total loss in the trainig process is shown
in Eq. S1.

In the reverse diffusion process, we measure the model performance of two
tasks: reconstruction and generation tasks. For the former task, $\mathbf{z}$
is obtained from Eq. (8) by using the ground-truth structure as an input of
the encoders. For the latter task, $\mathbf{z}\sim\mathcal{N}(0,\mathbf{I})$,
which is then used to predict $N_{a}$, $\mathbf{L}_{\mathbf{z}}$,
$\mathbf{A}_{\mathbf{z}}$, and concatenate with the node feature of the
crystal graph in the diffusion network. At the initial step, $t=T$, $Z_{T}$ is
sampled from the highest probability of $\mathbf{A}_{\mathbf{z}}$, and the
final-time coordinate is obtained from sampling a Gaussian distribution, i.e.
$\mathbf{r}_{T}\sim\mathcal{N}(0,\mathbf{I}).$ The coordinates can be denoised
using Eq. (7), and the predicted atomic types are updated in each reversed
time step by $\text{argmax}(\mathbf{A}_{\theta})$.

###  IV.4 DFT calculations

The Vienna ab initio Simulation Package (VASP) was employed for structural
relaxations and energy calculations based on DFT [43, 44]. The calculations
were conducted under the generalized gradient approximation (GGA) and the
project augmented wave (PAW) method [45, 46]. The thresholds for energy and
force convergence were set to $10^{-5}$ eV and $10^{-5}$ eV/Å, respectively.
The plane-wave energy cutoff was set to 800 eV, and the Brillouin zone
integration was carried out on a k-point mesh of 5 × 5 × 5 created by the
Monkhorst-Pack method [47, 48].

###### Acknowledgements.

This research project is supported by the Second Century Fund (C2F),
Chulalongkorn University. This Research is funded by Thailand Science research
and Innovation Fund Chulalongkorn University (IND66230002) and National
Research Council of Thailand (NRCT): (NRCT5-RSA63001-04). T.C. acknowledges
funding support from the NSRF via the Program Management Unit for Human
Resources and Institutional Development, Research and Innovation [grant number
B05F650024]. The authors acknowledge high performance computing resources
including NVIDIA A100 GPUs from Chula Intelligent and Complex Systems Lab,
Faculty of Science, and from the Center for AI in Medicine (CU-AIM), Faculty
of Medicine, Chulalongkorn University, Thailand. We acknowledge the supporting
computing infrastructure provided by NSTDA, CU, CUAASC, NSRF via PMUB
[B05F650021, B37G660013] (Thailand). URL:www.e-science.in.th. The
Computational Materials Physics (CMP) Project, SLRI, Thailand, is acknowledged
for providing computational resource.

## References

  * Needs and Pickard [2016] R. J. Needs and C. J. Pickard, Perspective: Role of structure prediction in materials discovery and design, APL Materials 4, 053210 (2016).
  * Kohn and Sham [1965] W. Kohn and L. J. Sham, Phys. Rev. 140, A1133 (1965).
  * Oganov and Glass [2006] A. R. Oganov and C. W. Glass, Crystal structure prediction using ab initio evolutionary techniques: Principles and applications, The Journal of Chemical Physics 124, 244704 (2006).
  * Wang  _et al._ [2010] Y. Wang, J. Lv, L. Zhu, and Y. Ma, Crystal structure prediction via particle-swarm optimization, Phys. Rev. B 82, 094116 (2010).
  * Pickard and Needs [2011] C. J. Pickard and R. J. Needs, Ab initio random structure searching, Journal of Physics: Condensed Matter 23, 053201 (2011).
  * Oganov  _et al._ [2019] A. R. Oganov, C. J. Pickard, Q. Zhu, and R. J. Needs, Structure prediction drives materials discovery, Nature Reviews Materials 4, 331 (2019).
  * Schön  _et al._ [2010] J. C. Schön, K. Doll, and M. Jansen, Predicting solid compounds via global exploration of the energy landscape of solids on the ab initio level without recourse to experimental information, physica status solidi (b) 247, 23 (2010).
  * Xie  _et al._ [2022] T. Xie, X. Fu, O.-E. Ganea, R. Barzilay, and T. S. Jaakkola, Crystal diffusion variational autoencoder for periodic material generation, in  _International Conference on Learning Representations_ (2022).
  * Shi  _et al._ [2021] C. Shi, S. Luo, M. Xu, and J. Tang, Learning gradient fields for molecular conformation generation, in  _Proceedings of the 38th International Conference on Machine Learning_ , Proceedings of Machine Learning Research, Vol. 139, edited by M. Meila and T. Zhang (PMLR, 2021) pp. 9558–9568.
  * Xu  _et al._ [2022] M. Xu, L. Yu, Y. Song, C. Shi, S. Ermon, and J. Tang, Geodiff: A geometric diffusion model for molecular conformation generation, in  _International Conference on Learning Representations_ (2022).
  * Guan  _et al._ [2023] J. Guan, W. W. Qian, X. Peng, Y. Su, J. Peng, and J. Ma, 3d equivariant diffusion for target-aware molecule generation and affinity prediction, in  _The Eleventh International Conference on Learning Representations_ (2023).
  * Kang and Cho [2019] S. Kang and K. Cho, Conditional molecular design with deep generative models, Journal of Chemical Information and Modeling 59, 43 (2019), pMID: 30016587.
  * Lim  _et al._ [2018] J. Lim, S. Ryu, J. W. Kim, and W. Y. Kim, Molecular generative model based on conditional variational autoencoder for de novo molecular design, Journal of Cheminformatics 10, 31 (2018).
  * Song  _et al._ [2022] Y. Song, L. Shen, L. Xing, and S. Ermon, Solving inverse problems in medical imaging with score-based generative models, in  _International Conference on Learning Representations_ (2022).
  * Cui  _et al._ [2019] A. Cui, K. Jiang, M. Jiang, L. Shang, L. Zhu, Z. Hu, G. Xu, and J. Chu, Decoding phases of matter by machine-learning raman spectroscopy, Phys. Rev. Appl. 12, 054049 (2019).
  * Carbone  _et al._ [2020] M. R. Carbone, M. Topsakal, D. Lu, and S. Yoo, Machine-learning x-ray absorption spectra to quantitative accuracy, Phys. Rev. Lett. 124, 156401 (2020).
  * Liang  _et al._ [2023] Z. Liang, M. R. Carbone, W. Chen, F. Meng, E. Stavitski, D. Lu, M. S. Hybertsen, and X. Qu, Decoding structure-spectrum relationships with physically organized latent spaces, Phys. Rev. Mater. 7, 053802 (2023).
  * Song and Ermon [2019] Y. Song and S. Ermon, Generative modeling by estimating gradients of the data distribution, in  _Advances in Neural Information Processing Systems_, Vol. 32, edited by H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (Curran Associates, Inc., 2019).
  * Song  _et al._ [2021] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole, Score-based generative modeling through stochastic differential equations, in  _International Conference on Learning Representations_ (2021).
  * Ho  _et al._ [2020] J. Ho, A. Jain, and P. Abbeel, Denoising diffusion probabilistic models, in  _Advances in Neural Information Processing Systems_, Vol. 33, edited by H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin (Curran Associates, Inc., 2020) pp. 6840–6851.
  * Bronstein  _et al._ [2021] M. M. Bronstein, J. Bruna, T. Cohen, and P. Velickovic, Geometric deep learning: Grids, groups, graphs, geodesics, and gauges, CoRR abs/2104.13478 (2021), 2104.13478 .
  * Cohen  _et al._ [2018] T. S. Cohen, M. Geiger, J. Köhler, and M. Welling, Spherical CNNs, in  _International Conference on Learning Representations_ (2018).
  * Thomas  _et al._ [2018] N. Thomas, T. Smidt, S. Kearnes, L. Yang, L. Li, K. Kohlhoff, and P. Riley, Tensor field networks: Rotation- and translation-equivariant neural networks for 3d point clouds (2018).
  * Kingma and Welling [2014] D. P. Kingma and M. Welling, Auto-encoding variational bayes, in  _International Conference on Learning Representations_ (2014).
  * Jiao  _et al._ [2023] R. Jiao, W. Huang, P. Lin, J. Han, P. Chen, Y. Lu, and Y. Liu, Crystal structure prediction by joint equivariant diffusion on lattices and fractional coordinates, in  _Workshop on ”Machine Learning for Materials” ICLR 2023_ (2023).
  * Okhotin  _et al._ [2023] A. Okhotin, D. Molchanov, V. Arkhipkin, G. Bartosh, A. Alanov, and D. Vetrov, Star-shaped denoising diffusion probabilistic models (2023), arXiv:2302.05259 [stat.ML] .
  * Gasteiger  _et al._ [2022] J. Gasteiger, S. Giri, J. T. Margraf, and S. Günnemann, Fast and uncertainty-aware directional message passing for non-equilibrium molecules (2022), arXiv:2011.14115 [cs.LG] .
  * Hu*  _et al._ [2020] W. Hu*, B. Liu*, J. Gomes, M. Zitnik, P. Liang, V. Pande, and J. Leskovec, Strategies for pre-training graph neural networks, in  _International Conference on Learning Representations_ (2020).
  * Schütt  _et al._ [2017] K. Schütt, P.-J. Kindermans, H. E. Sauceda Felix, S. Chmiela, A. Tkatchenko, and K.-R. Müller, Schnet: A continuous-filter convolutional neural network for modeling quantum interactions, in  _Advances in Neural Information Processing Systems_, Vol. 30, edited by I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (Curran Associates, Inc., 2017).
  * Castelli  _et al._ [2012a] I. E. Castelli, D. D. Landis, K. S. Thygesen, S. Dahl, I. Chorkendorff, T. F. Jaramillo, and K. W. Jacobsen, New cubic perovskites for one-and two-photon water splitting using the computational materials repository, Energy & Environmental Science 5, 9034 (2012a).
  * Castelli  _et al._ [2012b] I. E. Castelli, T. Olsen, S. Datta, D. D. Landis, S. Dahl, K. S. Thygesen, and K. W. Jacobsen, Computational screening of perovskite metal oxides for optimal solar light capture, Energy & Environmental Science 5, 5814 (2012b).
  * [32] C. J. Pickard, .
  * Jain  _et al._ [2013] A. Jain, S. P. Ong, G. Hautier, W. Chen, W. D. Richards, S. Dacek, S. Cholia, D. Gunter, D. Skinner, G. Ceder, _et al._ , Commentary: The materials project: A materials genome approach to accelerating materials innovation, APL materials 1, 011002 (2013).
  * Grosse-Kunstleve  _et al._ [2004] R. W. Grosse-Kunstleve, N. K. Sauter, and P. D. Adams, Numerically stable algorithms for the computation of reduced unit cells, Acta Crystallographica Section A 60, 1 (2004).
  * Ren  _et al._ [2022] Z. Ren, S. I. P. Tian, J. Noh, F. Oviedo, G. Xing, J. Li, Q. Liang, R. Zhu, A. G. Aberle, S. Sun, X. Wang, Y. Liu, Q. Li, S. Jayavelu, K. Hippalgaonkar, Y. Jung, and T. Buonassisi, An invertible crystallographic representation for general inverse design of inorganic crystals with targeted properties, Matter 5, 314 (2022).
  * Xu  _et al._ [2019] K. Xu, W. Hu, J. Leskovec, and S. Jegelka, How powerful are graph neural networks?, in  _International Conference on Learning Representations_ (2019).
  * Wu  _et al._ [2013] Y. Wu, P. Lazic, G. Hautier, K. Persson, and G. Ceder, First principles high throughput screening of oxynitrides for water-splitting photocatalysts, Energy Environ. Sci. 6, 157 (2013).
  * Ishikawa and Miyake [2020] T. Ishikawa and T. Miyake, Evolutionary construction of a formation-energy convex hull: Practical scheme and application to a carbon-hydrogen binary system, Phys. Rev. B 101, 214106 (2020).
  * Ektarawong  _et al._ [2023] A. Ektarawong, E. Johansson, T. Pakornchote, T. Bovornratanaraks, and B. Alling, Boron vacancy-driven thermodynamic stabilization and improved mechanical properties of alb2-type tantalum diborides as revealed by first-principles calculations, Journal of Physics: Materials 6, 025002 (2023).
  * Kingma  _et al._ [2021] D. P. Kingma, T. Salimans, B. Poole, and J. Ho, On density estimation with diffusion models, in  _Advances in Neural Information Processing Systems_, edited by A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan (2021).
  * Klicpera  _et al._ [2021] J. Klicpera, F. Becker, and S. Günnemann, Gemnet: Universal directional graph neural networks for molecules, in  _Advances in Neural Information Processing Systems_, edited by A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan (2021).
  * Joshi  _et al._ [2023] C. K. Joshi, C. Bodnar, S. V. Mathis, T. Cohen, and P. Lio, On the expressive power of geometric graph neural networks (2023).
  * Kresse and Furthmüller [1996a] G. Kresse and J. Furthmüller, Computat. Mater. Sci. 6, 15 (1996a).
  * Kresse and Furthmüller [1996b] G. Kresse and J. Furthmüller, Phys. Rev. B 54, 11169 (1996b).
  * Perdew  _et al._ [1996] J. P. Perdew, K. Burke, and M. Ernzerhof, Phys. Rev. Lett. 77, 3865 (1996).
  * Blöchl [1994] P. E. Blöchl, Phys. Rev. B 50, 17953 (1994).
  * Monkhorst and Pack [1976] H. J. Monkhorst and J. D. Pack, Phys. Rev. B 13, 5188 (1976).
  * Pack and Monkhorst [1977] J. D. Pack and H. J. Monkhorst, Phys. Rev. B 16, 1748 (1977).
  * Gasteiger  _et al._ [2020] J. Gasteiger, J. Groß, and S. Günnemann, Directional message passing for molecular graphs, in  _International Conference on Learning Representations_ (2020).

Diffusion probabilistic models combined with variational autoencoder for
crystal structure generation

Supplementary Materials

##  S1 Encoders

Figure S1: (a) is the encoder for predicting $\bm{\mu}_{\phi}$ and
$logvar_{\phi}$. (b) and (d) are DimeNet++ and GINE encoders, respectively,
that take pristine crystal structures as inputs. (c) is the $N_{a}$ encoder
that takes the number of atoms $N_{a}$ as an input. (e), (f), (g), and (h) are
encoders of DP-CDVAE+$N_{a}$, DP-CDVAE+GINE, DP-CDVAE, and DP-
CDVAE+$N_{a}$+GINE models, respectively. Blue, yellow, and pink boxes are the
latent features from (b), (c), and (d), respectively, and FC boxes are fully
connected layers.

##  S2 Model evaluation

Table S1: Reconstruction performance using Eq. 6. Models | Match rate (%) $\uparrow$ |  $\langle\delta_{\text{rms}}\rangle$ $\downarrow$
---|---|---
| Perov-5 | Carbon-24 | MP-20 | Perov-5 | Carbon-24 | MP-20
DP-CDVAE | 68.72 | 35.86 | 20.40 | 0.0177 | 0.2573 | 0.0715
DP-CDVAE+$N_{a}$ | 67.56 | 38.13 | 21.71 | 0.0222 | 0.2719 | 0.0774
DP-CDVAE+GINE | 49.56 | 35.27 | 19.98 | 0.0822 | 0.2683 | 0.0620
DP-CDVAE+$N_{a}$+GINE | 64.12 | 34.68 | 24.43 | 0.0250 | 0.3212 | 0.0697
Figure S2:  Histograms of energy difference between generated and relaxed
structures of (a) CDVAE, (b) CDVAE+Fourier, and (c) DP-CDVAE models. The modes
(highest counting number) of CDVAE, CDVAE+Fourier, and DP-CDVAE models are
19.2 – 32.0, 68.0 – 128, and 64.0 – 128 eV/atom, respectively, where the bin
size is 64.0 meV/atom.

##  S3 Fourier features

We appended the Fourier features to the node attributes for the input of the
diffusion network. The Fourier features are the concatenation of
$\sin(2^{n}\pi\mathbf{r}_{t})$ and $\cos(2^{n}\pi\mathbf{r}_{t})$ where
$n\in\\{n_{\text{min}},...,n_{\text{max}}\\}$. For every DP-CDVAE model in the
main text, $n_{\text{min}}=3$ and $n_{\text{max}}=8$.

##  S4 Loss Functions

Similar to the original CDVAE, the total loss function to train the model
consists of 5 sub-loss functions: the Kullback–Leibler divergence loss
($\mathcal{L}_{KLD}$) for training $\bm{\mu}_{\phi}$ and $logvar_{\phi}$, the
lattice loss ($\mathcal{L}_{latt}$) for training lattice parameters, the
composition loss ($\mathcal{L}_{comp}$) for training
$\mathbf{A}_{\mathbf{z}}$, the loss for training the number of atoms
($\mathcal{L}_{N_{a}}$), and the loss from the diffusion network
($\mathcal{L}_{diff}$) for training $\bm{\epsilon}_{\theta}$ and
$\mathbf{A}_{\theta}$. In particular, the total loss is

$\mathcal{L}=\mathcal{L}_{diff}+\lambda_{1}\mathcal{L}_{KLD}+\lambda_{2}\mathcal{L}_{latt}+\lambda_{3}\mathcal{L}_{comp}+\lambda_{4}\mathcal{L}_{N_{a}},$
(S1)

where $\lambda_{1}$, $\lambda_{2}$, $\lambda_{3}$, and $\lambda_{4}$ are
tunable loss scaling factors, $\mathcal{L}_{latt}$ is computed from the mean
square error of lattice parameters, and $\mathcal{L}_{comp}$ and
$\mathcal{L}_{N_{a}}$ are cross-entropy losses.