FERMILAB-PUB-23-698-CSAID current page.north east -2in -0.5in

Towards a data-driven model of hadronization using normalizing flows

Christian Bierlich1$\spadesuit$, Phil Ilten2$\dagger$, Tony Menzo2,3,4$\star$,
Stephen Mrenna2,5$\maltese$, Manuel Szewc2$\parallel$, Michael K.
Wilkinson2$\perp$, Ahmed Youssef2$\ddagger$, and Jure Zupan2,3,4$\mathsection$

1 Department of Physics, Lund University, Box 118, SE-221 00 Lund, Sweden

2 Department of Physics, University of Cincinnati, Cincinnati, Ohio 45221, USA

3 Berkeley Center for Theoretical Physics, University of California, Berkeley,
CA 94720, USA

4 Theoretical Physics Group, Lawrence Berkeley National Laboratory, Berkeley,
CA 94720, USA

5 Scientific Computing Division, Fermilab, Batavia, Illinois 60510, USA

<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
## Abstract

We introduce a model of hadronization based on invertible neural networks that
faithfully reproduces a simplified version of the Lund string model for meson
hadronization. Additionally, we introduce a new training method for
normalizing flows, termed magic, that improves the agreement between simulated
and experimental distributions of high-level (macroscopic) observables by
adjusting single-emission (microscopic) dynamics. Our results constitute an
important step toward realizing a machine-learning based model of
hadronization that utilizes experimental data during training. Finally, we
demonstrate how a Bayesian extension to this normalizing-flow architecture can
be used to provide analysis of statistical and modeling uncertainties on the
generated observable distributions.

###### Contents

  1. 1 Introduction
  2. 2 Hadronization
  3. 3 Normalizing flows
  4. 4 Fine-tuning microscopic fragmentation kinematics
  5. 5 Uncertainty estimation with Bayesian normalizing flows
  6. 6 Conclusions
  7. A Public code
  8. B Further details on normalizing flows
    1. B.1 Normalizing flows
      1. B.1.1 One-dimensional normalizing flows
      2. B.1.2 Two-dimensional flows
    2. B.2 Bayesian neural networks
    3. B.3 Bayesian normalizing flows
    4. B.4 Interpreting BNF ensembles
    5. B.5 Architecture impact on the BNF uncertainty

##  1 Introduction

Hadronization is one of the least understood ingredients in the simulation of
particle collisions. While the Lund-string [1, 2] and cluster-fragmentation
[3, 4, 5] models give reasonable overall descriptions of hadronization, there
are still significant discrepancies between both the two models [6] and the
models and data [7]. Augmenting these phenomenological models with a data-
driven description of hadronization may help to improve the predictions.

Hadronization models, such as the string and cluster models, serve two
distinct purposes. The first purpose is rooted in direct physics motivation.
We aim to enhance our understanding of QCD behavior beyond the approximations
afforded by lattice quantum chromodynamics (QCD) and perturbative QCD under
specific limits. For this purpose, reliance on models is essential. The second
purpose is to provide a realistic description of final state particles in high
energy collisions. This description allows for the detailed study of detector
responses, as well as realistic modeling of both background and signals for
high momentum transfer processes. This modeling is critical in most high
energy particle physics analyses both for interpreting the results as well as
estimating systematic uncertainties.

In the first scenario, discrepancies between models offer an opportunity to
utilize measurements to deepen our understanding of physics. In this context,
substituting a physics model with a machine learned (ML) model may obscure
these discrepancies or mask the foundational physics phenomena. However, a
carefully designed ML model could also provide insights into the physics model
by supplying a detailed description across all phase space of the physics
model at a granular level. In the second scenario, discrepancies with data do
not provide deeper insights into the hadronization process but rather produce
more ambiguous interpretations of experimental data with larger associated
systematic uncertainties. Here, data-driven models can reduce these
uncertainties while retaining the same physics motivation present in the
original models.

To develop these data-driven models the existing phenomenology can be
augmented, keeping the underlying strings or clusters as the building blocks,
but perturbing their dynamics to accommodate all relevant experimental
observables. This is a problem well suited for ML methods, which can form a
flexible basis for adjusting the underlying model dynamics to match
experimental data. The first attempts at providing an ML description of
simplified hadronizing systems have been carried out using both (MLhad)
conditional sliced Wasserstein autoencoders (cSWAEs) [8] and (HadML)
generative adversarial networks (GANs) [9, 10]. These two architectures have
reproduced key features of the Lund string model in Pythia 8 [11] and the
cluster model in Herwig 7 [12, 13], respectively, but both rely on training
data that is not available at the experimental level.111Here, MLhad [8] uses
the kinematics of first hadron emissions from a string, which is only
available at the generator level. The HadML model uses either the same
information from cluster decays [9] or the full hadron-level kinematic
information for collisions [10], which is not yet available in practice.

At present, there are three main challenges to performing ML training on real
experimental data: (1) develop a procedure to alter microscopic string
dynamics for parton systems produced from existing event generators; (2)
quantify the uncertainties associated with this procedure and propagate them
through detector and material simulations; and (3) identify and measure an
adequately large set of observables sufficiently sensitive to hadronization to
break model degeneracy. Here, we address the first two challenges.

Building on the work of the MLhad model from Ilten:2022jfm,refs. ref. [8], we
analyze a simplified version of the Lund string model, now in the context of
normalizing-flow (NF) ML architectures [14, 15, 16]. The NF architectures
transform a simple underlying probability density into the complex
hadronization probability density via mappings that can be inverted. This
specific feature of NFs is key to the work presented here; hadronization from
one NF model can be reweighted to another model with minimal computational
cost. The NF architectures presented here surpass the previous cSWAE
architecture in both efficiency and physics capabilities, and enables, in the
context of Bayesian NFs (BNFs) [17], a coherent analysis of uncertainties.
Furthermore, it provides a method for determining microscopic dynamics from
macroscopic observables via a novel training approach termed magic.

This paper is organized as follows: in section 2 we briefly review the
simplified Lund string model. This model is used in section 3 to train a NF
architecture for hadronization. In section 4 we use a modified NF to introduce
the magic method for fine-tuning hadronization models, while in section 5 we
use Bayesian NFs to estimate uncertainties. Finally, section 6 contains our
conclusions. Appendices contain details about the public code, appendix A, and
a pedagogical introduction to Bayesian normalizing flows, appendix B.

##  2 Hadronization

Hadronization describes the conversion of a partonic system consisting of
quarks $q$, antiquarks $\bar{q}$, and gluons $g$ into a final state consisting
of hadrons $h$. In what follows, we use the Lund string model of hadronization
[1, 2] to describe the simplest hadronizing partonic system, a
$q_{i}\bar{q}_{i}$ pair of massless quarks with flavor $i$. Specifically, we
consider the quark system in the center of mass frame, with the quark and
antiquark in close proximity and traveling apart with equal and opposite
momenta. In the Lund model, as the separation between the quark and antiquark
increases, the non-Abelian nature of the strong force causes an approximately
uniform string, or flux tube, of color field to form between the quark and the
antiquark, with an approximately uniform energy density $\kappa\approx
1\text{\,Ge\kern-1.00006ptV}/\text{fm}\approx
0.2\text{\,Ge\kern-1.00006ptV}^{2}$. The quark and antiquark are the endpoints
of this string.

The constant force between the quark and antiquark translates into a potential
energy that increases with their separation. As the kinetic energy of the
quark and antiquark at the endpoints of the string is converted into the
potential energy of the string, it can become energetically favorable to
create $q\bar{q}$ pairs out of the vacuum, thereby breaking the string. The
original string breaks into fragments, e.g., a composite hadron $h\equiv
q_{i}\bar{q}_{j}$ and a string fragment with endpoints $q_{j}\bar{q}_{i}$.
Multiple emissions can be implemented sequentially by boosting and rotating
into each hadronizing string fragment’s center-of-mass frame, emitting a
hadron while conserving energy and momentum, then boosting and rotating the
hadron and the new string fragment back to the rest frame of the initial
string.

The kinematics of the emitted hadron is determined through a correlated
sampling of transverse momentum $p_{{\perp}}$ and longitudinal momentum
fraction

$z\equiv(E\pm p_{z})_{\text{hadron}}/(E\pm p_{z})_{\text{string}}\,,$

where $E$ and $p_{{z}}$ are the energy and longitudinal momentum of the hadron
or string, as labeled, in the center-of-mass frame of the string fragment from
which the hadron is emitted, and the parton is moving in the $\pm\hat{z}$
direction in the same frame. After each hadron emission, the kinematics of the
string are updated. Although the $q_{j}\bar{q}_{j}$ pair has no net transverse
momentum, both the $q_{j}$ quark and $\bar{q}_{j}$ antiquark carry transverse
momentum $p_{{\perp}}\equiv\sqrt{p_{x}^{2}+p_{y}^{2}}$, where $p_{x}$ and
$p_{y}$ are perpendicular to each other and sampled from a Gaussian
probability distribution

$\mathcal{P}(p_{x},p_{y};\sigma_{p_{{\perp}}})=\frac{1}{2\pi\sigma_{p_{{\perp}}}^{2}}\exp\left(-\frac{p_{x}^{2}+p_{y}^{2}}{2\sigma_{p_{{\perp}}}^{2}}\right)\,,$
(1)

where the width parameter $\sigma_{p_{{\perp}}}$ is obtained from fits to
data.222Within Pythia 8, $\sigma_{p_{{\perp}}}$ is set with the parameter name
and default value of StringPT:sigma = 0.335.

The probability for a hadron to be emitted with longitudinal lightcone
momentum fraction $z$ is given by the Lund symmetric fragmentation function

$f(z)\propto\frac{\left({1-z}\right)^{a}}{z}\exp\left(-\frac{bm_{{\perp}}^{2}}{z}\right)\,,$
(2)

where $m_{{\perp}}^{2}\equiv m^{2}+p_{{\perp}}^{2}$ is the square of the
transverse mass, $m$ is the hadron mass, and $a$ and $b$ are fixed parameters
determined by fits to data.333The default parameter names and values as
implemented in Pythia 8 are StringZ:aLund = 0.68 and StringZ:bLund = 0.98, for
$a$ and $b$, respectively. Each iteration of causally disconnected string
fragmentations consists of: randomly selecting one string end; assigning
probabilistically a quark flavor to be pair produced during string breaking;
generating the transverse momentum of this pair; generating the light-cone
momentum fraction of the new hadron; and finally computing the longitudinal
momentum of the new hadron, by conserving the total energy and momentum of the
system. Iterative fragmentation continues until the energy of the string
system crosses a threshold. The remaining string piece is then combined into a
final pair of hadrons such that the energy of the initial two-parton system is
converted entirely into the emitted hadrons.

Working within the Lund string model, the phenomenology of hadronization is
largely determined by the probabilities with which different hadron species
are produced, i.e., (i) the forms of the probability distributions for the
hadron momenta traditionally determined by eqs. 2 and 1, (ii) the method of
determining the color singlet systems, and (iii) the process of flavor
selection. Here, we set aside (ii) and (iii) to focus on (i), the kinematics
of string fragmentations, for which we want to ultimately develop a data-
driven determination of the probability distributions.

##  3 Normalizing flows

Normalizing flows (NFs) are generative ML models that can produce high-quality
continuous approximations of probability distributions from a limited set of
data samples [14, 15, 16]. They accomplish this by concatenating a series of
$N$ independent, bijective transformations,
$F(\boldsymbol{z})=f_{N}(f_{N-1}(\ldots f_{2}(f_{1}(\boldsymbol{z}))\ldots))$,
which then map a latent space probability distribution
${\mathcal{P}_{Z}}(\boldsymbol{z})$ to a target distribution
${\mathcal{P}_{X}}(\boldsymbol{x})$. The latent space is typically chosen such
that it can be easily sampled. The form of the bijective functions $f_{i}$ is
adjusted by modifying the model parameters $\boldsymbol{\theta}$ and any
external parameters, including any provided conditional labels
$\boldsymbol{c}$.

