[论文解读] Accelerating HEP simulations with Neural Importance Sampling

Nicolas Deutschmann, Niklas Götz|arXiv (Cornell University)|Jan 17, 2024

Nuclear reactor physics and engineering被引用 6

一句话总结

本文中提出 ZüNIS，是一个基于 PyTorch 的神经重要性采样库，通过在梯度步骤之间重用样本以提高数据效率，并与 VEGAS 进行基准对比，在若干 HEP 类积分中显示显著的方差降低和无权重化增益，尽管训练成本较高。

ABSTRACT

Many 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 this approach suffers from poor performance in complex cases. As a result, there has been keen interest in HEP for modern machine learning to power adaptive sampling. While previous studies have shown the potential of normalizing-flow-powered neural importance sampling (NIS) over VEGAS, there remains a gap in accessible tools tailored for non-experts. In response, we introduce ZüNIS, a fully automated NIS library designed to bridge this divide, while at the same time providing the infrastructure to customise the algorithm for dealing with challenging tasks. After a general introduction on NIS, we first show how to extend the original formulation of NIS to reuse samples over multiple gradient steps while guaranteeing a stable training, yielding a significant improvement for slow functions. Next, we introduce the structure of the library, which can be used by non-experts with minimal effort and is extensivly documented, which is crucial to become a mature tool for the wider HEP public. We present systematic benchmark results on both toy and physics examples, and stress the benefit of providing different survey strategies, which allows higher performance in challenging cases. We show that ZüNIS shows high performance on a range of problems with limited fine-tuning.

研究动机与目标

• 动机：减少对在复杂 HEP 模拟中手工调优的 VEGAS 的依赖。
• 提出一种数据高效的神经重要性采样（NIS）训练方案，使样本在多次梯度更新中重用。
• 提供一个端到端、对非专家友好的基于 PyTorch 的库（ZüNIS），以便应用 NIS。

提出的方法

• 通过采用数据高效的训练方法扩展神经重要性采样，该方法使用辅助分布 q(x) 来估计损失，同时从 q 中采样并对训练中的 p(x, θ) 进行评估。
• 推导并实现一个训练算法（带缓存/前向采样），使样本在多次梯度更新中重用。
• 引入三种 q 采样策略：均匀（扁平）、冻结模型，以及 adaptive_variance，支持灵活训练。
• 提供一个基于 PyTorch 的库（ZüNIS），其模块组件包括：Flows、Trainers、Integrators。
• 将 NIS（使用 ZüNIS）与 VEGAS 和均匀采样在低维和高维积分上进行比较，包括类似 MadGraph 截面的问题。

(a)

实验结果

研究问题

• RQ1 在具有挑战性的 HEP 积分上，具有样本重用的神经重要性采样能否超越 VEGAS？
• RQ2 数据高效训练方案对 NIS 的收敛速度和数据使用有何影响？
• RQ3 ZüNIS 在不同维度和问题类型（低维函数和类似 MadGraph 的积分）上的实际性能如何？

主要发现

函数	相对于均匀采样的方差降低	相对于 VEGAS 的方差降低
Camel	1.8±0.4×10^3	7.0±1.4×10^2
Slashed Circle	8.9±0.9×10^1	8.8±0.9×10^1
Sinusoidal	2.0±0.5×10^2	1.6±0.5×10^2

• ZüNIS 在某些 2D–35D 函数（骆驼形、带缝圆、正弦形）下对均匀采样和 VEGAS 实现了最高可达 1000x 的方差降低。
• 在多个梯度步骤中使用带 q 分布的样本重用显著加速训练，在前向采样下收敛提升约 2–10 倍。
• 在某些问题上，尤其是高成本评估场景，ZüNIS 可以在积分收敛和去权效率方面优于 VEGAS，尽管总体训练和采样更慢。
• 在 MadGraph 截面测试中，ZüNIS 在不同过程上的提升具有变化；平滑函数更 favor KL 损失，而峰状函数受益于方差损失，建议使用自适应损失。
• ZüNIS 在训练和采样方面比 VEGAS 慢，但利用 GPU 加速，在函数评估代价高时提供实际收益。

(b)