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 investigate sampling problems in HEP simulations and neural importance sampling (NIS) based methods. In addition, an NIS algorithm developed based on the ZÜNIS, a PyTorch library, is introduced. The performance of the proposed algorithm is demonstrated by toy examples and some numerical experiments of HEP simulations.