Since each $f_{i}$ is a continuous function, the composite function $F$ is
also continuous. This allows for density estimation over the full phase space
including regions sparsely populated by the training data. In section 4 we use
this feature to introduce a method for fine-tuning the form of the microscopic
fragmentation function using measured observable quantities. Furthermore, in
the MLhad NF architecture we use Bayesian NFs (BNFs) [17], in which the model
parameters $\boldsymbol{\theta}$ themselves are random variables. They are
taken to be normally distributed, with average values and variances learned
from training data and which encode data uncertainties, as described in
section 5. Further details on both NFs and BNFs are given in appendix B.

The MLhad NF architecture is able to reproduce pseudo-data generated using a
simplified version of the Pythia 8 Lund string hadronization model. This
pseudo-data is produced using the same simplified model as in
Ilten:2022jfm,refs. ref. [8], in which only light-quark flavors are allowed as
endpoints, and isospin symmetry is required, i.e., only neutral and charged
pions at a single mass are generated. The MLhad NF is trained on a dataset of
$N$ hadron emissions from a string with energy
$E_{\text{ref}}=200\text{\,Ge\kern-1.00006ptV}$. That is, the training dataset
consists of $N$ two-dimensional arrays of first hadron emission kinematics
$\boldsymbol{x}_{n}=\\{p_{z,n},p_{{\perp},n}\\}$, where
$n\in\left\\{1,\ldots,N\right\\}$ and $p_{{z}}$ and $p_{{\perp}}$ are,
respectively, the longitudinal and transverse components of the emitted
hadron’s momentum in the-center-of-mass frame of the string. To generate
hadron kinematics for strings with energies other than $E_{\text{ref}}$, we
use the rescaling property of the Lund string fragmentation function to render
$p_{{z}}$ independent of the string energy, transforming the generated value
of $p_{{z}}$ according to $p_{z}\to p_{z}E_{\text{ref}}/E$, where $E$ is the
energy of the quark in the initial string’s center-of-mass frame [8].

Unlike in Ilten:2022jfm,refs. ref. [8], we train the MLhad NF on a dataset
containing different transverse masses, $m_{{\perp}}$. For this, we construct
labeled training datasets $\\{x_{n},c_{n}\\}_{n=1}^{N}$, where

$c_{n}\equiv\frac{m_{{\perp},\max}-m_{{\perp},n}}{m_{{\perp},\max}-m_{{\perp},\min}}\,,$
(3)

$m_{{\perp},\min}=m_{\pi^{\pm}}\approx 0.140\text{\,Ge\kern-1.00006ptV}$ and
$m_{{\perp},\max}=1.3\text{\,Ge\kern-1.00006ptV}$ are, respectively, the
minimal and maximal values of $m_{{\perp}}$ used in training. The maximum is
chosen such that none of the hadronization chains considered will produce an
$m_{{\perp}}$ above this value. That is, the conditional labels $c_{n}$ are
functions of hadronic transverse mass $m_{{\perp}}$ such that $c_{n}\in[0,1]$,
where the boundaries correspond to the minimum and maximum $m_{{\perp}}$.
Here, $m_{{\perp}}$ is used rather than mass to ensure the independence of the
$z$ and $p_{{\perp}}$ probability distributions for a given $m_{{\perp}}$
value, see eq. 2. The training dataset is split into 15 different conditional
labels, where each label corresponds to a different fixed $m_{{\perp}}$. For
each conditional label $5\times 10^{5}$ first hadron emissions are used, for a
total of $N=7.5\times 10^{6}$ emissions in the full training dataset.

Figure 1: A comparison between the (histograms) Pythia 8 and (solid lines)
MLhad NF generated single emission $p_{{z}}$ distributions produced at four
different fixed values of $m_{{\perp}}$ which were not used in the training of
the model.

Figure 1 shows a comparison between Pythia 8 generated kinematic distributions
and the learned MLhad NF kinematic distributions for different values of the
transverse mass (we have set the NF model parameters $\boldsymbol{\theta}$ to
their fixed average values). We observe that the NF model can fully reproduce
the $p_{{z}}$ of the hadronizing $q\bar{q}$ system of the Lund string model
for arbitrary hadron mass. The result of a full hadronization chain, where
hadrons are sequentially emitted from the string fragments, is shown in fig.
2; the string fragmentation terminates at
$E_{\text{cutoff}}=25\text{\,Ge\kern-1.00006ptV}$ in this case in order to
avoid the final combination step, which requires a separate treatment. We
observe excellent agreement between the hadron multiplicities produced by
Pythia 8 and the MLhad NF.

Figure 2: (left) Comparison of hadron multiplicity distributions generated
with (orange) Pythia 8 and the (blue) MLhad NF, constructed from
hadronizations of $10^{5}$ strings, all with an initial energy of
$E_{\text{string}}=100\text{\,Ge\kern-1.00006ptV}$ and an energy cutoff of
$E_{\text{cutoff}}=25\text{\,Ge\kern-1.00006ptV}$. (right) Scaling of average
hadron multiplicity $\langle N_{h}\rangle$ is given as a function of the
starting string energy. Each marker represents an averaging of hadron
multiplicity over the hadronization of $5\times 10^{3}$ strings. The dotted
lines show the average multiplicity $\langle N_{h}\rangle$ and the bands the
corresponding $1\sigma$ range.

Next, we show how the MLhad NF model trained on single hadron emission pseudo-
data can be used to adjust the microscopic model of single hadron emission
kinematics, so that it reproduces experimental data that has no direct single
hadron emission measurements.

##  4 Fine-tuning microscopic fragmentation kinematics

In this section we introduce a training method for fine tuning NF-based models
of hadronization. The method, termed microscopic alterations generated from
infrared collections444Here, infrared collections refer to any hadronization-
sensitive high-level observable distribution that can be obtained from
experiment. (magic), allows for the fine tuning of microscopic dynamics to
describe a set of macroscopic observables. Practically, the microscopic
dynamics are produced from an underlying phenomenological model, while the
macroscopic observables are from experiment.

magic is a natural extension of the traditional tuning techniques, such as
manual tuning [18] or automated regression techniques [19, 20, 21, 22].
Crucially, while approaches such as deep neutral networks using classification
for tuning and reweighting (DCTR) [22] do parameter reweighting and tuning,
they do not directly modify the underlying parametric Lund model used for
training. magic is able to increase the flexibility of the model beyond the
parametric form, eqs. 1 and 2, while remaining physically meaningful by
keeping the emission-by-emission paradigm described in section 2.

This method works by adding data-driven perturbations to an analytic solution,
the Lund symmetric fragmentation function of eq. 2 in the case of
hadronization, by augmenting it with an over-parameterized function such as an
NF that can be modified arbitrarily to accommodate data. The Lund symmetric
fragmentation function already provides a good description of experimental
data; we seek to learn data-driven perturbations to obtain even better
agreement with experiment.

As a toy example we take a simplified one-dimensional NF model, consisting of
a weighted mixture of Gaussian distributions, trained on the $z$ component of
the momentum of first-emission hadrons, as described in section 3. magic
consists of two training phases: in the first phase, the NF is trained on
simulated kinematics as described in section 3; in the second phase, the NF is
modified to match the experimental data. The initial NF model, or base model,
provides high fidelity sampling of single-hadron emission kinematics
$\boldsymbol{x}=p_{{z}}$. Here we omit $p_{{\perp}}$, in contrast to section
3, for simplicity of the model. From these $\boldsymbol{x}$ kinematics, one
can obtain predictions for measurable, hadronization-sensitive observables
$\boldsymbol{y}$, e.g., hadron multiplicity. In the second phase of training,
the base model is fine-tuned by reweighting the dataset of $\boldsymbol{y}$
values generated by the base model to statistically match the experimentally
observable dataset. This second phase explicitly relies on the ability to
reweight distributions.

The training data for the second phase of magic consists of three components:
(i) the hadronization-chain-level kinematics $\boldsymbol{x}$, i.e., the
hadron kinematics $p_{{z}}$ from simulated emissions produced by the base
model; (ii) the desired measurable observables from simulated hadronization
chains produced by the base model $\boldsymbol{y}_{\text{sim}}$, e.g., the
hadron multiplicity $N_{h}$ predicted by the base model; and (iii) values of
the same observables measured experimentally $\boldsymbol{y}_{\text{exp}}$. As
a proof of principle, we use just a single observable, the total number of
hadrons for a single hadronization chain, i.e., the hadron multiplicity such
that $\boldsymbol{y}$ is $N_{h}$.

An example of the training data, consisting of $N$ hadronization chains, is
therefore555While not denoted explicitly in eq. 4, each array
$\boldsymbol{x}_{i}$ is zero-padded to a fixed length of size
$\max(\boldsymbol{y}_{\text{sim}})=\max(N_{h,n})$.

$\boldsymbol{x}=\begin{pmatrix}\boldsymbol{x}_{1}=\big{\\{}p_{{z,h_{1}}},p_{{z,h_{2}}},p_{{z,h_{3}}}\big{\\}}_{1}\\\
\boldsymbol{x}_{2}=\big{\\{}p_{{z,h_{1}}},p_{{z,h_{2}}},p_{{z,h_{3}}},p_{{z,h_{4}}}\big{\\}}_{2}\\\
\vdots\\\
\boldsymbol{x}_{N}=\big{\\{}p_{{z,h_{1}}},p_{{z,h_{2}}}\big{\\}}_{N}\end{pmatrix},\quad\boldsymbol{y}_{\text{sim}}=\begin{pmatrix}\boldsymbol{y}_{1}=N_{h,1}=3\\\
\boldsymbol{y}_{2}=N_{h,2}=4\\\ \vdots\\\
\boldsymbol{y}_{N}=N_{h,N}=2\end{pmatrix}\,.$ (4)

The fine tuning modifies the hadronization model, i.e., the distribution
governing $p_{{z}}$ for each emission, to minimize the difference between the
two ensembles $\boldsymbol{y}_{\text{sim}}$ and $\boldsymbol{y}_{\text{exp}}$.
We do not match specific measured hadronization chains to a given
hadronization history but instead compare the two ensembles at the statistical
level.

magic does not require regenerating hadron emissions for each perturbation of
the NF, and instead reweights the hadronization chains, which is
computationally advantageous. We make use of the fact that NFs give explicit
access to the model likelihood and reweight events that were originally
sampled from the base model to events sampled from the updated model. Each
hadronization-chain weight can be computed in terms of the likelihood ratio
between the updated, or perturbed, model likelihood
${\mathcal{P}_{X}}(\boldsymbol{x}_{n},\boldsymbol{\theta}_{P})$, and the base
model likelihood
${\mathcal{P}_{X}}(\boldsymbol{x}_{n},\boldsymbol{\theta}_{B})$. Written in
terms of single emissions, the likelihood
${\mathcal{P}_{X}}(\boldsymbol{x}_{n},\boldsymbol{\theta})$ can be factorized
as,

${\mathcal{P}_{X}}(\boldsymbol{x}_{n},\boldsymbol{\theta})=\prod_{i=1}^{N_{h,n}}{\mathcal{P}_{X}}(\boldsymbol{x}_{n,i},\boldsymbol{\theta})\,,$
(5)

where $N_{h,n}$ is the number of hadrons in hadronization chain $n$, and
$\boldsymbol{x}_{n,i}$ is emission $i$ of chain $n$.

We introduce a hadronization-chain weight array $\boldsymbol{w}$, where each
weight is computed as the product of the likelihood ratios for all emissions
in a chain

$\boldsymbol{w}=\begin{pmatrix}w_{1}\\\ w_{2}\\\ \vdots\\\
w_{N}\end{pmatrix},\text{ where
}w_{n}=\prod_{i=1}^{N_{h,n}}\frac{{\mathcal{P}_{X}}(\boldsymbol{x}_{n,i},\boldsymbol{\theta}_{P})}{{\mathcal{P}_{X}}(\boldsymbol{x}_{n,i},\boldsymbol{\theta}_{B})}\,.$
(6)

Explicitly for this example, $\boldsymbol{x}_{n,i}$ is just the $p_{{z}}$ of
hadron $i$ from hadronization chain $n$. The reduction in training time
associated with this use of hadronization-chain weights is crucial for the
technical feasibility of the magic approach to fine-tuning.

