[논문 리뷰] Fast Point Cloud Generation with Diffusion Models in High Energy Physics

V. M. Mikuni, Benjamin Nachman외 1명|arXiv (Cornell University)|2023. 04. 03.

Generative Adversarial Networks and Image Synthesis인용 수 4

한 줄 요약

이 논문은 고에너지 물리학 제트 시뮬레이션을 위한 고속 포인트 클러스터 디퓨전(FPCD)을 소개한다. FPCD는 높은 정밀도로 길이가 변하는 포인트 클러스터를 생성하는 디퓨전 모델이며, 점진적 정련을 통해 기준 모델 대비 추론 시간을 450배 감소시켜 단일 스텝 추론으로도 최신 기술 수준의 성능을 달성하면서도 토프 쿼크, W 보손, 글루온을 포함한 여러 제트 유형에서 정확도를 유지한다.

ABSTRACT

Many particle physics datasets like those generated at colliders are described by continuous coordinates (in contrast to grid points like in an image), respect a number of symmetries (like permutation invariance), and have a stochastic dimensionality. For this reason, standard deep generative models that produce images or at least a fixed set of features are limiting. We introduce a new neural network simulation based on a diffusion model that addresses these limitations named Fast Point Cloud Diffusion (FPCD). We show that our approach can reproduce the complex properties of hadronic jets from proton-proton collisions with competitive precision to other recently proposed models. Additionally, we use a procedure called progressive distillation to accelerate the generation time of our method, which is typically a significant challenge for diffusion models despite their state-of-the-art precision.

연구 동기 및 목표

• 고에너지 물리학 데이터 시뮬레이션에서 격자 기반 및 고정 길이의 생성 모델의 한계를 해결하기 위해, 예를 들어 하드론 제트와 같은 길이가 변하는, 순열에 불변인 포인트 클러스터를 효과적으로 처리하는 것.
• 픽셀화나 고정 크기 제약 없이 제트 데이터의 연속성, 확률적 성격, 대칭성을 유지하는 디퓨전 모델을 개발하는 것.
• 일般적으로 느린 추론을 보이는 디퓨전 모델의 추론 속도를 점진적 정련을 통해 가속화하여 대규모 데이터셋에 대한 실시간 또는 근접 실시간 시뮬레이션을 가능하게 하는 것.
• 각각의 제트 유형에 대해 별도의 모델을 훈련시키지 않고도 하나의 조건부 모델로 여러 입자 유형(예: 토프 쿼크, 글루온, W 보손)의 제트 생성을 통합하는 것.

제안 방법

• FPCD는 제트의 운동학적 특성(예: 에너지, 횡방향 운동량)을 생성하는 모듈과 제트 유형 및 운동학적 특성에 조건부로 입자 수준의 운동학적 특성(예: 4모멘타)을 생성하는 모듈으로 구성된 이중 구조의 디퓨전 모델을 사용한다.
• 모델은 로그 확률 밀도의 기울기를 학습하는 스코어 기반 디퓨전 프레임워크를 사용하여 볼록 최적화와 안정적인 훈련이 가능하며, 역전환 흐름이 필요로 하지 않는다.
• 마스킹 전략을 통해 훈련 중에 마스킹된 입력에 조건을 줌으로써 다양한 입자 수(30~150개 입자)를 가진 제트를 생성할 수 있다.
• 점진적 정련을 적용하여 전체 512단계의 디퓨전 과정을 8단계 또는 심지어 1단계로 생성하는 모델로 정련함으로써 추론 속도를 크게 향상시켰다.
• 입자 유형(예: 토프 쿼크, 글루온)에 조건을 줌으로써 하나의 통합 아키텍처로 다양한 제트 유형을 생성할 수 있도록 하였다.
• 모델은 순열 불변성을 다루고 포인트 클러스터 데이터의 복잡한 위상적 관계를 포착하기 위해 그래프 신경망을 활용한다.

실험 결과

연구 질문

• RQ1 표준 디퓨전 모델이 고에너지 물리학의 하드론 제트를 대표하는 길이가 변하는, 순열에 불변인 포인트 클러스터를 효과적으로 생성할 수 있는가?
• RQ2 표준 디퓨전 모델의 높은 계산 비용을 줄이되, 제트 시뮬레이션에서 물리적 정밀도를 손상시키지 않고서는 어떻게 할 수 있는가?
• RQ3 입자 유형에 조건을 줌으로써 하나의 통합 모델이 높은 정확도로 여러 제트 유형(예: 토프 쿼크, W 보손, 글루온)을 생성할 수 있는가?
• RQ4 디퓨전 추론 단계를 512단계에서 1단계로 줄였을 때, 점진적 정련이 생성 품질을 어느 정도 유지하는가?
• RQ5 기존의 생성 모델인 EPiC-GAN과 MP-GAN에 비해 FPCD는 제트 생성 과제에서 정밀도와 속도 측면에서 어떻게 비교되는가?

주요 결과

• FPCD는 물리적 의미 있는 메트릭에서 최신 기술 수준의 성능을 달성하였으며, 150개 입자 데이터셋에서 글루온의 입자 다중도 생성 오차는 0.157 ± 0.036, Z 보손의 경우 0.241 ± 0.090을 기록하였다.
• 단일 추론 단계로 생성하는 정련된 FPCD 모델(FPCD 1)은 높은 정밀도를 유지하여 글루온의 입자 다중도 오차는 0.157 ± 0.035, Z 보손의 경우 0.239 ± 0.090를 기록하였다.
• 150개 입자를 가진 FPCD의 생성 시간은 기준 모델의 31 × 10³ μs에서 FPCD 1의 66 μs로 감소하여 전체 모델 대비 450배 빨라졌다.
• FPCD 1은 MP-GAN(150개 입자 기준 11 μs 대비 12 μs)보다 빠르며, 단일 스텝 디퓨전 모델이 다단계 GAN을 능가하는 속도를 확보하면서도 정밀도는 유사하거나 초월함을 보였다.
• 단일 아키텍처로 토프 쿼크, W 보손, Z 보손, 글루온, 경량 쿼크 등 다섯 가지 제트 유형을 성공적으로 생성하여 각 유형별 별도의 모델 훈련이 필요 없음을 입증하였다.
• 부록 B의 히스토그램 비교 결과, 정련된 모델이 토프 쿼크 기반 제트의 모든 운동학적 분포에서 전체 시뮬레이션과 밀도적으로 유사함을 확인하였다.