[논문 리뷰] PC-JeDi: Diffusion for Particle Cloud Generation in High Energy Physics

Matthew Leigh, D. Sengupta외 4명|arXiv (Cornell University)|2023. 03. 09.

Particle physics theoretical and experimental studies참고 문헌 86인용 수 5

한 줄 요약

PC-JeDi는 고에너지물리학에서 고해상도 입자군집 제트를 생성하기 위해 스코어 기반 학습과 트랜스포머를 활용한 확산 기반 생성 모델을 소개한다. GAN보다 느린 추론 속도를 지닌다 하더라도, 제트의 하위구조 충실도와 조건부 생성에서 최신 기술 수준의 성능을 달성하였으며, 전통적인 몬테카를로 방법 대비 10배 빠른 속도로 복잡한 제트 시뮬레이션을 위한 유연하고 순열에 대한 등장성 보장 구조를 제공한다.

ABSTRACT

In this paper, we present a new method to efficiently generate jets in High Energy Physics called PC-JeDi. This method utilises score-based diffusion models in conjunction with transformers which are well suited to the task of generating jets as particle clouds due to their permutation equivariance. PC-JeDi achieves competitive performance with current state-of-the-art methods across several metrics that evaluate the quality of the generated jets. Although slower than other models, due to the large number of forward passes required by diffusion models, it is still substantially faster than traditional detailed simulation. Furthermore, PC-JeDi uses conditional generation to produce jets with a desired mass and transverse momentum for two different particles, top quarks and gluons.

연구 동기 및 목표

• 기존 몬테카를로 방법 대비 계산 비용을 줄일 수 있는 고속이고 정확한 제트 시뮬레이션 대체 모델을 개발하는 것.
• 제트 이벤트에서 순서 없는 입자군집 데이터를 다룰 수 있도록 설계된 순열에 대한 등장성 보장 아키텍처를 활용하여 현실적인 제트 하위구조를 생성하는 도전 과제를 해결하는 것.
• 특정 물리적 성질을 가진 제트, 예를 들어 원하는 운동량과 질량을 가진 톰 쿼크 제트나 글루온 제트를 조건부로 생성할 수 있도록 하는 것.
• 이전에 생성 모델에서 보고되지 않은 하위구조 관측량을 포함한 종합적인 지표를 사용하여 모델 성능을 평가하는 것.
• 확산 기반 입자군집 생성에서 해상도와 속도 간의 상충 관계를 탐색하고 향후 최적화를 위한 잠재력을 탐색하는 것.

제안 방법

• PC-JeDi는 표준 정규 분포 사전에서 시작하여 노이즈가 첨가된 궤적을 역행하는 스코어 기반 확산 과정을 사용하며, 단계별로 입자 운동량을 노이즈 제거하는 방식으로 학습한다.
• 모델은 순서 없는 입자군집 데이터를 다루는 데 자연스럽게 적합한 순열에 대한 등장성 보장 특성을 지닌 트랜스포머 기반 아키텍처를 사용한다.
• 조건부 생성은 목표 제트 질량과 횡운동량에 조건을 붙여 확산 과정을 조절함으로써 특정 유형의 제트를 제어적으로 샘플링할 수 있도록 한다.
• 노이즈 제거 과정은 적응형 단계 크기를 가진 ODE/SDE 해법기를 사용하며, 데이터 분포 상에서 스코어 매칭을 통해 추정된 스코어 함수를 활용한다.
• 학습은 시간 단계의 시퀀스 동안 노이즈 제거 손실을 최소화하는 방식으로 이루어지며, 네트워크는 각 단계에서 스코어를 예측하여 역행 과정을 이끄는 데 기여한다.
• 추론은 200단계에 걸쳐 반복적으로 노이즈 제거를 수행하며, 초기에 무작위 노이즈에서 시작하여 입자 운동량을 점차 정교화하여 현실적인 제트 구성으로 도달한다.

Figure 2 : Diagram of the PC-JeDi generation procedure. Random noise is sampled at the beginning of the loop from a standard normal distribution. Then any chosen integration sampler is iteratively applied in order to fully denoise the input towards actual data, using the conditional model as a noise

실험 결과

연구 질문

• RQ1 트랜스포머 아키텍처를 갖춘 확산 기반 생성 모델은 실제 제트의 통계적 성질을 잘 따르는 고해상도 입자군집 제트를 생성할 수 있는가?
• RQ2 PC-JeDi는 하위구조 관측량과 충실도 지표 측면에서 최신 기술 수준의 MPGAN 모델 대비 어떤 성능을 보이는가?
• RQ3 PC-JeDi는 특정 물리적 성질을 가진 제트, 예를 들어 원하는 질량과 횡운동량를 가진 톰 쿼크 제트나 글루온 제트를 얼마나 잘 생성할 수 있는가?
• RQ4 PC-JeDi의 계산 비용은 GAN 기반 모델과 기존 몬테카를로 시뮬레이션 대비 어떻게 되는가?
• RQ5 PC-JeDi는 비포함된 톰 제트나 이차원 마진 분포와 같은 복잡한 제트 하위구조 상관관계를 얼마나 잘 포착할 수 있는가?

주요 결과

• PC-JeDi는 제트 하위구조 충실도와 분포 유사도를 측정하는 여러 지표에서 최신 기술 수준의 MPGAN 모델과 경쟁 가능한 성능을 달성한다.
• 특히 N-서브젯니스 및 평면 흐름과 같은 하위구조 관측량의 분포를 포착하는 데서 MPGAN을 뛰어넘는 성능을 보였다.
• 지정된 횡운동량와 질량을 가진 조건부 제트를 성공적으로 생성하여 출력의 물리적 성질에 대해 정확한 제어를 가능하게 했다.
• 샘플당 약 200회의 순방향 전파가 필요한 것으로 나타났지만, 전통적인 몬테카를로 시뮬레이션 대비 약 10배 빠른 속도로 제트 생성이 가능했다.
• GPU에서 1000개의 제트 생성에 대해 PC-JeDi의 추론 시간은 약 500ms이며, 몬테카를로의 약 4721ms에 비해 뚜렷한 속도 향상을 보였다.
• 트랜스포머 기반 아키텍처는 배치 처리에서 더 나은 확장성을 보였으며, 배치 크기 증가에 따라 추론 시간이 선형적으로 증가하는 반면, 일부 GAN에서는 비선형적 증가를 보였다.

Figure 3 : The relative transverse momentum (left) and invariant mass (right) of gluon jets generated with MPGAN (orange) and PC-JeDi (DDIM solver, green; EM solver, red) compared to the MC simulation (shaded blue). Calculated from the leading 30 $p_{\mathrm{T}}$ constituents using the constituent $