The learning objective of the fine-tuning phase is to minimize the statistical
distance between $\boldsymbol{y}_{\text{sim}}$, weighted by $\boldsymbol{w}$,
and the target distribution $\boldsymbol{y}_{\text{exp}}$. In our toy example,
we use the Wasserstein distance [23, 24, 25, 26], or Earth mover’s distance
(EMD), as a measure of the similarity between the two samples666In many cases,
when training from real experimental distributions, one may only have access
to binned datasets. The magic paradigm may equivalently be used in these
scenarios by simply utilizing a binned statistical distance such as
$\chi^{2}$. and define the loss function as

${\mathcal{L}_{\text{EMD}}}(\boldsymbol{y}_{\text{sim}},\boldsymbol{y}_{\text{exp}})=\sum_{n=1}^{N}\sum_{m=1}^{M}f^{*}_{n,m}d_{n,m}\,,$
(7)

where the elements of the flow matrix $f_{n,m}$ encode the fractional amount
of weight to be transferred between event $\boldsymbol{y}_{\text{sim},n}$ and
$\boldsymbol{y}_{\text{exp},m}$ and
$d_{n,m}=||\boldsymbol{y}_{\text{sim},n}-\boldsymbol{y}_{\text{exp},m}||_{2}$
is the distance between these two hadronization chains. Here, $M$ is the
number of hadronization chains observed in the experimental dataset.

Once the loss has been computed, back-propagation algorithms update the NF
parameters $\boldsymbol{\theta}$ using PyTorch’s automatic differentiation
engine autograd. The autograd engine utilizes differential programming
paradigms with dynamic computational graphs to yield the gradients of the loss
function with respect to all parameters
$\nabla_{\kern-2.10002pt\boldsymbol{\theta}}\,{\mathcal{L}_{\text{EMD}}}$ by
tracking the impact of the hadronization-chain weights
$\boldsymbol{w}(\boldsymbol{\theta})$ on the loss. We can then find a model
likelihood that produces the targeted observable distribution because updating
$\boldsymbol{\theta}$ corresponds to updating every
${\mathcal{P}_{X}}(\boldsymbol{x}_{n,i},\boldsymbol{\theta})$. The only
dynamical object in the fine-tuning phase of magic is the hadronization-chain
weight array $\boldsymbol{w}$; the base model, $\boldsymbol{x}$,
$\boldsymbol{y}_{\text{sim}}$, and $\boldsymbol{y}_{\text{exp}}$ all remain
fixed.

In our toy example, we use a one-dimensional NF base model, see section B.1,
trained on $N=5\times 10^{5}$ Pythia 8 generated hadronization chains with the
Lund string parameter $a$ set to $0.68$. Each of the transverse momentum
components of the emitted hadrons is sampled from a Gaussian distribution, and
the correlation between $p_{{z}}$ and $p_{{\perp}}$ is neglected for
simplicity. We create a pseudo-experimental hadron multiplicity observable
$\boldsymbol{y}_{\text{exp}}$ with $M=5\times 10^{5}$ samplings of a second NF
trained on Pythia 8 hadronization chains, produced with $a=1.5$.777We do not
use Pythia 8 directly to generate the targeted pseudo-experimental
multiplicity training dataset due to the included correlations between
$p_{{z}}$ and $p_{{\perp}}$, unlike the simplified base model.

We then perform 15 independent magic fine-tunings888Training using magic is
computationally inexpensive, for example, the training for the presented toy
example can be performed on a modern laptop CPU with training times of
$\mathcal{O}(1~{}\text{hour})$ to achieve similar accuracy to the results
shown in fig. 3. of the base model, where each fine-tuning is trained over
three epochs in batches of $10^{4}$ samples with the learning rate fixed to
$\delta=2.5\times 10^{-2}$ in the first epoch, and reduced to $\delta=1\times
10^{-2}$ for the remaining two epochs. The results of the independent
trainings can be seen in the left panel in fig. 3; we see that the fine-tuned
models all successfully learn the softer hadron emission spectrum used to
generate the pseudo-experimental values of hadron multiplicity. In the right
of fig. 3, we see that the corresponding hadron multiplicity distributions
also agree.

Figure 3: (left) Comparison between the single emission kinematic distribution
$p_{{z}}$ of the (blue dashed) base model, the (green dashed-dotted)
distribution used to generate $\boldsymbol{y}_{\text{exp}}$, and (red) an
ensemble of solutions learned using magic. (right) Comparison of
hadronization-chain level hadron multiplicity between the (blue) base, (red)
one of the fine tuned, and the (green) target distributions where each
histogram is constructed from $5\times 10^{5}$ hadronization chains. The mean
of each histogram is shown as a vertical dashed line.

We note some observations regarding the application of magic. First, the fine-
tuned models of $p_{{z}}$ shown in fig. 3 do not exactly match the target, and
are not unique. This is expected, and is ultimately a consequence of finite
training data and time. In multi-dimensional magic fine-tunings with just a
single observable, we found that qualitatively different solutions are
obtained, all of which are able to reproduce the observable distributions
matching those of the target. This degeneracy is presumably broken once
additional sufficiently-orthogonal observables are included in
$\boldsymbol{y}$.

The magic fine-tuned models could be included within existing event generation
pipelines, either directly as a kinematic generator or as a reweighter that
utilizes the learned likelihood ratio between the base, e.g., the default
Pythia 8, and fine-tuned models, to provide event weights similar to those of
Bierlich:2023fmh,refs. ref. [27]. We leave a full exploration and
incorporation of the magic method within Pythia 8 for future work.

##  5 Uncertainty estimation with Bayesian normalizing flows

When generating binned distributions using an ML model, there are typically
two sources of uncertainty to consider: an uncertainty $\sigma_{\text{gen}}$
due to the limited statistics of the generated data, as well as the systematic
uncertainties due to the ML model. The systematic uncertainties can be further
separated into $\sigma_{\text{train}}$ and $\sigma_{\text{data}}$, where
$\sigma_{\text{train}}$ captures the uncertainties due to the size of the
training dataset and the ML architecture, while $\sigma_{\text{data}}$ are the
additional systematic uncertainties inherent to the training data. For
example, when training on experimental data, the statistical uncertainty of
that data contributes to $\sigma_{\text{train}}$, while the systematic
uncertainties, such as the detector resolution, contribute to
$\sigma_{\text{data}}$. While $\sigma_{\text{gen}}$ is just the statistical
uncertainty of the generated sample, $\sigma_{\text{train}}$ and
$\sigma_{\text{data}}$ are typically much more difficult to quantify. Here, we
propose methods to evaluate both $\sigma_{\text{train}}$ and
$\sigma_{\text{data}}$.

Using a Bayesian NF, $\sigma_{\text{gen}}$ and $\sigma_{\text{train}}$ can be
evaluated simultaneously, see section B.4. In this framework, the posterior of
the network parameters $\boldsymbol{\theta}$ should capture how the training
uncertainties result in different neural network choices that are all
compatible with the training data. We first demonstrate that BNFs can capture
both $\sigma_{\text{gen}}$ and $\sigma_{\text{train}}$ by training a BNF on a
dataset of $N_{\text{train}}=10^{5}$ two-dimensional vectors of first-emission
hadron kinematics given by $\\{p_{{z}},p_{{\perp}}\\}$. Then, $M=5\times
10^{4}$ sets of randomly-selected network parameters $\boldsymbol{\theta}_{m}$
are sampled from the BNF posterior. For each $\boldsymbol{\theta}_{m}$, we
generate an independent dataset of $N_{\text{gen}}=10^{5}$ first emissions.

We compute the average number of emissions $\langle N_{\text{bin}}\rangle$ in
each bin of hadron $p_{{z}}$, and its variance $\sigma_{\text{bin}}$ across
all $M$ BNFs:

$\displaystyle\langle N_{\text{bin}}\rangle$
$\displaystyle\equiv\frac{1}{M}\sum_{m=1}^{M}\langle
N_{\text{bin}}\rangle_{\boldsymbol{\theta}_{m}}\,,$ (8)
$\displaystyle\sigma^{2}_{\text{bin}}$ $\displaystyle\equiv\langle
N_{\text{bin}}\rangle+\frac{1}{M}\sum_{m=1}^{M}\left(\langle
N_{\text{bin}}\rangle_{\boldsymbol{\theta}_{m}}-\langle
N_{\text{bin}}\rangle\right)^{2},$

where $\langle N_{\text{bin}}\rangle_{\boldsymbol{\theta}_{m}}$ is the
expected number of first emissions that fall in this particular $p_{{z}}$ bin,
estimated from the sample of $N_{\text{gen}}$ emissions generated using the NF
with parameters $\boldsymbol{\theta}_{m}$. The left panel in fig. 4 compares
the $\langle N_{\text{bin}}\rangle$ from the learned model with the training
dataset. We observe that for most bins, the model and the training dataset are
consistent within uncertainty, although this degrades for sparsely populated
bins where the model has not been trained with sufficient data.

The right panel in fig. 4 compares $\sigma_{\text{gen}}=\sqrt{\langle
N_{\text{bin}}\rangle}$ with the total uncertainty $\sigma_{\text{bin}}$ in a
given bin. The total uncertainty includes both the generated uncertainty as
well as the training uncertainty,
$\sigma_{\text{bin}}^{2}=\sigma_{\text{gen}}^{2}+\sigma_{\text{train}}^{2}$.
For illustrative purposes, we compare in fig. 4 (right) the relative
uncertainties in each bin, $\sigma_{\text{bin}}/\langle N_{\text{bin}}\rangle$
and $\sigma_{\text{gen}}/\langle N_{\text{bin}}\rangle$, rather than the
absolute ones. The BNF estimate of the total uncertainty $\sigma_{\text{bin}}$
is always larger than $\sigma_{\text{gen}}$ because it contains also
$\sigma_{\text{train}}$, the additional model uncertainty due to finite
training statistics and model choice. We observe that in the sparsely
populated bins the relative uncertainty is large, with $\sigma_{\text{bin}}$
significantly larger than $\sigma_{\text{gen}}$, signaling that the BNF model
is a poor predictor in this kinematic regime, primarily due to the lack of
training data in the corresponding bins. Conversely, in densely populated bins
the model is an accurate approximation of the data and the estimated BNF
uncertainty $\sigma_{\text{bin}}$ approaches $\sigma_{\text{gen}}$. This
implies that either the posterior is relatively certain of the underlying
probability distribution, or that the model has been over-fit. To avoid over-
fitting, we regularized the training to stop when the model performance on a
validation dataset does not improve after $50$ epochs.

Figure 4: Generated and training uncertainties captured by the BNF. (left)
Comparison of the $p_{{z}}$ distribution generated by (green) Pythia 8 and the
(blue) MLhad BNF. The Pythia 8 distribution was generated with $10^{5}$ single
emissions, and the uncertainty is given by $\sqrt{N_{\text{bin}}}$. The BNF
distribution is generated from an ensemble average over $5\times 10^{4}$ BNFs,
each consisting of $10^{5}$ single emissions, with means and uncertainties
calculated using eq. 8. (right) Comparison of the (blue) BNF estimate of the
relative uncertainty, $\sigma_{\text{bin}}/\langle N_{\text{bin}}\rangle$, and
the (red) relative generated uncertainty, $\sigma_{\text{gen}}/\langle
N_{\text{bin}}\rangle$, for each bin.

A more detailed study of the uncertainties as a function of the number of
generated events $N_{\text{gen}}$ is presented in fig. 5 for two
representative $p_{z}$ bins from fig. 4: the densely populated bin,
$p_{{z}}\in[6.67,8.33)\text{\,Ge\kern-1.00006ptV}$ and the scarcely populated
bin, $p_{{z}}\in[48.33,50)\text{\,Ge\kern-1.00006ptV}$. The expected values
$\sigma_{i}$ for each $N_{\text{gen}}$ in fig. 5 were obtained by first
computing $\langle N_{\text{bin}}\rangle_{\boldsymbol{\theta}_{m}}$ by
approximating the integral in eq. 36 of section B.4 with a very large number
of emissions, $10^{7}$, and then using $\langle
N_{\text{bin}}\rangle_{\boldsymbol{\theta}_{m}}$ to obtain $\sigma_{i}$ as in
eq. 8. This strategy ensures that the integration uncertainty is negligible
and independent of $N_{\text{gen}}$. For the largest bin $\sigma_{\text{gen}}$
dominates even when $N_{\text{gen}}>N_{\text{train}}$ until for sufficiently
large $N_{\text{gen}}$, $\sigma_{\text{train}}$ becomes the leading
uncertainty on the generated dataset. This is expected, since the model is
acting as a fit and, if sufficiently constrained, need not return the
statistical uncertainty of the bin counts in the training data. For densely
populated bins, $\sigma_{\text{train}}$ may thus be smaller than just the
statistical uncertainty in a particular bin of the training dataset. In
contrast, for the smallest bin, $\sigma_{\text{train}}$ is the dominant
uncertainty even when $N_{\text{gen}}\lesssim N_{\text{train}}$, indicating
that the model is a poor predictor with the posterior yielding a larger
variance for $\boldsymbol{\theta}$.

