Accelerating HEP simulations with Neural Importance Sampling

1 INTRODUCTION . 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 fights 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 significantly . 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 final 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 fields , 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 scientific 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 inefficient , 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 fields 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 flows 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-efficiency 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 difficult . 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 efficient . • 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 . 2 BACKGROUND . 2.1 IMPORTANCE SAMPLING AS AN OPTIMIZATION PROBLEM . Importance sampling relies on the interpretation of integrals as expectation values . Indeed , let us consider an integral over a finite volume : I = ∫ Ω dxf ( x ) , where V ( Ω ) = ∫ Ω dx is finite . ( 1 ) Let p be a strictly positive probability distribution over Ω , we can re-express our integral as the expectation of a random variable I = ∫ Ω p ( x ) dx f ( x ) p ( x ) = E Xi∼p 1 N N∑ i=1 f ( Xi ) p ( Xi ) , ( 2 ) 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 : σ2 ( f , p ) = E x∼p ( ( f ( x ) p ( x ) ) 2 ) − I2 . ( 3 ) The problem statement of importance sampling is to find 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.2 NORMALIZING FLOWS AND COUPLING CELLS . Normalizing flows ( 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 flows are diffeomorphisms : invertible , ( nearly-everywhere ) differentiable mappings with a differentiable inverse . Indeed , if u ∼ p ( u ) , then T ( u ) = x ∼ q ( x ) where q ( x = T ( u ) ) = p ( u ) |JT ( u ) |−1 , ( 4 ) where JT is the Jacobian determinant of T : JT ( u ) = det ∂Ti ∂uj ( u ) . ( 5 ) In the world of machine learning , T is called a normalizing flow and is typically part of a parametric family of diffeomorphisms ( T ( · , θ ) ) such that gradients∇θJT 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.3 NEURAL IMPORTANCE SAMPLING . Neural importance sampling was introduced in the context of computer graphics ( Müller et al. , 2018 ) and proposes to use normalizing flows as a family of probability distributions over which to solve the minimization problem of importance sampling . L ( θ ) = ∫ Ω dx f2 ( x ) p ( x , θ ) . ( 6 ) Of course , to actually do so , one needs to find 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 ∇θL ( θ ) = − ∫ Ω dx f2 ( x ) p ( x , θ ) ∇θ log p ( x , θ ) , ( 7 ) for which an estimator is proposed as ∇̂θL ( θ ) = − N∑ i=0 ( f ( Xi ) p ( Xi , θ ) ) 2 ∇θ log p ( Xi , θ ) , Xi ∼ p. ( 8 ) 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 DKL 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 flows provides a tractable probability distribution function from which to sample points and estimate the integral . 3 CONCEPTS 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 flexible and data-efficient way of training normalizing flows 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.1 EFFICIENT 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 defined eq . ( 8 ) works , its main drawback is that it samples points from the same distribution that it tries to optimize . As a result , new pointsXi must be sampled from p after every gradient step , which is very inefficient 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 define an auxiliary probability distribution function q ( x ) , independent from θ , from which we sample to estimate our loss function : ∫ dx f ( x ) 2 p ( x , θ ) = E x∼q f ( x ) 2 p ( x , θ ) q ( x ) . ( 9 ) 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 . ( 8 ) . 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 fig . 4 that reusing data indeed has a very significant impact on data efficiency . Algorithm 1 : Backward sampling training in ZÜNIS Data : Parametric PDF p ( x , θ ) Result : Trained PDF p ( x , θ ) 1 for M steps do 2 Sample a batch x1 , . . . , xnbatch from q ; 3 Compute the sampling PDF values q ( xi ) ; 4 Compute function values f ( xi ) ; 5 Start tracking gradients with respect to θ ; 6 for N steps do 7 Compute the PDF values from the parametric PDF p ( xi , θ ) ; 8 Estimate the loss L̂ = 1 N nbatch∑ i=1 f ( xi ) 2 p ( x , θ ) q ( x ) ; 9 Compute ∇θL̂ using backpropagation ; 10 Set θ ← θ − η∇θL̂ ; 11 Reset gradients ; 12 end 13 end 14 return p ( x , θ ) ; 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 while1 ( 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 '' ) . An important point to notice is that the original NIS formulation in eq . ( 8 ) can be restated as a limiting case of our approach . Indeed , if we take q to be a frozen version of the model p ( x , θ0 ) , which we update everytime we sample points ( setting N = 1 in algorithm 2 ) , the gradient update in line 9 is ∇θ [ E x∼p ( x , θ0 ) f ( x ) 2 p ( x , θ ) p ( x , θ0 ) ] ∣∣∣∣ θ0=θ = − E x∼p ( x , θ ) f ( x ) 2 p ( x , θ ) 2 ∇θ log p ( x , θ ) . ( 10 )

Summary: The paper describes ZüNIS, a neural importance sampling (NIS) library. The main application area is high-energy physics (HEP) simulation which is used to, for example, interpret experiments at the LHC. Traditionally, integrals arising in HEP simulations are solved numerically with the VEGAS algorithm. VEGAS is an adaptive importance sampler that neglects correlations between variables and optimizes a coordinate-wise discrete transformation of the integration domain. VEGAS potentially suffers from large variances in the Monte Carlo approximation and violations of the underlying assumptions. NIS offers a remedy in that it can adapt itself more flexibly to the integrand and thereby achieves smaller Monte Carlo errors. The authors propose a simple extension of NIS that allows them to reuse samples when the evaluation of the integrand is very complex and expensive (as is the case in HEP simulations). They show that ZüNIS outperforms the VEGAS algorithm on various toy examples and simple HEP processes.