[논문 리뷰] Accelerating HEP simulations with Neural Importance Sampling

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

Nuclear reactor physics and engineering인용 수 6

한 줄 요약

이 논문은 PyTorch 기반의 신경 중요도 샘플링 라이브러리 ZüNIS를 소개하며, 샘플을 gradient 단계들 간에 재사용하여 데이터 효율성을 개선하고 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에 대한 의존도를 줄이고자 한다.
• 샘플을 gradient 단계들 간에 재사용하는 데이터 효율적인 NIS(Neural Importance Sampling) 학습 체계를 제안한다.
• 비전문가가 NIS를 적용할 수 있도록 엔드-투-엔드의 사용자 친화적인 PyTorch 기반 라이브러리(ZüNIS)를 제공한다.

제안 방법

• 보조 분포 q(x)로 손실을 추정하고 학습을 위해 q에서 샘플링하며 p(x, θ)를 평가하는 데이터 효율적 학습 방식으로 Neural Importance Sampling을 확장한다.
• 여러 gradient 스텝에 걸쳐 샘플을 재사용하는 버퍼링/포워드 샘플링 알고리즘을 도출하고 구현한다.
• uniform(평탄한), frozen model, adaptive_variance의 세 가지 q-샘플링 전략을 도입하여 유연한 학습을 가능하게 한다.
• 모듈형 컴포넌트로 Flows, Trainers, Integrators를 갖춘 PyTorch 기반 라이브러리(ZüNIS)를 제공한다.
• NIS (ZüNIS 포함)를 저차원 및 고차원 적분에서 VEGAS 및 균일 샘플링과 비교하며 MadGraph 교차섹션 유사 문제를 포함한다.

(a)

실험 결과

연구 질문

• RQ1 Neural Importance Sampling with sample reuse가 도전적인 HEP 적분에서 VEGAS를 능가할 수 있는가?
• RQ2 데이터 효율적 학습 체계가 NIS의 수렴 속도와 데이터 사용량에 어떤 영향을 미치는가?
• RQ3 다양한 차원수와 문제 유형(저차원 함수 및 MadGraph 유사 적분)에서 ZüNIS의 실용적 성능은 어떠한가?

주요 결과

함수	일반 샘플링 대비 분산 감소	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 함수들(Camel, Slashed Circle, Sinusoidal)에서 uniform sampling과 VEGAS에 비해 최대 1000x의 분산 감소를 달성한다.
• q 분포를 사용해 여러 gradient 스텝에 걸쳐 샘플을 재사용하면 학습 속도가 크게 빨라지며 forward sampling에서 수렴이 2–10x 개선된다.
• 일부 문제에서 적분 수렴 및 언웨팅 효율에서 VEGAS를 능가할 수 있지만 학습 및 샘플링은 전반적으로 느리다.
• MadGraph 교차섹션 테스트에서 ZüNIS의 개선은 프로세스에 따라 다르며, 더 매끄러운 함수는 KL 손실에 유리하고, 피크가 있는 함수는 분산 손실이 이익을 준다, 적응형 손실 선택이 바람직하다고 시사한다.
• ZüNIS는 VEGAS보다 학습 및 샘플링 속도가 느리지만, GPU 가속을 활용해 함수 평가가 비용이 큰 경우 실용적 이점을 제공한다.

(b)