Figure 5: Study of the component uncertainties obtained with the BNF: the
generated uncertainty $\sigma_{\text{gen}}$, training uncertainty
$\sigma_{\text{train}}$, and the total uncertainty $\sigma_{\text{bin}}$
evaluated by the BNF for the (left) most populated and (right) least populated
bins of the $p_{{z}}$ distribution in fig. 4.

The NF architecture of section 3 can also be used to efficiently assess the
systematic uncertainties of experimental data, $\sigma_{\text{data}}$, as long
as the NF is conditioned on these systematic uncertainties during training.
Here, we consider a toy example where we treat the uncertainties on the
parameters of the hadronization model as systematic uncertainties.
Specifically, we introduce an uncertainty on the parameter $b$ of the Pythia 8
Lund string model, see eq. 2. We train the BNF on pseudo-data produced by
Pythia 8, where the value of $b$ is set to three discrete values encoded by
the conditional labels: the base value $b_{B}=0.98$ and two perturbed values
$b_{P}\in\\{0.8,1.4\\}$, corresponding to a systematic uncertainty envelope on
$b$.

After training, a large dataset of hadronization chains can be generated with
this BNF using the conditional label $b_{P}$. However, since the probability
for each hadron emission,
$\boldsymbol{x}_{n}\sim{\mathcal{P}_{X}}(\boldsymbol{x}_{n},b)$, is now a
known function of $b$ from the conditioned training, we can also calculate and
track the probability for each perturbed conditional label $b_{P}$ per hadron
emission. The weight for a hadronization chain to be produced with $b_{P}$
rather than $b_{B}$ is

$w\approx\prod_{n=1}^{N_{h}}\frac{{\mathcal{P}_{X}}(\boldsymbol{x}_{n},\boldsymbol{\theta}^{*},b_{P})}{{\mathcal{P}_{X}}(\boldsymbol{x}_{n},\boldsymbol{\theta}^{*},b_{B})}\,,$
(9)

where $N_{h}$ is the number of hadrons emitted in the particular hadronization
chain, i.e., the hadron multiplicity for the hadronization of that particular
string. We can then quickly produce large datasets for the values of $b_{P}$,
by simply reweighting the initial dataset generated with $b_{B}$.

It is important to note that eq. 9 would be an equality, rather than an
approximate relation, had we marginalized over all possible network parameters
$\boldsymbol{\theta}$. However, for the sake of expediency, we use a fixed set
$\boldsymbol{\theta}^{*}$ instead. Although $\boldsymbol{\theta}^{*}$ can be
chosen to be the parameter set that provides a maximum a posteriori (MAP)
value, in this example we simply sample an arbitrary $\boldsymbol{\theta}^{*}$
per emission for the posterior distribution of the BNF.

Figure 6: Distribution of the number of hadrons produced per hadronization
chain for a sample of $5\times 10^{4}$ strings. These hadronization chains are
generated using three BNFs with fixed network weights conditioned on three
distinct datasets: (blue) the base dataset $b_{B}=0.98$, (orange) a perturbed
dataset with $b_{P}=0.8$, and (red) a perturbed dataset with $b_{P}=1.4$. The
distribution from the BNF for the base dataset is then reweighted using eq. 9
to the perturbed values of $b$, (green) $b_{P}=0.8$ and (blue) $b_{P}=1.4$.

In fig. 6 we illustrate the reweighting process by plotting the multiplicity
distribution, the number of hadrons produced per hadronization chain, from a
large sample of hadronization chains produced by our BNF. This distribution is
then reweighted to the two perturbed values of $b$. We see that reweighting
each hadronization chain using (9) leads to multiplicity distributions that
match the Pythia 8 generated ones within statistical uncertainty. Beyond this
toy example, eq. 9 can be used to reweight any sample generated by a BNF to
variations of the systematic uncertainties, assuming that the BNF is initially
trained with conditional labels for each systematic uncertainty to be
considered.

##  6 Conclusions

In this manuscript, we have shown that the normalizing flow architecture is
well suited for modeling the non-perturbative process of hadronization. A key
feature of the normalizing flows is the analytic knowledge of the probability
distribution for individual hadron emissions. The architecture is able to
generate high fidelity first hadron emission kinematic samples, and Bayesian
normalizing flows can be used to provide estimates of the hadronization
uncertainty.

We have also introduced a novel training method, magic, which allows for the
systematic alteration of microscopic fragmentation dynamics such that the
predictions best fit macroscopic observables, e.g., hadron multiplicities. The
magic method avoids the need to produce additional samples each time the
symmetric Lund string fragmentation is modified, and calculating instead
appropriate weights for existing data samples. The use of automatic
differentiation to update the model parameters makes the training numerically
efficient.

We have showcased the potential of both Bayesian normalizing flows and magic
in the context of hadronization using toy examples, but there are a number of
steps that still need to be made before data can be used in training, such as
considering gluon-strings and flavor selection, both of which are part of
ongoing work.

Acknowledgments. AY, JZ, MS, and TM acknowledge support in part by the DOE
grant DE-SC1019775, and the NSF grant OAC-2103889. JZ acknowledges support in
part by the Miller Institute for Basic Research in Science, University of
California Berkeley. SM is supported by the Fermi Research Alliance, LLC under
Contract No. DE-AC02-07CH11359 with the U.S. Department of Energy, Office of
Science, Office of High Energy Physics. CB acknowledges support from the Knut
and Alice Wallenberg foundation, contract number 2017.0036. PI is supported by
NSF grant OAC-2103889 and NSF-PHY-2209769. TM acknowledges support in part by
the U.S. Department of Energy, Office of Science, Office of Workforce
Development for Teachers and Scientists, Office of Science Graduate Student
Research (SCGSR) program. The SCGSR program is administered by the Oak Ridge
Institute for Science and Education for the DOE under contract number DE-
SC0014664.

##  Appendix A Public code

The public code can be found at https://gitlab.com/uchep/mlhad/BNF. The
repository consists of a hierarchical structure with three major components.
The first component contains the implementation of the one- and two-
dimensional NF network, with and without conditioning, the second component
constitutes the integration of these NFs into fragmentation chains, and the
third component is the implementation of magic. Detailed explanations and
examples of each component can be found within the code documentation and
example notebooks. All code is written in Python, developed using v3.11, and
heavily utilizes the PyTorch, developed using v2.1, deep learning library.
Finally, all training datasets were produced using Pythia v8.309.

##  Appendix B Further details on normalizing flows

In this appendix we give further details on the Bayesian normalizing flows
that we use in MLhad. Section B.1 reviews the basics of normalizing flows,
section B.2 contains a brief review of Bayesian neural networks, and Section
B.3 describes Bayesian normalizing flows.

###  B.1 Normalizing flows

Normalizing flows (NFs) [14, 15, 16] are a class of generative ML
architectures that can produce high fidelity continuous approximations of
complex probability distributions using a finite collection of data samples.
This is achieved by learning a composition of $n$ independent bijective
transformations that relate a probability distribution
${\mathcal{P}_{Z}}(\boldsymbol{z})$ on a chosen latent space $Z$ to the target
distribution $p_{X}(\boldsymbol{x})$ on target space $X$.

Figure 7: (top) Schematic of the input and output of the NF architecture. Here
$\boldsymbol{x}_{i}$ and $\tilde{\boldsymbol{z}}$ represent, respectively, the
input and output samples obtained from the network when mapping in the forward
direction and $\boldsymbol{z}$ and $\tilde{\boldsymbol{x}_{i}}$ represent,
respectively, input and output samples of the network obtained when traversing
the network in the backward direction. The backward direction is used for
event generation while the forward direction is used for training. The forward
(backward) direction consists of a series of $n$ successive transformations
$f_{i+1}(\boldsymbol{z}_{i})$ ($f_{i}^{-1}(\boldsymbol{x}_{i})$).

More precisely, given a multivariate random variable
$\boldsymbol{z}\in\mathbb{R}^{d}$ and an invertible map
$f:\mathbb{R}^{d}\to\mathbb{R}^{d}$, the probability distribution for the
random variable $\boldsymbol{x}=f(\boldsymbol{z})$ is given by

${\mathcal{P}_{X,f}}(\boldsymbol{x})={\mathcal{P}_{Z}}(\boldsymbol{z})|\det
J_{f}(\boldsymbol{z})|^{-1}\,,$ (10)

where $J_{f}=\partial f/\partial\boldsymbol{z}$ is the Jacobian of the
differentiable transformation $f$. The full map $F$ produced by the NF
architecture is composed from a sequence of $n$ such transformations
$\boldsymbol{z}\equiv\boldsymbol{z}_{0}\to\boldsymbol{z}_{1}\equiv
f_{1}(\boldsymbol{z}_{0})\to\cdots\to\boldsymbol{x}\equiv\boldsymbol{z}_{n}\equiv
f_{n}(\boldsymbol{z}_{n-1})$, as shown in fig. 7, with the final distribution
given by

${\mathcal{P}_{X}}(\boldsymbol{x})={\mathcal{P}_{Z}}(\boldsymbol{z})\prod_{i=1}^{n}|\det
J_{f_{i}}(\boldsymbol{z}_{i-1})|^{-1}\,.$ (11)

The NF architecture provides a continuous map from the latent space $Z$ to the
target space $X$ and vice versa. In order to train the network to generate
high fidelity mappings of samples from the latent distribution
${\mathcal{P}_{Z}}(\boldsymbol{z})$ to samples of the target distribution
${\mathcal{P}_{X}}(\boldsymbol{x})$ we require a learning objective that will
drive our model distribution
${\mathcal{P}_{X}}(\boldsymbol{x};\boldsymbol{\theta})$ towards
${\mathcal{P}_{X}}(\boldsymbol{x})$. Given training samples
$\boldsymbol{x}_{a}$ of $N$ data points,
$\boldsymbol{x}_{a}=\\{\boldsymbol{x}_{1},\boldsymbol{x}_{2},\ldots,\boldsymbol{x}_{N}\\}$,
with conditional labels
$\boldsymbol{c}_{a}=\\{\boldsymbol{c}_{1},\boldsymbol{c}_{2},\dots,\boldsymbol{c}_{N}\\}$,
we use the minimization of the negative log likelihood as our learning
objective,

$\displaystyle{\mathcal{L}_{\text{NF}}}$
$\displaystyle=\mathbb{E}_{{\mathcal{P}_{X}}(\boldsymbol{x},\boldsymbol{c})}\left[-\log{\mathcal{P}_{X}}(\boldsymbol{x};\boldsymbol{\theta},\boldsymbol{c})\right]=-\sum_{a=1}^{N}\log{\mathcal{P}_{X}}(\boldsymbol{x}_{a};\boldsymbol{\theta},\boldsymbol{c}_{a})$
(12)
$\displaystyle=\sum_{a=1}^{N}\left\\{-\log{\mathcal{P}_{Z}}\left(F^{-1}(\boldsymbol{x}_{a};\boldsymbol{\theta},\boldsymbol{c}_{a})\right)+\log\left|\det
J_{F^{-1}}(\boldsymbol{x}_{a};\boldsymbol{\theta},\boldsymbol{c}_{a})\right|\right\\}\,,$

where $F(\boldsymbol{x};\boldsymbol{\theta},\boldsymbol{c})$ denotes the full
network, parameterized by weights $\boldsymbol{\theta}$ and conditioned on
labels $\boldsymbol{c}$. For a latent space sampled from a two-dimensional
normal distribution, the loss function is given by

