A CCELERATING HEP SIMULATIONS WITH N EURAL I MPORTANCE S AMPLING
Anonymous authors
Paper under double-blind review
A BSTRACT
Virtually all high-energy-physics (HEP) simulations for the LHC rely on Monte
Carlo using importance sampling by means of the VEGAS algorithm.
However, complex high-precision calculations have become a challenge for the standard toolbox.
As a result, there has been keen interest in HEP for modern machine learning to power adaptive sampling.
Despite previous work proving that normalizing-ﬂow-powered neural importance sampling (NIS) sometimes outperforms VEGAS, existing research has still left major questions open, which we intend to solve by introducing ZüNIS, a fully automated NIS library.
We ﬁrst show how to extend the original formulation of NIS to reuse samples over multiple gradient steps, yielding a signiﬁcant improvement for slow functions.
We then benchmark ZüNIS over a range of problems and show high performance with limited ﬁne-tuning.
The library can be used by non-experts with minimal effort, which is crucial to become a mature tool for the wider HEP public.
I NTRODUCTION
High-Energy-Physics (HEP) simulations are at the heart of the Large Hadron Collider (LHC) program for studying the fundamental laws of nature.
Most HEP predictins are expressed as expectation values, evaluated numerically as Monte Carlo (MC) integrals.
This permits both the integration of the very complex functions and the reproduction of the data selection process by experiments.
Most HEP simulations tools (Alwall et al., 2014) perform MC integrals using importance sampling, which allows to adaptively sample points to speed up convergence while keeping independent and identically distributed samples, crucial to reproduce experimental analyses which can only ingest uniformly-weighted data, typically produced by rejection sampling (see appendix A).
The most popular tool to optimize importance sampling in HEP is by far
the VEGAS algorithm (Lepage, 1980; 2021), which ﬁghts the curse of dimensionality by assuming no correlations between the variables.
While this is rarely the case in general, a good understanding of the integrand function can help signiﬁcantly.
Indeed optimized parametrizations using multichannelling (Kleiss et al., 1986; Ohl, 1999; Kleiss & Pittau, 1994)
have become bread-and-butter tools for HEP event generation simulators, with good success for leading-order (LO) calculations.
However, as simulations get more complex, either by having more complex ﬁnal states or by including higher orders in perturbation theory, performance degrades fast.
While newer approaches to adaptive importance sampling have been developped, like Population
Monte Carlo (Bugallo et al., 2017; 2015; Cappé et al., 2004; Iba, 2001) and its extensions (Cappé et al., 2008; Koblents & Míguez, 2015; Elvira et al., 2017; Douc et al., 2007; Cornuet et al., 2012), and have been successfully applied to other ﬁelds, VEGAS remains the standard tool used by essentially all modern simulation tools (Alwall et al., 2014; Bothmann et al., 2019; Bellm et al., 2016).
We suspect that the reasons for this are twofold.
First of all the lack of exchange between the two scientiﬁc communities is probably at cause and should probably not be neglected.
Second of all however, the impetus in the HEP community is to move toward more "black box" approaches where little is known of the structure of the integrand, which seems to be in tension with the situation with PMC, which is sensitive to to the initial proposal density (Beaujean & Caldwell, 2013; Cappé et al., 2008). As the main practical goal of this paper is to reduce the reliance of HEP simulations on the careful tuning of integrators, we will focus on comparing our work with the de facto HEP standard: the VEGAS algorithm. There is much room for investing computational time into improving sampling (ATLAS Collaboration, 2020): modern HEP theoretical calculations are taking epic proportions and can require hours for a single function evaluation (Jones, 2018).
Furthermore, unweighting samples can be extremely inefﬁcient, with upwards of 90% sampled points discarded (HEP Software Foundation, 2020).
More powerful importance sampling algorithms would therefore be a welcome improvement (Buckley, 2020; WG et al., 2021).
First attempts to use machine learning (ML) to address this challenge explored using classical neural networks to sample (Bendavid, 2017; Klimek & Perelstein, 2020; Chen et al., 2021) but typically suffer from excessive computational costs.
Another avenue of research has been to leverage generative models successful in other ﬁelds such as generative adversarial networks (Butter et al., 2019;
Di Sipio et al., 2019; Butter et al., 2020; Ahdida et al., 2019; Hashemi et al., 2019; Carrazza &
Dreyer, 2019).
While such approaches do improve sampling speed by a large factor, they have major limitations.
In particular, they have no theoretical guarantees of providing a correct answer on average (Matchev et al., 2021) and poor control of uncertainties.
To avoid these disadvantages, our work exploits Neural Importance Sampling (NIS)
(Müller et al., 2019; Zheng & Zwicker, 2019), which relies on normalizing ﬂows and has strong theoretical guarantees.
A number of exploratory papers have been published on using NIS for LHC simulations (Gao et al., 2020b; Bothmann et al., 2020; Gao et al., 2020a), as well as closely related variations (Bellagente et al., 2021; Stienen & Verheyen, 2021), but most studies have focused on preliminary investigation of performance without much concern for the practical usability of the method.
Indeed, training requires function evaluations, which we are trying to minimize and data-efﬁciency training is therefore an important but under-appreciated concern.
Furthermore, few authors have provided usable open source code, making the adoption of the technique in the HEP community difﬁcult.
Our contribution to improve this situation can be summarized in two items: • The introduction of a new loss function and an associated new training algorithm for NIS.
This permits the re-use of sampled points over multiple gradient descent steps, therefore making NIS much more data efﬁcient.
• The introduction of Z Ü NIS, a PyTorch -based library providing robust and usable NIS tools, usable by non-experts. It implements previously-developped ideas as well as our new training procedure and is extensively documented.
B ACKGROUND 2. I MPORTANCE SAMPLING AS AN OPTIMIZATION PROBLEM
Importance sampling relies on the interpretation of integrals as expectation values.
Indeed, let us consider an integral over a ﬁnite volume:
[Equation]
Let p be a strictly positive probability distribution over Ω , we can re-express our integral as the expectation of a random variable
[Equation]
whose mean is indeed I and whose standard deviation is σ ( f, p ) √
N , where σ ( f, p ) is the standard deviation of f ( X ) /p
( X ) for X ∼ p :
[Equation]
The problem statement of importance sampling is to ﬁnd the probability distribution function p that minimizes the variance of our estimator for a given N .
In Neural Importance Sampling, we rely on
Normalizing Flows to approximate the optimal distribution, which we can optimize using stochastic gradient descent. 2. N ORMALIZING FLOWS AND COUPLING CELLS
Normalizing ﬂows (Tabak & Vanden-Eijnden, 2010; Tabak & Turner, 2013; Rippel & Adams, 2013;
Rezende & Mohamed, 2015) provide a way to generate complex probability distribution functions from simpler ones using parametric changes of variables that can be learned to approximate a target distribution.
As such, normalizing ﬂows are diffeomorphisms: invertible, (nearly-everywhere) differentiable mappings with a differentiable inverse.
Indeed, if u ∼ p ( u ) , then T ( u ) = x ∼ q ( x ) where
[Equation]
where J T is the Jacobian determinant of T :
[Equation]
In the world of machine learning, T is called a normalizing ﬂow and is typically part of a parametric family of diffeomorphisms ( T ( · , θ ) ) such that gradients ∇ θ J T are tractable.
Coupling cell mappings perfectly satisfy this requirement (Dinh et al., 2015; 2017; Müller et al., 2018): they are neural-network-parametrized bijections whose Jacobian factor can be obtained analytically without backpropagation or expensive determinant calculation.
As such, they provide a good candidate for importance sampling as long as they can be trained to learn an unnormalized target function, which is exactly what neural importance sampling proposes.
2. N EURAL IMPORTANCE SAMPLING
Neural importance sampling was introduced in the context of computer graphics (Müller et al., 2018) and proposes to use normalizing ﬂows as a family of probability distributions over which to solve the minimization problem of importance sampling.
[Equation]
Of course, to actually do so, one needs to ﬁnd a way to explicitly evaluate L ( θ ) and the original neural importance sampling approach proposes to approximate it using importance sampling.
One needs to be careful that the gradient of the estimator of the loss need not be the estimator of the gradient of the loss.
The gradient of the loss can be expressed as
[Equation]
for which an estimator is proposed as
[Equation]
The authors also observed that other loss functions are possible which share the same global minimum as the variance based loss: for example, the Kullback-Leibler divergence D KL between two functions is also minimized when they are equal.
Such alternative loss functions are not guaranteed to work for importance sampling, but they prove quite successful in practice.
After training to minimize the loss estimator of eq.
(8), the normalizing ﬂows provides a tractable probability distribution function from which to sample points and estimate the integral.
C ONCEPTS AND ALGORITHMS
In this section we describe the original contributions of this paper.
The major conceptual innovation we provide in Z Ü NIS is a more ﬂexible and data-efﬁcient way of training normalizing ﬂows in the context of importance sampling.
This relies on a more rigorous formulation of the connection between the theoretical expression of ideal loss functions in terms of integrals and their practical realizations as random estimators than in previous literature.
We describe this improvement in section 3.1.
We also give a high-level overview of the organization of the Z Ü NIS library, which implements this new training procedure.
3. E FFICIENT TRAINING FOR IMPORTANCE SAMPLING
In this section, we describe how we train probability distributions within Z Ü NIS using gradientbased optimizers.
While the solution proposed in the original formulation of NIS deﬁned eq. (8) works, its main drawback is that it samples points from the same distribution that it tries to optimize. As a result, new points X i must be sampled from p after every gradient step, which is very inefﬁcient for slow integrands.
Our solution to this problem is to remember that the loss function is an integral, which can be evaluated by importance sampling using any PDF, not only p .
We will therefore deﬁne an auxiliary probability distribution function q ( x ) , independent from θ , from which we sample to estimate our loss function:
[Equation]
This is the basis for the general method we use for training probability distributions within Z Ü NIS, described in algorithm 2.
Because the sampling distribution is separated from the model to train, the same point sample can be reused for multiple training steps, which is not possible when using eq.
This is particularly important for high-precision particle physics predictions that involve high-order perturbative calculations or complex detector simulations because function evaluations can be extremely costly.
We show in section 4, in particular in ﬁg. 4 that reusing data indeed has a very signiﬁcant impact on data efﬁciency.
Data: Parametric PDF p ( x, θ ) Result: Trained PDF p ( x, θ ) 1 for M steps do After training, q is discarded and the integral is estimated from the optimized p .
The only constraint on q is that it yields a good enough estimate so that gradient steps improve the model. Much like in other applications of neural networks, we have found that stochastic gradient descent can yield good results despite noisy loss estimates. We propose three schemes for q : • A uniform distribution ( survey_strategy="flat" )
• A frozen copy of the model, which can be updated once in a while 1 ( survey_strategy= "forward" ) • An adaptive scheme starting with a uniform distribution and switching to sampling from a frozen model when it is expected to yield a more accurate loss estimate ( survey_strategy="adaptive_variance" ). 3. T HE Z Ü NIS LIBRARY On the practical side, Z Ü NIS is a PyTorch -based library which implements many ideas formulated in previous work but organizes them in the form of a modular software library with an easy-to-use interface and well-deﬁned building blocks.
We believe this structure will help non-specialist use it without understanding all the nuts and bolts, while experts can easily add new features to responds to their needs. The Z Ü NIS library relies on three layers of abstractions which steer the different aspects of using normalizing ﬂows to learn probability distributions from un-normalized functions and compute their integrals: • Flows , which implement a bijective mapping which transforms points and computes the corresponding Jacobian are described in appendix F. • Trainers , which provide the infrastructure to perform training steps and sample from ﬂow models are described in appendix F. • Integrators , which use trainers to steer the training of a model and compute integrals are described in appendix F.  R ESULTS
In this section, we evaluate Z Ü NIS on a variety of test functions to assess its performance and compare it to the commonly used VEGAS
algorithm
(Peter Lepage, 1978; Ohl, 1999).
We ﬁrst produce a few low dimensional examples for illustrative purposes, then move on to integrating parametric functions in various dimensions and ﬁnally evaluate performance on particle scattering matrix elements.
4. L OW - DIMENSIONAL EXAMPLES
Let us start by illustrating the effectiveness of Z Ü NIS in a low dimensional setting where we can readily visualize results.
We deﬁne three functions on the 2D unit hypercube which each illustrate some failure mode of VEGAS (see appendix C).
We ran the Z Ü NIS Integrator with default arguments over ten repetitions for each function and report the performance of the trained integral estimator compared to a ﬂat estimator and to VEGAS in table 1.
Overall, Z Ü NIS Integrators learn to sample from their target function extremely well: we outperform VEGAS by a factor 100 for the camel and the slashed circle functions and a factor 30 for the sinusoidal function and VEGAS itself provides no advantage over uniform sampling for the latter two.
[Table]
We further illustrate the performance of our trained models by comparing the target functions and density histogram for points sampled from the normalizing ﬂows in ﬁg. 1, which shows great qualitative agreement. [Figure]
4. S YSTEMATIC BENCHMARKS
Let us now take a more systematic approach to benchmarking Z Ü NIS.
We compare Z Ü NIS Integrators against uniform integration and VEGAS using the following metrics: integrand variance (a measure of convergence speed, see section 2.1), unweighting efﬁciency (a measure of the efﬁciency of exact sampling with rejection, see appendix A) and wall-clock training and sampling 2 .
Z Ü NIS improves convergence rate compared to VEGAS For this experiment, we focus on the camel function deﬁned in eq.
(11) and scan a 35 conﬁgurations spanning from 2 to 32 dimensions over function variances between 10 − 2 and 10 2 as shown in table 3.
Except in the low variance limit, Z Ü NIS can reduce the required number of points sampled to attain a given precision on integral estimates without any parameter tuning, attaining speed-ups of up to × 1000 both compared to uniform sampling and VEGAS-based importance sampling, as shown in ﬁg. 2a-2b and table 4.
Unweighting efﬁciencies are also boosted signiﬁcantly, although more mildly than variances, as shown in ﬁg. 2c-2d, which we could attribute to PDF underestimation in regions with low point density; the nature of the veto algorithm makes it very sensitive to a few bad behaving points in the whole dataset. [Figure] [Figure]
Z Ü NIS does not, however, outclass VEGAS on all metrics by far: as shown in ﬁg. 3, training is a few hundred times slower than VEGAS and sampling is 10times slower, all while Z Ü NIS runs on GPUs.
This is to be expected given the much increased computational complexity of normalizing ﬂows compared to the VEGAS algorithm.
As such, Z Ü NIS is not a general replacement for VEGAS, but provides a clear advantage for integrating time-intensive functions, where sampling is a negligible overhead, such as precision high-energy-physics simulations.
The new loss function introduced in Z Ü NIS improves data efﬁciency We have shown that
Z Ü NIS is a very performant importance sampling and event generation tool and provides signiﬁcant [Figure] improvements over existing tools, while requiring little ﬁne tuning from users.
Another key result is that the new approach to training we introduced in section 3.1 has a large positive impact on performance.
Indeed, as can be seen in ﬁg. 4, re-using samples for training over multiple epochs provides a 2- to 10-fold increase in convergence speed, making training much more data-efﬁcient.
For this experiment, we use forward sampling, where the frozen model is used to sample a batch of points which are then used for training over multiple epochs before resampling from an update of the frozen model.
As a result, we essentially reproduce the original formulation of NIS in eq.
(8) when we use a single epoch.
4. M AD G RAPH CROSS SECTION INTEGRALS
Cross-sections are integrals of quantum transition matrix elements for a a scattering process such as a LHC collision and express the probability that speciﬁc particles are produced.
Matrix elements themselves are un-normalized probability distributions for the conﬁguration of the outgoing particles: it is therefore both valuable to integrate them to know the overall frequency of a given scattering process, and to sample from them to understand how particles will be spatially distributed as they ﬂy off the collision point.
We study the performance of Z Ü NIS in comparison to VEGAS by studying three simple processes at leading order in perturbation theory, e − µ → e − µ via Z , d ¯ d → d ¯ d via Z and uc → ucg (with 3-jet cuts based on ∆ R ), see table 2 and ﬁg. 5.
We use the ﬁrst process as a very easy reference while the two other, quark-initiated processes are used to illustrate speciﬁc points.
Indeed, both feature narrow regions of their integration space with large enhancements, due respectively to Z -boson resonances and infra-red divergences.
[Figure]
We evaluate the matrix elements for these three processes by using the F ORTRAN standalone interface of M AD G RAPH 5_ A MC@NLO
(Alwall et al., 2014).
The two hadronic processes are convolved with parton-distribution functions from LHAPDF6 (Buckley et al., 2015).
We parametrize e − µ → e − µ via Z d ¯ d → d ¯ d via Z uc → ucg dimensions   normalized standard deviation 1 . 45 × 10 −6 .
57 × 10 −0 . phase space (the particle conﬁguration space) using the RAMBO on diet algorithm (Plätzer, 2013) implemented for PyTorch in T ORCH PS (Götz, 2021).
We report benchmark results in ﬁg. 6, in which we trained over 500,000 points for each process using near-default conﬁguration, scanning only over variance and Kullback-Leibler losses.
[Figure]
As previously observed, little convergence acceleration is achieved for low variance integrands like e − µ → e − µ , but unweighting still beneﬁts from NIS.
The two hadronic processes illustrate typical features for cross sections: training performance is variable and different processes are optimized by different loss function choices 3 .
The performance of d ¯ d → d ¯ d shows nice improvement with Z Ü NIS while that of uc → ucg is more limited.
This can be understood by comparing to importance sampling (see appendix D.3): it is in fact VEGAS, which performs signiﬁcantly better on uc → ucg compared to d ¯ d → d ¯ d because the parametrization of RAMBO is based on splitting invariant masses, making them aligned with the enhancements in the ucg phase space and allowing great VEGAS performance.
This drives a key conclusion for the potential role of Z Ü NIS in the HEP simulation landscape: not to replace VEGAS, but to ﬁll in the gaps where it fails due to inadequate parametrizations, as we illustrate here by using non-multichanneled d ¯ d → d ¯ d as a proxy for more complex processes.
C ONCLUSION
We have showed that Z Ü NIS can outperform VEGAS both in terms of integral convergence rate and unweighting efﬁciency on speciﬁc cases, at the cost of a signiﬁcant increase in training and sampling time, which is an acceptable tradeoff for high-precision HEP computations with high costs.
In this context, the introduction of efﬁcient training is a key element to making the most of the power of neural importance sampling where function evaluation costs are a major concern.
While further testing is required to ascertain how far NIS can ﬁll the gaps left by VEGAS for integrating complex functions, there is already good evidence that Z Ü NIS can provide needed improvements in speciﬁc cases.
We hope that the publication of a usable toolbox for NIS such as Z Ü NIS will stir a wider audience within the HEP community to apply the method so that the exact boundaries its applicability can be more clearly ascertained.
R EPRODUCIBILITY S TATEMENT
An anonymized version of the library is available on Github.
The recommended procedure to reproduce the experiments is to clone the repository and install the
Python requirements using pip install -r requirements.txt .
The data to reproduce the experiments can be generated using scripts provided in the repository at experiments/benchmarks , in which Jupyter notebooks are also available to reproduce the ﬁgures of the paper.
The following scripts are available:
• benchmarks_03/camel/run_benchmark_defaults.sh to generate camel integration data • benchmarks_04/camel/run_benchmark_defaults.sh to generate camel sampling speed data • benchmark_madgraph/ex_benchmark_emu.sh to generate e − µ → e − µ via Z integration data • benchmark_madgraph/ex_benchmark_dd.sh to generate d ¯ d → d ¯ d via Z integration data • benchmark_madgraph/ex_benchmark_ucg.sh to generate uc → ucg integration data
These scripts assume that 5 CUDA GPUs are available and run 5 benchmarks in parallel. If fewer GPUs are available, it is recommended to modify the scripts to run the benchmarking scripts sequentially (by removing the ampersand) and to adapt the --cuda=N option. A E VENT GENERATION WITH THE VETO ALGORITHM
Generating i.i.d points following an arbitrary probability distributions in high dimensions is not a priori a trivial task.
A straightforward way to obtain such data is to use rejection sampling, which can be based on any distribution q from which we can sample.
Given an i.i.d sample x 1 , . . .
, x N ∼ q , we can deﬁne weights w ( x i ) =
p ( x i ) /q
( x i ) and keep/reject points with probability w ( x i ) /w max .
The main metric for evaluating the performance of this algorithm is the unweighting efﬁciency: how much data is kept from an original sample of size N on average, which is expressed as
[Equation]
B B ACKWARD S AMPLING A LGORITHM
Data: Parametric PDF p ( x, θ )
Result:
Trained PDF p ( x, θ ) 1 for M steps do
Sample a batch x 1 , . . .
, x n batch from q ;
Compute the sampling PDF values q ( x i ) ;
Compute function values f ( x i ) ;
Start tracking gradients with respect to θ ; for N steps do
Compute the PDF values from the parametric PDF p ( x i , θ ) ;
Estimate the loss ˆ L = 1 N n batch (cid:88) i = f
( x i )p ( x, θ ) q ( x ) ;
Compute ∇ θ ˆ L using backpropagation;
Set θ ← θ − η ∇ θ ˆ L ;
Reset gradients; end 13 end 14 return p ( x, θ ) ;
C F UNDAMENTAL LIMITATIONS OF THE VEGAS ALGORITHM
We deﬁne three functions over the two dimensional hypercube: [Equation] (cid:32) − (cid:18) [Equation]
µ 1 ) σ (cid:19) 2 (cid:33) + exp (cid:32) − (cid:18) x − µ 2 ) σ (cid:19) 2 (cid:33) , (11) f
[Equation]
( x )
[Equation]
min (cid:34) 1 , exp (cid:32) − (cid:18) | x | − r σ ∅ (cid:19) 2 (cid:33) + exp (cid:32) − (cid:18) a · x σ ∅ (cid:19) 2 (cid:33)(cid:35) (12) f
[Equation]
( x )
[Equation]
[Equation]
) 2 , (13) to which we refer respectively as the camel, slashed circle and sinusoidal target functions.
We set their parameters as follows
[Equation]
We chose these functions because they illustrate different failure modes of the VEGAS algorithm:
• Because of the factorized PDF of VEGAS, the camel function leads to ’phantom peaks’ in the off diagonal.
This problem grows exponentially with the number of dimensions but can be addressed by a change of variable to align the peaks with an axis of integration.
• The sinuoidal function makes it nearly impossible for VEGAS to provide any improvement: the marginal PDF over each variable of integration is nearly constant.
Again this type of issue can be addressed by a change of variable provided one knows how to perform it.
• The slashed circle function is an example of a hard-for-VEGAS function that cannot be improved signiﬁcantly by a change of variables.
One can instead use multi-channeling, but this requires a lot of prior knowledge and has a computational costs since each channel is its own integral.
D S UPPLEMENTARY R ESULTS
D. S YSTEMATIC BENCHMARK DETAILS
[Table]
C a m e l p a r a m .
[Table] [Table]
D. C OMPARING Z Ü NIS WITH UNIFORM SAMPLING ON MATRIX ELEMENTS
[Figure]
D. E FFECT OF SURVEY STRATEGIES
In the following, we want to investigate how the integration of the three example processes in 4.
with Z Ü NIS behaves relative to VEGAS in dependence of the choice of the loss function, the survey strategy and the number of epochs during training.
For all other options, the default values are chosen again, except for the number of points during the survey phase, which is set to 500,000.
Figure 9a shows that for a simple process like e − µ → e − µ via Z , where no correlations exist, Z ÜNIS cannot reach the speed-up achieved by VEGAS.
Variance loss seems to lead to higher variance improvements than D KL loss.
Contrary, 9b shows that Z Ü NIS can greatly improve the unweighting efﬁciency for this process.
The effect is again consistently stronger when using variance loss.
Using a ﬂat survey strategy suffers for both loss functions from overﬁtting, whereas adaptive sampling in average performs slightly better and does not show overﬁtting.
d ¯ d → d ¯ d via Z presents a more realistic use case, as the parton distribution functions introduce correlations between the integration variables which present a challenge to the VEGAS algorithm.
For this process, both the speed-up and the unweighting efﬁciency ratio clearly favor the variance loss again, which outperforms in both metrics the D KL loss when multiple epochs are used, as can be seen in 10.
The richer structure of the integrand reduces the effect of overﬁtting.
Therefore, the performance increases or stays approximately constant except for the combination of D KL loss and the forward survey strategy.
For the variance loss, it becomes apparent that the ﬂat survey strategy increases much slower in performance than alternative strategies as a function of the number of epochs.
However, for combination of losses and survey strategies, VEGAS could be outperformed substantially in terms of integration speed.
[Figure] [Figure]
An opposite picture is drawn by the process uc → ucg in ﬁgure 11, for which, apart from the ﬂat survey strategy, D KL loss is in general favored both for integration speed as well as unweighting efﬁciency ratio.
The adaptive survey strategy is here giving the best results, although for a high number of epochs causes overﬁtting for the unweighting efﬁciency.
[Figure]
The take-home message of this section is, one the one hand, that the ﬂat survey strategy is in general not recommended.
Apart from this, the most important mean to improve the quality of importance sampling are testing whether, independent of the survey strategy, the loss function should be chosen differently.
E E XACT MINIMZATION OF THE NEURAL IMPORTANCE SAMPLING ESTIMATOR VARIANCE
Let us show that the optimal probability distribution for importance sampling is the function itself.
Explicitly, as discussed in section 2.1, we want to ﬁnd the probability distribution p deﬁned over some domain Ω which minimizes the variance of the estimator f
( X ) /p
( X ) , for X ∼ p
( X ) .
We showed that this amounts to solving the following optimization problem:
[Equation]
such that (cid:90) dx p ( x ) = 1 ,
[Equation]
15) which we can encode using Lagrange multipliers
[Equation]
We can now solve this problem by ﬁnding extrema with functional derivatives
[Equation]
[Equation]
which indeed is zero if p ( x ) ∝ | f ( x ) | .
Furthermore, this extremum is certainly a minimum because the loss function is positive and unbounded.
Indeed, if we separate Ω into two disjoint measurable subdomains Ω 1 and Ω 2 , and deﬁne p α ( x ) such that points are drawn uniformly over Ω 1 with probability α and uniformly over Ω 2 with probability 1 − α , then the resulting loss function would be
[Equation]
which can be made arbitrarily large by sending α to 0.
F H
F. N ORMALIZING FLOWS WITH F LOW CLASSES
Flows map batches of points and their densities.
The Z Ü NIS library implements normalizing ﬂows by specifying a general interface deﬁned as a Python abstract class: GeneralFlow .
All ﬂow models in Z Ü NIS are child classes of GeneralFlow , which itself inherits from the Pytorch nn.
Module interface.
As deﬁned in section 2.2, a normalizing ﬂow in Z Ü NIS is a bijective mapping between two d dimensional spaces, which in practice are always the unit hypercube [0 , 1] d or R d , with a tractable
Jacobian so that it can be used to map probability distributions.
To this end, the GeneralFlow interface deﬁnes normalizing ﬂows as a callable Python object which acts on batches of point drawn from a known PDF p .
A batch of points x i with their PDF values is encoded as a Pytorch Tensor object X organized as follows
[Equation]
where each X i corresponds to a points stacked with its negative log density
[Equation]
Encoding point densities by their negative logarithm makes their transformation under normalizing ﬂows additive.
Indeed if we have a mapping Q with Jacobian determinant j Q , then x ∼ p ( x ) is mapped to y
= Q ( x ) ∼ ˜ p
( y )
such that
[Equation]
Coupling Cells are ﬂows deﬁned by an element-wise transform and a mask.
All ﬂow models used in Z Ü NIS in practice are implemented as a sequence of coupling cell transformations acting on a subset of the variables.
The abstract class GeneralCouplingCell and its child InvertibleCouplingCell speciﬁes the general organization of coupling cells as needing to be instantiated with
• a dimension d • a mask deﬁned as a list of boolean specifying which
coordinates
are transformed or not • a transform that implements the mapping of the non-masked variables
In release v1.0 of Z Ü NIS two such classes are provided:
PWLinearCoupling and
PWQuadraticCoupling , which respectively implement the piecewise linear and piecewise quadratic coupling cells proposed in (Müller et al., 2018).
New coupling cells can easily be implemented, as explained in appendix F.4.
Both existing classes rely on dense neural networks for the prediction of the shape of their one-dimensional piecewise-polynomial mappings, whose parameters are set at instantiation.
Here is how one can use a piecewise-linear coupling cell for sampling points import torch from zunis.models.flows.coupling_cells.piecewise_coupling.
piecewise_quadratic import PWQuadraticCoupling d=
N_batch=10mask=[True,False] x=torch.zeros((N_batch,d+1))
# Sample the d first entries uniformly, keep 0.
for the negative log jacobian x[...,:-1].uniform_() print(x[0])#[0.3377, 0.4362, 0.]
f = PWQuadraticCoupling(d=d,mask=mask)
y = f(x)
print(y[0])
We provide further details of the use and possible parameters of ﬂows in the documentation of
Z Ü NIS: https://zunis-anonymous.github.io/zunis/ .
F. T RAINING WITH THE T RAINER CLASS
The design of the Z Ü NIS library intentionally restricts Flow models to being bijective mappings instead of being ways to evaluate and sample from PDFs so as not to restrict their applicability (see Brehmer & Cranmer (2020) for an example).
The speciﬁc application in which one uses a normalizing ﬂow, and in our case how precisely one samples from it, is intimately linked to how such a model is trained.
As a result, Z Ü NIS bundles together the low-level training tools for Flow models together with sampling tools inside the Trainer classes.
The general blueprint for such classes is deﬁned in the GenericTrainerAPI abstract class while the main implementation for users is provided as StatefulTrainer .
At instantiation, all trainers expect a Flow model and flow_prior which samples point from a ﬁxed PDF in latent space.
These two elements together deﬁne a probability distribution over the target space from which one can sample.
There are two main ways one interacts with Trainers :
• One can sample points from the PDF deﬁned by the model and the prior using the sample_forward method.
• One can train over a pre-sample batch of points, their sampling PDF and the corresponding function values using train_on_batch(self, x, px, fx)
In practice, we expect that the main way users will use Trainers is for sampling pre-trained models.
In the context of particle physics simulations for example, unweighting is a common task, which aims at sampling exactly from a known function f .
A common approach is the Hit-or-miss algorithm (James, 1980), whose efﬁciency is improved by sampling from a PDF approaching f .
This is how one would use a trainer trained on f :
[Equation]
or
[Equation]
a trainer ‘pretrained
[Equation]
import torch
[Equation]
-1]).exp()
[Table]
[Table]
F. I NTEGRATING WITH THE I NTEGRATOR CLASS
Integrators are intended as the main way for standard users to interact with Z Ü NIS.
They provide a high-level interface to the functionalities of the library and only optionally require users to know to what lower levels of abstractions really entail and to what their options correspond.
From a practical point of view, the main interface of Z Ü NIS for integration is implemented as the Integrator , which is a factory function that instantiates the appropriate integrator class based on a choice of options.
All integrators follow the same pattern, deﬁned in the SurveyRefineIntegratorAPI and
BaseIntegrator abstract classes.
Integration start by performing a survey phase, in which it optimizes the way it samples points and then a reﬁne phase, in which it computes the integral by using its learned sampler.
Each phase proceeds through a number of steps, which can be set at instantiation or when integrating:
# Setting values at instantiation time integrator =
Integrator(d=d, f=f, n_iter_survey=3, n_iter_refine=5)
# Override at integration time integral_data =
integrator.integrate(n_survey=10, n_refine=10)
For both the survey and the reﬁne phases, using multiple steps is useful to monitor the stability of the training and of the integration process: if one step is not within a few standard deviations of the next, either the sampling statistics are too low, or something is wrong.
For the reﬁne stage, this is the main real advantage of using multiple steps.
On the other hand, at each new survey step, a new batch of points is re-sampled, which can be useful to mitigate overﬁtting.
By default, only the integral estimates obtained during the reﬁne stage are combined to compute the ﬁnal integral estimate, and their combination is performed by taking their average.
Indeed, because the model is trained during the survey step, the points sampled during the reﬁne stage are correlated in an uncontrolled way with the points used during training.
Ignoring the survey stage makes all estimates used in the combination independent random variables, which permits us to build a formally correct estimator of the variance of the ﬁnal result.
F. I MPLEMENTING NEW COUPLING CELLS
To implement a new invertible coupling cell inheriting from InvertibleCouplingCell , one must provide an InvertibleTransform object and deﬁne a callable attribute T computing the parameters of the transform.
For example, consider a very simple linear coupling cell over R y
[Equation]
Q
[Equation]
x
[Equation] [Equation]
(22) where T ( x A ) is a scalar strictly positive value.
This can be deﬁned in the following way in Z Ü NIS import torch from zunis.models.flows.coupling_cells.general_coupling import
InvertibleCouplingCell from zunis.models.flows.coupling_cells.transforms import
InvertibleTransform from zunis.models.layers.trainable import ArbitraryShapeRectangularDNN class LinearTransform(InvertibleTransform): def forward(self,x,T): alpha=torch.exp(T) logj =
T * x.shape[-1] return x*alpha, logj def backward(self,x,T): alpha=torch.exp(-T)logj=-T*x.shape[-1] returnx*alpha,logj class LinearCouplingCell(InvertibleCouplingCell): def __init__(self, d, mask, nn_width, nn_depth): transform = LinearTransform() super(LinearCouplingCell, self).__init__(d=
d, mask=mask,transform= transform) d_in = sum(mask) self.
T =
ArbitraryShapeRectangularDNN(d_in=d_in,out_shape=(1,), d_hidden=nn_width,n_hidden=nn_depth)
G H ARDWARE SETUP DETAILS
The computations presented in this work were performed on a computing cluster using a Intel(R)
Xeon(R) Gold 5218 CPU @ 2.30GHz with 376 GB RAM.
Processes which could be performed on the GPU were done on a GeForce RTX 2080 having 12 GB memory and running on CUDA 11.0.