${\mathcal{L}_{\text{NF}}}=\sum_{a=1}^{N}\left\\{\frac{1}{2}\bigl{|}\bigl{|}F^{-1}(\boldsymbol{x}_{a};\boldsymbol{\theta},\boldsymbol{c}_{a})\bigr{|}\bigr{|}^{2}_{2}-\log\bigl{|}\det
J_{F}\left(F^{-1}(\boldsymbol{x}_{a};\boldsymbol{\theta},\boldsymbol{c}_{a})\right)\bigr{|}\right\\}\,.$
(13)

where $||\cdots||^{2}_{2}$ denotes the squared $\ell^{2}$-norm. Because each
operation is differentiable, the gradient of $\mathcal{L}_{\text{NF}}$ with
respect to each model parameter $\boldsymbol{\theta}$ may be computed using
standard auto-differentiation software and optimized through stochastic
gradient descent. Intuitively, the loss function in eq. 13 ensures that the
latent variables obtained from mapping the training data samples through the
network, i.e., propagated from $\boldsymbol{x}_{a}\to\boldsymbol{z}_{a}$, are
normally distributed.

For practical applications, the latent distribution
${\mathcal{P}_{Z}}(\boldsymbol{z})$ is chosen such that it can be easily
evaluated and sampled, while the transformations $f_{i}$ are chosen such that
(1) they are expressive enough to sufficiently approximate the transformation
${\mathcal{P}_{Z}}(\boldsymbol{z})\to{\mathcal{P}_{X}}(\boldsymbol{x})$ and
(2) they have computationally inexpensive Jacobians. For example, in the one-
dimensional examples presented in section 4, we use a uniform latent
distribution, $\boldsymbol{z}\sim\mathcal{U}_{[0,1]}$, and a mixture of
Gaussian cumulative distribution functions as the transformations $f_{i}$. For
the two-dimensional examples presented in section 3, we use a two-dimensional
unit Gaussian latent distribution,
$\boldsymbol{z}\sim\mathcal{N}(\vec{0},\mathbb{I}_{2\times 2})$, and real-
valued non-volume preserving (real NVP) transformations, as implemented in the
FrEIA software library [28], for $f_{i}$. In the following subsections, we
provide additional details regarding the architectures used in the one and
two-dimensional models presented in the main text.

####  B.1.1 One-dimensional normalizing flows

For one-dimensional distributions and a single map
$f:\mathbb{R}^{1}\to\mathbb{R}^{1}$, the transformation formula of eq. 10 can
be rewritten as

$\log{\mathcal{P}_{X,f}}=\log{\mathcal{P}_{Z}}(f(x))+\log\left|\frac{\text{d}f(x)}{\text{d}x}\right|\,.$
(14)

In one-dimension we can utilize a cumulative distribution function (CDF) as
the invertible transformation $f$. Because CDFs are continuous, non-decreasing
functions, they are guaranteed to have a unique inverse. Additionally, because
CDFs satisfy

$\lim_{x\to-\infty}\text{CDF}(x)=0\text{ and
}\lim_{x\to\infty}\text{CDF}(x)=1\,,$ (15)

a function that consists of a linear sum of CDFs is a CDF itself, i.e.,

$\lim_{x\to\infty}\left(\sum_{i}w_{i}\text{CDF}_{i}(x)\right)=\sum_{i}w_{i}\left(\lim_{x\to\infty}\text{CDF}_{i}(x)\right)=\sum_{i}w_{i}=1\,,$
(16)

as long as the weights $\boldsymbol{w}$ are normalized such that the right
most equality is true and

$\lim_{x\to-\infty}\left(\sum_{i}w_{i}\text{CDF}_{i}(x)\right)=\sum_{i}w_{i}\left(\lim_{x\to-\infty}\text{CDF}_{i}(x)\right)=\sum_{i}w_{i}\times
0=0\,.$ (17)

In the main text we use a weighted linear mixture of $K$ Gaussian CDFs,
$\Phi(x;\mu_{i},\sigma_{i})$, as the invertible transformation where the
weights $w_{i}$, means $\mu_{i}$, and standard deviations $\sigma_{i}$ of each
Gaussian component are tunable parameters learned by the network. This setup
is commonly referred to as a Gaussian mixture model (GMM). The transformation
and its derivative can be written explicitly as

$\displaystyle f(x)$
$\displaystyle=\sum_{i}^{K}w_{i}\Phi(x;\mu_{i},\sigma_{i})=\frac{1}{2}\sum_{i}^{K}w_{i}\left[1+\text{erf}\left(\frac{x-\mu_{i}}{\sqrt{2}\sigma_{i}}\right)\right]\,,$
(18) $\displaystyle\frac{\text{d}f(x)}{\text{d}x}$
$\displaystyle=\sum_{i}^{K}w_{i}\mathcal{N}(x;\mu_{i},\sigma_{i}),$ (19)

where we have used the fact that
$\text{d}\Phi(x;\mu_{i},\sigma_{i})/\text{d}x=\mathcal{N}(x;\mu,\sigma)$.

Once a latent distribution ${\mathcal{P}_{Z}}$ is specified, the network can
be trained by maximizing eq. 14). After training, samples can be obtained in
the inverse direction via inverse transform sampling. For $n$ successive
transformations eq. 14, is modified to

$\log{\mathcal{P}_{X}}(x)=\log{\mathcal{P}_{Z}}(F(x))+\sum_{i=1}^{n}\log\left|\frac{\text{d}f_{i}(x_{i-1})}{\text{d}x}\right|\,.$
(20)

To increase the flexibility of the network, we insert a total of $n-1$
intermediate non-linear functions $y_{i}$ between each transformation $f_{i}$
and $f_{i+1}$ such that the full transformation $F$ is given by
$F=f_{n}(y_{n-1}(f_{n-1}(\cdots(y_{1}(f_{1}(x)))\cdots)))$ and the sum in eq.
20 runs from $i=1,\ldots,2n-1$.

Specifically, we use a logit transformation defined as

$y(x)=\text{logit}\left(\frac{\alpha}{2}+(1-\alpha)x\right)\text{, where
logit}(x)=\log\frac{1}{1-x}\,,$ (21)

with the derivative

$\frac{\text{d}y(x)}{\text{d}x}=\frac{1-\alpha}{x(1-x)}\,.$ (22)

Here, $\alpha$ is a hyperparameter of the network that is set to $0.01$. The
full architecture used in section 4 utilizes a uniformly distributed latent
distribution ${\mathcal{P}_{Z}}\sim\mathcal{U}_{[0,1]}$ and $n=5$
transformations $f_{i}$, where each transformation contains $K=500$ Gaussian
components. The learnable parameters are the Gaussian means and variances and
the weights of the components, which are constrained to sum to unity per
transformation.

####  B.1.2 Two-dimensional flows

In section 3 we use two-dimensional real NVP transformations for $f_{i}$. Real
NVP transformations consist of modular blocks containing two affine coupling
layers. That is, given an input $\boldsymbol{z}_{i-1}$, the coupling block
splits the input into two channels999Because we will deal with two-dimensional
random variables, in our case both $\boldsymbol{z}_{1}$, $\boldsymbol{z}_{2}$
are one-dimensional. Which component of $\boldsymbol{z}$ is assigned as
$\boldsymbol{z}_{1}$ or $\boldsymbol{z}_{2}$ is randomly chosen for each real
NVP but is kept consistent over the complete dataset and stored for inference.
$\boldsymbol{z}_{i-1}=\\{\boldsymbol{z}_{i-1,1},\boldsymbol{z}_{i-1,2}\\}$ and
applies a sequential affine transformation to each channel as follows

$\displaystyle\boldsymbol{z}_{i,1}=\boldsymbol{z}_{i-1,1}\odot\exp\bigl{(}s_{i,1}\left(\boldsymbol{z}_{i-1,2}\right)\bigr{)}+t_{i,1}(\boldsymbol{z}_{i-1,2})\,,$
(23)
$\displaystyle\boldsymbol{z}_{i,2}=\boldsymbol{z}_{i-1,2}\odot\exp\bigl{(}s_{i,2}\left(\boldsymbol{z}_{i,1}\right)\bigr{)}+t_{i,2}(\boldsymbol{z}_{i,1})\,,$

where $s_{i,a}$ and $t_{i,a}$ are scale and translation transformation
operators, respectively, parameterized by fully-connected multi-layer-
perceptrons, while $\odot$ denotes the element-wise direct product.

Once passed through the coupling layer, the output
$\boldsymbol{z}_{i}=\\{\boldsymbol{z}_{i,1},\boldsymbol{z}_{i,2}\\}$ of the
two channels is concatenated to the final output
$f_{i}(\boldsymbol{z}_{i-1})=\boldsymbol{z}_{i}$. The full architecture
consists of $n$ sequential coupling blocks. Note that in the inverse
direction,

$\displaystyle\boldsymbol{z}_{i,2}=\bigl{(}\boldsymbol{z}_{i+1,2}-t_{i,2}\left(\boldsymbol{z}_{i+1,1}\right)\bigr{)}\odot\exp\bigl{(}-s_{i,2}\left(\boldsymbol{z}_{i+1,1}\right)\bigr{)}\,,$
(24)
$\displaystyle\boldsymbol{z}_{i,1}=\bigl{(}\boldsymbol{z}_{i+1,1}-t_{i,1}\left(\boldsymbol{z}_{i,2}\right)\bigr{)}\odot\exp\bigl{(}-s_{i,1}\left(\boldsymbol{z}_{i,2}\right)\bigr{)}\,,$

the $s_{i,a}$ and $t_{i,a}$ transformations are still evaluated in the forward
direction and thus do not require a tractable inverse.

By construction, the Jacobian matrix $J_{f}$ for each coupling block is upper
triangular, which allows for an efficient computation of its determinant

$\displaystyle\det J_{f}(\boldsymbol{z})=\det\frac{\partial
f_{ij}}{\partial\boldsymbol{z}}$
$\displaystyle=\det\begin{pmatrix}\text{diag}\left\\{\exp\bigl{(}s_{i,1}(\boldsymbol{z}_{i,2})\bigr{)}\right\\}&\cdots\\\
0&\text{diag}\left\\{\exp\bigl{(}s_{i,2}(\boldsymbol{z}_{i+1,1})\bigr{)}\right\\}\end{pmatrix}$
(25)
$\displaystyle=\prod\exp\bigl{(}s_{i,1}(\boldsymbol{z}_{i,2})\bigr{)}\prod\exp\bigl{(}s_{i,2}(\boldsymbol{z}_{i+1,1})\bigr{)}\,.$

Because the transformations $s_{i,a}$ and $t_{i,a}$ can be arbitrarily
complicated, the full architecture can be conditioned by concatenating labels
$\boldsymbol{c}$ to the inputs of $s_{i,a}$ and $t_{i,a}$, i.e.,
$s_{i,a}(\boldsymbol{z}),t_{i,a}(\boldsymbol{z})\to
s_{i,a}(\boldsymbol{z},\boldsymbol{c}),t_{i,a}(\boldsymbol{z},\boldsymbol{c})$.

###  B.2 Bayesian neural networks

When fitting a model to a data sample, it is often useful to understand the
correlations and uncertainties related to the best-fit parameters. These
uncertainties provide both information on the stability of the fit as well as
information on the statistical variations within the data sample. Training a
generative neural network is akin to a model fit, involving the optimization
of network parameters to minimize a learning objective that produces samples
matching the training dataset. As such, it is informative to understand the
uncertainties associated with the network parameters. In deterministic neural
networks, model parameters are single valued and remain fixed after training.
There are a number of proposed methods for evaluating model uncertainties in
deterministic networks, including the incorporation of drop-out layers [29],
$k$-folding cross-validation [30], network ensemble averaging [31], etc.

However, these methods are either prescription-dependent, e.g., a choice of
the drop-out scheme, with no guarantee of comprehensive coverage, or require
additional training cycles, making them computationally prohibitive for
sufficiently complex networks. An alternative to these methods, which provides
a systematic and statistically coherent assignment of uncertainties to model
output, can be provided by Bayesian neural networks (BNNs) [32, 33]. The major
difference between deterministic neural networks and their Bayesian counter-
parts resides in the conversion of single-valued network parameters to
parameters that are sampled according to a posterior, approximated as a
product of normal distributions with mean and variance learned from the
training dataset. In this context, the stability and uncertainty of model
output is understood over an ensemble of samplings in network parameter space.

Consider a BNN whose goal is to accurately model the functional relationship
$\boldsymbol{y}=f(\boldsymbol{x})$. The BNN is parameterized by model
parameters $\boldsymbol{\theta}$ distributed before training according to a
prior ${\mathcal{P}}(\boldsymbol{\theta})$ and the model output
$f(\boldsymbol{x})$ is understood as a likelihood describing the probability
${\mathcal{P}}(\boldsymbol{y}|\boldsymbol{x},\boldsymbol{\theta})$ of output
$\boldsymbol{y}$, given the model parameters $\boldsymbol{\theta}$ and input
$\boldsymbol{x}$. After training the model with a labeled dataset of $N$ pairs
$\\{(\boldsymbol{x}_{n},\boldsymbol{y}_{n})\\}_{n=1}^{N}$, the probability of
a particular output $\boldsymbol{y}$ for a given $\boldsymbol{x}$, the
posterior predictive, may be written as a marginalization over the model
parameter space

${\mathcal{P}}(\boldsymbol{y}|\boldsymbol{x})=\int\text{d}\boldsymbol{\theta}\,{\mathcal{P}}(\boldsymbol{y}|\boldsymbol{x},\boldsymbol{\theta}){\mathcal{P}}(\boldsymbol{\theta}|\\{(\boldsymbol{x}_{n},\boldsymbol{y}_{n})\\}_{n=1}^{N})\,,$
(26)

where
${\mathcal{P}}(\boldsymbol{\theta}|\\{(\boldsymbol{x}_{n},\boldsymbol{y}_{n})\\}_{n=1}^{N})$
represents the posterior distribution determined by the training data.

In practice, the actual form of
${\mathcal{P}}(\boldsymbol{\theta}|\\{(\boldsymbol{x}_{n},\boldsymbol{y}_{n})\\}_{n=1}^{N})$
is analytically intractable and thus difficult to extract, although possible
using Markov chain Monte Carlo techniques. Because of this, it is common
practice to approximate this posterior through variational inference, where we
approximate the posterior with a proposal distribution
${\mathcal{Q}}(\boldsymbol{\theta})$, typically chosen as a product of per
parameter Gaussians, with all means and variances collectively denoted by
$\boldsymbol{\phi}$. An accurate model of the posterior is one where the
difference between
${\mathcal{P}}(\boldsymbol{\theta}|\\{(\boldsymbol{x}_{n},\boldsymbol{y}_{n})\\}_{n=1}^{N})$
and ${\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi})$ is minimized. This
can be achieved by minimizing the KL divergence

$\text{min
}\text{KL}\bigl{(}{\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi}),{\mathcal{P}}(\boldsymbol{\theta}|\\{(\boldsymbol{x}_{n},\boldsymbol{y}_{n})\\}_{n=1}^{N})\bigr{)}\,,$
(27)

where

$\text{KL}\bigl{(}{\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi}),{\mathcal{P}}(\boldsymbol{\theta}|\\{(\boldsymbol{x}_{n},\boldsymbol{y}_{n})\\}_{n=1}^{N})\bigr{)}=-\int\text{d}\boldsymbol{\theta}\,{\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi})\log\frac{{\mathcal{P}}(\boldsymbol{\theta}|\\{(\boldsymbol{x}_{n},\boldsymbol{y}_{n})\\}_{n=1}^{N})}{{\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi})}\,.$
(28)

Bayes’ theorem allows us to rewrite the intractable posterior

${\mathcal{P}}(\boldsymbol{\theta}|\\{(\boldsymbol{x}_{n},\boldsymbol{y}_{n})\\}_{n=1}^{N})=\frac{\prod_{n=1}^{N}{\mathcal{P}}(\boldsymbol{y}_{n}|\boldsymbol{\theta},\boldsymbol{x}_{n}){\mathcal{P}}(\boldsymbol{\theta})}{{\mathcal{P}}(\\{\boldsymbol{y}_{n^{\prime}}\\}_{n^{\prime}=1}^{N}|\\{\boldsymbol{x}_{n^{\prime}}\\}_{n^{\prime}=1}^{N})}$
(29)

where
${\mathcal{P}}(\boldsymbol{y}_{n}|\boldsymbol{\theta},\boldsymbol{x}_{n})$ is
the per-event likelihood and ${\mathcal{P}}(\boldsymbol{\theta})$ denotes the
assigned prior on the model parameters. Introducing this into the KL
divergence we obtain

$\displaystyle\text{KL}\bigl{(}{\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi}),{\mathcal{P}}(\boldsymbol{\theta}|$
$\displaystyle\\{(\boldsymbol{x}_{n},\boldsymbol{y}_{n})\\}_{n=1}^{N})\bigr{)}=$
(30) $\displaystyle=$
$\displaystyle-\int\text{d}\boldsymbol{\theta}\,{\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi})\log\frac{\prod_{n=1}^{N}{\mathcal{P}}(\boldsymbol{y}_{n}|\boldsymbol{\theta},\boldsymbol{x}_{n}){\mathcal{P}}(\boldsymbol{\theta})}{{\mathcal{P}}(\\{\boldsymbol{y}_{n^{\prime}}\\}_{n^{\prime}=1}^{N}|\\{\boldsymbol{x}_{n^{\prime}}\\}_{n^{\prime}=1}^{N}){\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi})}$
$\displaystyle=$
$\displaystyle\log{\mathcal{P}}(\\{\boldsymbol{y}_{n^{\prime}}\\}_{n^{\prime}=1}^{N}|\\{\boldsymbol{x}_{n^{\prime}}\\}_{n^{\prime}=1}^{N})-\mathcal{L}_{\text{ELBO}}\,.$

Because the KL divergence is non-negative and the evidence $\log{\mathcal{P}}$
is not a function of $\boldsymbol{\phi}$, the minimization in eq. 27 is
achieved by maximizing the so-called evidence lower bound (ELBO) contribution,
defined within the square brackets,

$\displaystyle\mathcal{L}_{\text{ELBO}}$
$\displaystyle=\int\text{d}\boldsymbol{\theta}\,{\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi})\left[\sum_{n=1}^{N}\log{\mathcal{P}}(\boldsymbol{y}_{n}|\boldsymbol{\theta},\boldsymbol{x}_{n})+\log\frac{{\mathcal{P}}(\boldsymbol{\theta})}{{\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi})}\right]$
(31)
$\displaystyle=\sum_{n=1}^{N}\int\text{d}\boldsymbol{\theta}\,{\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi})\log{\mathcal{P}}(\boldsymbol{y}_{n}|\boldsymbol{\theta},\boldsymbol{x}_{n})-\text{KL}\bigl{(}{\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi}),{\mathcal{P}}(\boldsymbol{\theta})\bigr{)}$
$\displaystyle\approx\frac{1}{M}\sum_{j=1}^{M}\sum_{n=1}^{N}\log{\mathcal{P}}(\boldsymbol{y}_{n}|\boldsymbol{\theta}_{j},\boldsymbol{x}_{n})-\text{KL}\bigl{(}{\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi}),{\mathcal{P}}(\boldsymbol{\theta})\bigr{)}\text{,
where
}\boldsymbol{\theta}_{j}\sim{\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi})\,.$

Above, we have approximated in the last line the expectation value of
$\sum_{n=1}^{N}\log{\mathcal{P}}(\boldsymbol{y}_{n}|\boldsymbol{\theta},\boldsymbol{x}_{n})$
sampled over ${\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi})$ as a
summed average of
$\sum_{n=1}^{N}\log{\mathcal{P}}(\boldsymbol{y}_{n}|\boldsymbol{\theta},\boldsymbol{x}_{n})$
evaluated at $M$ points of $\boldsymbol{\theta}$ distributed according to
${\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi})$. Maximization of the
ELBO loss $\mathcal{L}_{\text{ELBO}}$ also minimizes the KL divergence in eq.
27, with the benefit that it requires no knowledge about the intractable
distributions
${\mathcal{P}}(\boldsymbol{\theta}|\\{(\boldsymbol{x}_{n},\boldsymbol{y}_{n})\\}_{n=1}^{N})$
and
${\mathcal{P}}(\\{\boldsymbol{y}_{n^{\prime}}\\}_{n^{\prime}=1}^{N}|\\{\boldsymbol{x}_{n^{\prime}}\\}_{n^{\prime}=1}^{N})$.
The first term in the ELBO loss drives the model to provide an accurate fit to
the training data, while the second term acts as a regulator by weighting
possible model parameters with a chosen prior
${\mathcal{P}}(\boldsymbol{\theta})$.

###  B.3 Bayesian normalizing flows

Incorporating the Bayesian framework into the normalizing flow architecture,
section B.1, only requires replacing the deterministic transformations
$s_{i,a}$ and $t_{i,a}$ in eq. 23 with their Bayesian counterparts, i.e., the
BNNs. The full normalizing flow network architecture with BNN subnetworks is
referred to as a Bayesian normalizing flow (BNF) [17]. The learning objective
is equivalent to eq. 31 with the likelihood
${\mathcal{P}}(\boldsymbol{y}_{n}|\boldsymbol{\theta}_{j},\boldsymbol{x}_{n})$
now the NF model likelihood
${\mathcal{P}_{X}}(\boldsymbol{x};\boldsymbol{\theta},\boldsymbol{c})$, i.e.,
with the replacement $\boldsymbol{y}\to\boldsymbol{x}$ in the notation, and
with no additional input measurements, only the model parameters
$\boldsymbol{\theta}$ which determine $F(\boldsymbol{x}|\boldsymbol{\theta})$.
The ELBO loss function in eq. 31 is therefore replaced by the following loss
function, which includes a minus sign to minimize rather than maximize,

$\displaystyle\mathcal{L}_{\text{BNF}}=$
$\displaystyle-\sum_{n=1}^{N}\mathbb{E}_{\boldsymbol{\theta}\sim{\mathcal{Q}}(\boldsymbol{\theta},\boldsymbol{\phi})}\left[\log
p_{X}^{F}(\boldsymbol{x}_{n};\boldsymbol{\theta},\boldsymbol{c}_{n})\right]+\text{KL}\bigl{(}{\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi}),{\mathcal{P}}(\boldsymbol{\theta})\bigr{)}$
(32) $\displaystyle=$
$\displaystyle-\sum_{n=1}^{N}\mathbb{E}_{\boldsymbol{\theta}\sim{\mathcal{Q}}(\boldsymbol{\theta},\boldsymbol{\phi})}\biggl{[}\log{\mathcal{P}_{Z}}\bigl{(}F^{-1}(\boldsymbol{x}_{n};\boldsymbol{\theta},\boldsymbol{c}_{n})\bigr{)}$
$\displaystyle\quad+\log\bigl{|}\det
J_{F}\left(F^{-1}(\boldsymbol{x}_{n};\boldsymbol{\theta},\boldsymbol{c}_{n})\right)\bigr{|}\biggr{]}+\text{KL}\bigl{(}{\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi}),{\mathcal{P}}(\boldsymbol{\theta})\bigr{)}$
$\displaystyle\approx$
$\displaystyle-\sum_{n=1}^{N}\frac{1}{M}\sum_{m=1}^{M}\biggl{\\{}\log{\mathcal{P}_{Z}}\bigl{(}F^{-1}(\boldsymbol{x}_{n};\boldsymbol{\theta}_{m},\boldsymbol{c}_{n})\bigr{)}$
$\displaystyle\quad+\log\bigl{|}\det
J_{F}\left(F^{-1}(\boldsymbol{x}_{n};\boldsymbol{\theta}_{m},\boldsymbol{c}_{n})\right)\bigr{|}\biggr{\\}}+\text{KL}\bigl{(}{\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi}),{\mathcal{P}}(\boldsymbol{\theta})\bigr{)}\,.$

We assume a two dimensional standard normal latent space, along with both a
Gaussian prior ${\mathcal{P}}(\boldsymbol{\theta})$ and a Gaussian variational
distribution ${\mathcal{Q}}(\boldsymbol{\theta})$, i.e.,

${\mathcal{P}}(\boldsymbol{\theta};\boldsymbol{\mu}_{{\mathcal{P}}},\boldsymbol{\sigma}_{{\mathcal{P}}})=\mathcal{N}(\boldsymbol{\theta};\boldsymbol{\mu}_{{\mathcal{P}}},\boldsymbol{\sigma}_{{\mathcal{P}}}),\qquad{\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\mu}_{{\mathcal{Q}}},\boldsymbol{\sigma}_{{\mathcal{Q}}})=\mathcal{N}(\boldsymbol{\theta};\boldsymbol{\mu}_{{\mathcal{Q}}},\boldsymbol{\sigma}_{{\mathcal{Q}}})\,.$
(33)

The KL divergence for $D$ parameters is then given by

$\text{KL}\bigl{(}{\mathcal{Q}}(\boldsymbol{\theta};\boldsymbol{\phi}),{\mathcal{P}}(\boldsymbol{\theta})\bigr{)}=\sum_{d=1}^{D}\frac{\sigma_{{\mathcal{Q}},d}^{2}-\sigma_{{\mathcal{P}},d}^{2}+(\mu_{{\mathcal{Q}},d}-\mu_{{\mathcal{P}},d})^{2}}{2\sigma_{{\mathcal{P}},d}^{2}}+\log\frac{\sigma_{{\mathcal{P}},d}}{\sigma_{{\mathcal{Q}},d}}.$
(34)

Choosing a standard normal prior, $\mu_{{\mathcal{P}},d}=0$,
$\sigma_{{\mathcal{P}},d}=1$ for $d=1,\ldots,D$, leaves us with our final BNF
loss function

$\begin{split}{\mathcal{L}_{\text{BNF}}}=&\sum_{n=1}^{N}\frac{1}{M}\sum_{m=1}^{M}\left\\{\frac{||F^{-1}(\boldsymbol{x}_{n};\boldsymbol{\theta}_{m},\boldsymbol{c}_{n})||^{2}_{2}}{2}-\log\bigl{|}\det
J_{F}\left(F^{-1}(\boldsymbol{x}_{n};\boldsymbol{\theta}_{m},\boldsymbol{c}_{n})\right)\bigr{|}\right\\}\\\
&+\sum_{d=1}^{D}\left[\frac{1}{2}(\sigma_{{\mathcal{Q}},d}^{2}+\mu_{{\mathcal{Q}},d}^{2}-1)-\log\sigma_{{\mathcal{Q}},d}\right]\,.\end{split}$
(35)

Above, the expression in the first line represents the same NF loss function
as $\mathcal{L}_{\text{NF}}$ in eq. 13, but now averaged over $M$ samplings of
the network parameters. In practice, it is typical to set $M=1$ in order to
reduce the computational cost of training, with the understanding that the
constraints imposed on the Jacobian structure will ensure that the mapping to
and from the latent space will remain stable with non-divergent gradients. The
second term represents the optimization of the individual network weight
distributions to best replicate the intractable posterior
${\mathcal{P}}(\boldsymbol{\theta}|\\{\boldsymbol{x}_{n}\\}_{n=1}^{N})$ via a
product of Gaussians with tunable $\boldsymbol{\mu}_{{\mathcal{Q}}}$ and
$\boldsymbol{\sigma}_{{\mathcal{Q}}}$ that are not too far away from the prior
values 0 and 1.

###  B.4 Interpreting BNF ensembles

After training the BNF network one obtains a probabilistic model defined over
a distribution of network parameters $\boldsymbol{\theta}$. This distribution
should, in principle, contain information about both the model stability and
uncertainties due to the training dataset. That is, the BNF defines an
envelope of possible NFs defined by different values of $\boldsymbol{\theta}$
distributed according to probability distribution
${\mathcal{Q}}(\boldsymbol{\theta})$. Concrete values of $\boldsymbol{\theta}$
give a particular realization of the NF, i.e., a particular map between the
latent and target spaces. All observable quantities should be computed as
averages over many samplings of the network parameters $\boldsymbol{\theta}$.

To assess how the BNF encodes the training uncertainties into the learned
densities, we follow Butter:2021csz,refs. ref. [34] and consider the
statistical distribution of an observable $g$ defined over a generated
dataset, $g(\\{\boldsymbol{x}_{i}\\}_{i=1}^{N_{\text{gen}}})$, where
$N_{\text{gen}}$ denotes the sample size of the generated data. For a given
set of parameters $\boldsymbol{\theta}$,
$g(\\{\boldsymbol{x}_{i}\\}_{i=1}^{N_{\text{gen}}})$ will have a likelihood
${\mathcal{P}}(g(\\{\boldsymbol{x}_{i}\\}_{i=1}^{N_{\text{gen}}})|\boldsymbol{\theta})$.
An example of such an observable $g$ is, e.g., the number of times an emission
falls into a given bin, in which case the likelihood
${\mathcal{P}}(g(\\{\boldsymbol{x}_{i}\\}_{i=1}^{N_{\text{gen}}})|\boldsymbol{\theta})$
is a Poisson distribution with the average rate

$\langle
N_{\text{bin}}\rangle_{\boldsymbol{\theta}}=N_{\text{gen}}\int_{\boldsymbol{x}\in\text{bin}}\text{d}\boldsymbol{x}\,{\mathcal{P}}(\boldsymbol{x}|\boldsymbol{\theta})\,,$
(36)

and the usual related variance

$\sigma_{\text{bin}|\boldsymbol{\theta}}^{2}=\langle
N_{\text{bin}}\rangle_{\boldsymbol{\theta}}\,.$ (37)

The integral in $\boldsymbol{x}$ can be approximated by sampling from
${\mathcal{P}}(\boldsymbol{x}|\boldsymbol{\theta})$. The likelihood
${\mathcal{P}}(g(\\{\boldsymbol{x}_{i}\\}_{i=1}^{N_{\text{gen}}})|\boldsymbol{\theta})$
will in turn determine the mean value of the observable,
$\mathbb{E}_{X[N_{\text{gen}}]|\boldsymbol{\theta}}[g]$, and its variance,
$\sigma^{2}_{X[N_{\text{gen}}]|\boldsymbol{\theta}}[g]$, where
$X[N_{\text{gen}}]$ is the space of all possible datasets of size
$N_{\text{gen}}$. In general, these will not be analytic functions of
$\boldsymbol{\theta}$ and will need to be determined numerically.

Since one needs to take into account all possible networks this affects the
expectation value and the variance of $g$,

$\displaystyle\mathbb{E}_{X[N_{\text{gen}}],\Theta}[g]$
$\displaystyle=\int\text{d}\boldsymbol{\theta}\,{\mathcal{P}}(\boldsymbol{\theta}|\\{x_{j}\\}_{j=1}^{N_{\text{train}}})\mathbb{E}_{X[N_{\text{gen}}]|\boldsymbol{\theta}}[g]\,,$
(38) $\displaystyle\sigma^{2}_{X[N_{\text{gen}}],\Theta}$
$\displaystyle=\int\text{d}\boldsymbol{\theta}\,{\mathcal{P}}(\boldsymbol{\theta}|\\{x_{j}\\}_{j=1}^{N_{\text{train}}})\mathbb{E}_{X[N_{\text{gen}}]|\boldsymbol{\theta}}[(g-\mathbb{E}_{X[N_{\text{gen}}],\Theta}[g])^{2}]$
$\displaystyle=\int\text{d}\boldsymbol{\theta}\,{\mathcal{P}}(\boldsymbol{\theta}|\\{x_{j}\\}_{j=1}^{N_{\text{train}}})\bigl{(}\mathbb{E}_{X[N_{\text{gen}}]|\boldsymbol{\theta}}[g^{2}]-\mathbb{E}_{X[N_{\text{gen}}]|\boldsymbol{\theta}}[g]^{2}$
$\displaystyle\quad+\left(\mathbb{E}_{X[N_{\text{gen}}]|\boldsymbol{\theta}}[g]-\mathbb{E}_{X[N_{\text{gen}}],\Theta}[g]\right)^{2}\bigr{)}$
$\displaystyle=\sigma^{2}_{\text{gen}}+\sigma^{2}_{\text{train}}\,,$

where $\mathbb{E}_{X[N_{\text{gen}}],\Theta}$ is the expectation value over
the whole distribution of $\boldsymbol{\theta}$, $N_{\text{train}}$ denotes
the sample size of the training data, and $\sigma^{2}_{\text{gen}}$ and
$\sigma^{2}_{\text{train}}$ are given in the third and the fourth lines of eq.
38, respectively. For the relevant example where $g$ is the number of times an
emission falls into a given bin or region, we have

$\displaystyle\mathbb{E}_{X[N_{\text{gen}}],\Theta}[N_{\text{bin}}]$
$\displaystyle=\int\text{d}\boldsymbol{\theta}\,{\mathcal{P}}(\boldsymbol{\theta}|\\{x_{j}\\}_{j=1}^{N_{\text{train}}})\langle
N_{\text{bin}}\rangle_{\boldsymbol{\theta}}=\langle N_{\text{bin}}\rangle\,,$
(39) $\displaystyle\sigma_{\text{gen}}^{2}$
$\displaystyle=\int\text{d}\boldsymbol{\theta}\,{\mathcal{P}}(\boldsymbol{\theta}|\\{x_{j}\\}_{j=1}^{N_{\text{train}}})\sigma_{\text{bin}|\boldsymbol{\theta}}^{2}=\langle
N_{\text{bin}}\rangle\,,$ $\displaystyle\sigma_{\text{train}}^{2}$
$\displaystyle=\int\text{d}\boldsymbol{\theta}\,{\mathcal{P}}(\boldsymbol{\theta}|\\{x_{j}\\}_{j=1}^{N_{\text{train}}})\left(\langle
N_{\text{bin}}\rangle_{\boldsymbol{\theta}}-\langle
N_{\text{bin}}\rangle\right)^{2}\,.$

We observe how $\sigma_{\text{gen}}$ is the uncertainty of a Poisson
distributed random variable whose expected rate is given by $\langle
N_{\text{bin}}\rangle$. That is, in this case $\sigma_{\text{gen}}$ will not
depend on the characteristics of the posterior except for its mean, and thus
does not include any training uncertainties. This is reflected in the fact
that although the posterior depends on the training dataset of size
$N_{\text{train}}$, $\sigma_{\text{gen}}$ itself depends exclusively on the
size of the generated dataset of size $N_{\text{gen}}$. In general, the
statistical variance of a randomly distributed $g$ whose distribution is
determined by the BNF is captured by $\sigma_{\text{gen}}$ and does not
vanish, even if we collapse the posterior
${\mathcal{P}}(\boldsymbol{\theta}|\\{x_{j}\\}_{j=1}^{N_{\text{train}}})$ to a
delta function. For instance, for the example where $g$ is the number of times
and emission falls into a given bin, we have

$\displaystyle\sigma_{\text{gen}}^{2}$
$\displaystyle=\int\text{d}\boldsymbol{\theta}\,\delta\bigl{(}\boldsymbol{\theta}-\boldsymbol{\theta}^{\text{MAP}}(\\{x_{j}\\}_{j=1}^{N_{\text{train}}})\bigr{)}\left(\mathbb{E}_{X[N_{\text{gen}}]|\boldsymbol{\theta}}[g^{2}]-\mathbb{E}_{X[N_{\text{gen}}]|\boldsymbol{\theta}}[g]^{2}\right)$
(40)
$\displaystyle=\int\text{d}\boldsymbol{\theta}\,\delta\bigl{(}\boldsymbol{\theta}-\boldsymbol{\theta}^{\text{MAP}}(\\{x_{j}\\}_{j=1}^{N_{\text{train}}})\bigr{)}\langle
N_{\text{bin}}\rangle_{\boldsymbol{\theta}}$ $\displaystyle=\langle
N_{\text{bin}}\rangle_{\boldsymbol{\theta}^{\text{MAP}}}=N\int_{\boldsymbol{x}\in\text{bin}}\text{d}\boldsymbol{x}\,{\mathcal{P}}(\boldsymbol{x}|\boldsymbol{\theta}^{\text{MAP}})\,,$

where $\boldsymbol{\theta}^{\text{MAP}}$ is the maximum a posteriori (MAP)
estimate. That is, we recover the Poisson uncertainty
$\sigma_{\text{gen}}^{2}=\langle N_{\text{bin}}\rangle$. Again, we note that
although $\boldsymbol{\theta}^{\text{MAP}}$ depends on the training dataset of
size $N_{\text{train}}$, the Poisson uncertainty of the observable itself
reflects the size of the studied dataset of size $N_{\text{gen}}$.

The fourth line of eq. 38, $\sigma_{\text{train}}^{2}$, captures how the
expected values of the observables change due to uncertainties in
$\boldsymbol{\theta}$ and does vanish if we collapse the posterior to a delta
function. The larger $\sigma_{\text{train}}^{2}$ is, the larger the variations
in possible networks sampled from the posterior and the larger the set of
$\boldsymbol{\theta}$ parameter values resulting in outputs consistent with
the training dataset.

###  B.5 Architecture impact on the BNF uncertainty

In section 3, we showed how the BNF captures the training uncertainty.
However, this uncertainty also depends on the BNF architecture. In fig. 8 we
show the uncertainty variations as a function of the number of nodes per
hidden layer in the BNF, while fig. 9 shows two examples of learned $p_{{z}}$
distributions and the values of the associated BNF parameters
$\boldsymbol{\mu}_{{\mathcal{Q}}}$ and $\boldsymbol{\sigma}_{{\mathcal{Q}}}$
(see section B.4 for the details about the notation). We especially highlight
the model with 32 nodes, which was used in the main text for both fig. 4 and
fig. 5. In fig. 8 this model is denoted with a star. Although the results in
fig. 8 and fig. 9 do not represent an exhaustive scan, since the number of
hidden layers per module and number of modules remained fixed to 2 modules
with 4 layers each, we can nevertheless begin to characterize the behavior we
obtain for alternative models.

Figure 8: Relative total uncertainty for the most populous bin, as a function
of the model architecture. We scan over the number of nodes per hidden layer
and obtain the total uncertainty for the same bin as a function of said nodes
per hidden layer. The chosen bin is the one with the largest expected event
count for the model with 32 nodes per hidden layer, the default architectural
choice, denoted with a star.

Figure 9: Impact of the model architecture on the quality of the model. The
top (bottom) row corresponds to a model with 1 (32) nodes per hidden layer.
(left) Comparison of the training dataset and the BNF predictions. (right)
Visualization of the inferred parameters of the BNF posterior distribution for
the model weights, the means $\boldsymbol{\mu}_{{\mathcal{Q}}}$ and standard
deviations $\boldsymbol{\sigma}_{{\mathcal{Q}}}$ of the approximate Gaussian
posterior.

When the number of nodes is too low, the BNF model is simultaneously both
inflexible (the $p_{{z}}$ distribution does not match the training data) and
overly certain. The learned standard deviations
$\boldsymbol{\sigma}_{{\mathcal{Q}}}$ for the posterior are too small,
collapsing the posterior for $\boldsymbol{\mu}_{{\mathcal{Q}}}$ to a delta
function. The BNF then always samples very similar NFs, whose parameters are
given by the learned means $\boldsymbol{\mu}_{{\mathcal{Q}}}$. This behavior
can be seen in the two upper panels in fig. 9, which shows the results for a
learned BNF model with only 1 node per hidden layer.

Conversely, when the number of nodes is large enough, and we appropriately
regularize, the learned average distribution matches much better the data, and
the posterior is broad enough to appropriately capture the model
uncertainties. We observe how $\boldsymbol{\sigma}_{{\mathcal{Q}}}$ is now
large and thus allows a higher variation on the sampled $\boldsymbol{\theta}$.
The learned means $\boldsymbol{\mu}_{{\mathcal{Q}}}$ should be non-zero to
yield a non-trivial MAP distribution, although the specific values are hard to
analyze due to the inherent complexity of the neural network. Additionally,
the model shows a subset of weights which are effectively removed by having
very low mean and variance. This indicates redundant parameters which could be
removed by pruning or by a refinement of the architecture. The lower two
panels in fig. 9 demonstrate this behavior for a learned model with 32 nodes
per hidden layer.

We observe that the uncertainty as displayed in fig. 8 is a good reflection of
training quality but not necessarily a good metric for model selection. This
is exemplified by comparing the 16 node model to the 32 node model. We observe
that the former has a larger uncertainty despite providing a reasonable
description of data. In this case, the uncertainty increases because the model
provides a poor prediction in the example bin. Since the model is both
sufficiently expressive and well trained, it recognizes this mis-modeling and
reflects it in the increased uncertainty. If we were interested in lighter
models, we could choose the 16 node model at the expense of a slightly larger
uncertainty for the largest bin. In general, model selection should take into
account all bins, as well as other considerations, such as the model size. In
this work, however, we were not interested in selecting the optimal model but
in selecting a descriptive model with sensible uncertainties.

## References

  * [1] B. Andersson, G. Gustafson, G. Ingelman and T. Sjostrand,  _Parton Fragmentation and String Dynamics_ ,  Phys. Rept. 97, 31 (1983),  10.1016/0370-1573(83)90080-7.
  * [2] B. Andersson,  _The Lund model_ ,  Camb. Monogr. Part. Phys. Nucl. Phys. Cosmol. 7, 1 (1997).
  * [3] R. D. Field and S. Wolfram,  _A QCD Model for e+ e- Annihilation_ ,  Nucl. Phys. B 213, 65 (1983),  10.1016/0550-3213(83)90175-X.
  * [4] T. D. Gottschalk,  _An Improved Description of Hadronization in the{QCD} Cluster Model for $e^{+}e^{-}$ Annihilation_,  Nucl. Phys. B 239, 349 (1984),  10.1016/0550-3213(84)90253-0.
  * [5] B. Webber,  _A QCD Model for Jet Fragmentation Including Soft Gluon Interference_ ,  Nucl. Phys. B 238, 492 (1984),  10.1016/0550-3213(84)90333-X.
  * [6] G. Corcella, R. Franceschini and D. Kim,  _Fragmentation Uncertainties in Hadronic Observables for Top-quark Mass Measurements_ ,  Nucl. Phys. B 929, 485 (2018),  10.1016/j.nuclphysb.2018.02.012,  1712.05801.
  * [7] N. Fischer and T. Sjöstrand,  _Thermodynamical String Fragmentation_ ,  JHEP 01, 140 (2017),  10.1007/JHEP01(2017)140,  1610.09818.
  * [8] P. Ilten, T. Menzo, A. Youssef and J. Zupan,  _Modeling hadronization using machine learning_ ,  SciPost Phys. 14, 027 (2023),  10.21468/SciPostPhys.14.3.027,  2203.04983.
  * [9] A. Ghosh, X. Ju, B. Nachman and A. Siodmok,  _Towards a deep learning model for hadronization_ ,  Phys. Rev. D 106(9), 096020 (2022),  10.1103/PhysRevD.106.096020,  2203.12660.
  * [10] J. Chan, X. Ju, A. Kania, B. Nachman, V. Sangli and A. Siodmok,  _Fitting a Deep Generative Hadronization Model_ (2023),  2305.17169.
  * [11] C. Bierlich _et al._ ,  _A comprehensive guide to the physics and usage of PYTHIA 8.3_ (2022),  10.21468/SciPostPhysCodeb.8,  2203.11601.
  * [12] M. Bahr _et al._ ,  _Herwig++ Physics and Manual_ ,  Eur. Phys. J. C 58, 639 (2008),  10.1140/epjc/s10052-008-0798-9,  0803.0883.
  * [13] J. Bellm _et al._ ,  _Herwig 7.0/Herwig++ 3.0 release note_ ,  Eur. Phys. J. C 76(4), 196 (2016),  10.1140/epjc/s10052-016-4018-8,  1512.01178.
  * [14] I. Kobyzev, S. J. Prince and M. A. Brubaker,  _Normalizing Flows: An Introduction and Review of Current Methods_ ,  IEEE Transactions on Pattern Analysis and Machine Intelligence 43(11), 3964 (2021),  10.1109/tpami.2020.2992934.
  * [15] D. Rezende and S. Mohamed,  _Variational Inference with Normalizing Flows_ ,  In F. Bach and D. Blei, eds., _Proceedings of the 32nd International Conference on Machine Learning_ , vol. 37 of _Proceedings of Machine Learning Research_ , pp. 1530–1538. PMLR, Lille, France (2015), 1505.05770.
  * [16] L. Dinh, D. Krueger and Y. Bengio,  _NICE: Non-linear Independent Components Estimation_ ,  In Y. Bengio and Y. LeCun, eds., _3rd International Conference on Learning Representations_ (2015), 1410.8516.
  * [17] B. Trippe and R. Turner,  _Conditional Density Estimation with Bayesian Normalising Flows_ (2018),  1802.04908.
  * [18] P. Skands, S. Carrazza and J. Rojo,  _Tuning PYTHIA 8.1: the Monash 2013 Tune_ ,  Eur. Phys. J. C 74(8), 3024 (2014),  10.1140/epjc/s10052-014-3024-y,  1404.5630.
  * [19] P. Abreu _et al._ ,  _Tuning and test of fragmentation models based on identified particles and precision event shape data_ ,  Z. Phys. C 73, 11 (1996),  10.1007/s002880050295.
  * [20] M. Krishnamoorthy, H. Schulz, X. Ju, W. Wang, S. Leyffer, Z. Marshall, S. Mrenna, J. Müller and J. B. Kowalkowski,  _Apprentice for Event Generator Tuning_ ,  EPJ Web Conf. 251, 03060 (2021),  10.1051/epjconf/202125103060,  2103.05748.
  * [21] P. Ilten, M. Williams and Y. Yang,  _Event generator tuning using Bayesian optimization_ ,  JINST 12(04), P04028 (2017),  10.1088/1748-0221/12/04/P04028,  1610.08328.
  * [22] A. Andreassen and B. Nachman,  _Neural Networks for Full Phase-space Reweighting and Parameter Tuning_ ,  Phys. Rev. D 101(9), 091901 (2020),  10.1103/PhysRevD.101.091901,  1907.08209.
  * [23] P. T. Komiske, E. M. Metodiev and J. Thaler,  _The Hidden Geometry of Particle Collisions_ ,  JHEP 07, 006 (2020),  10.1007/JHEP07(2020)006,  2004.04159.
  * [24] C. e. Villani,  _Optimal transport, old and new_ ,  Springer, Berlin (2008).
  * [25] F. Santambrogio,  _Optimal Transport for Applied Mathematicians_ ,  Springer, Switzerland (2015).
  * [26] A. Ramdas, N. García Trillos and M. Cuturi,  _On wasserstein two-sample testing and related families of nonparametric tests_ ,  Entropy 19(2), 47 (2017).
  * [27] C. Bierlich, P. Ilten, T. Menzo, S. Mrenna, M. Szewc, M. K. Wilkinson, A. Youssef and J. Zupan,  _Reweighting Monte Carlo Predictions and Automated Fragmentation Variations in Pythia 8_ (2023),  2308.13459.
  * [28] L. Ardizzone, T. Bungert, F. Draxler, U. Köthe, J. Kruse, R. Schmier and P. Sorrenson,  _Framework for Easily Invertible Architectures (FrEIA)_ (2018).
  * [29] Y. Gal and Z. Ghahramani,  _Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning_ ,  In M. F. Balcan and K. Q. Weinberger, eds., _Proceedings of The 33rd International Conference on Machine Learning_ , vol. 48 of _Proceedings of Machine Learning Research_ , pp. 1050–1059. PMLR, New York, New York, USA (2016).
  * [30] C. M. Bishop,  _Pattern Recognition and Machine Learning (Information Science and Statistics)_ ,  Springer-Verlag, Berlin, Heidelberg,  ISBN 0387310738 (2006).
  * [31] U. Naftaly, N. Intrator and D. Horn,  _Optimal ensemble averaging of neural networks_ ,  Network: Computation in Neural Systems 8(3), 283 (1997),  10.1088/0954-898X/8/3/004.
  * [32] D. J. C. Mackay,  _Probable networks and plausible predictions — a review of practical Bayesian methods for supervised neural networks_ ,  Network: Computation in Neural Systems 6(3), 469 (1995),  10.1088/0954-898X_6_3_011,  https://doi.org/10.1088/0954-898X_6_3_011.
  * [33] R. M. Neal,  _Bayesian Learning for Neural Networks_ ,  Springer-Verlag, Berlin, Heidelberg,  ISBN 0387947248 (1996).
  * [34] A. Butter, T. Heimel, S. Hummerich, T. Krebs, T. Plehn, A. Rousselot and S. Vent,  _Generative Networks for Precision Enthusiasts_ (2021),  2110.13632.