[논문 리뷰] EPiC-ly Fast Particle Cloud Generation with Flow-Matching and Diffusion

Erik Buhmann, Cedric Ewen외 8명|arXiv (Cornell University)|2023. 09. 29.

Generative Adversarial Networks and Image Synthesis인용 수 14

한 줄 요약

이 논문은 두 가지 permutation-equivariant jet cloud 생성기인 EPiC-JeDi(score-based diffusion)와 EPiC-FM(flow-matching CNF)을 제시하여 JetNet top-jet 데이터셋에서 빠른 생성 속도와 함께 최첨단 재현도(fidelity)를 달성한다. 또한 무조건적 대 조건적 생성 분석을 수행하고 확산(diffusion)과 흐름(flow) 기반 시점을 하나의 연속 시간 프레임워크로 연결한다.

ABSTRACT

Jets at the LHC, typically consisting of a large number of highly correlated particles, are a fascinating laboratory for deep generative modeling. In this paper, we present two novel methods that generate LHC jets as point clouds efficiently and accurately. We introduce \epcjedi, which combines score-matching diffusion models with the Equivariant Point Cloud (EPiC) architecture based on the deep sets framework. This model offers a much faster alternative to previous transformer-based diffusion models without reducing the quality of the generated jets. In addition, we introduce \epcfm, the first permutation equivariant continuous normalizing flow (CNF) for particle cloud generation. This model is trained with {\it flow-matching}, a scalable and easy-to-train objective based on optimal transport that directly regresses the vector fields connecting the Gaussian noise prior to the data distribution. Our experiments demonstrate that \epcjedi and \epcfm both achieve state-of-the-art performance on the top-quark JetNet datasets whilst maintaining fast generation speed. Most notably, we find that the \epcfm model consistently outperforms all the other generative models considered here across every metric. Finally, we also introduce two new particle cloud performance metrics: the first based on the Kullback-Leibler divergence between feature distributions, the second is the negative log-posterior of a multi-model ParticleNet classifier.

연구 동기 및 목표

• Permutation-equivariant 아키텍처를 활용한 빠르고 확장 가능한 샘플링으로 제트 구름 생성을 향상시킨다.
• 확산 기반 및 흐름 기반의 연속 시간 모델을 제트 생성에 대해 단일 프레임워크로 통합할 수 있음을 보인다.
• JetNet top-jet 데이터셋에서 기존 모델과의 벤치마크를 통해 재현도와 속도 이점을 입증한다.
• 조건부 대 무조건부 생성을 비교하고 제트 레벨 조건화가 성능에 미치는 영향을 조사한다.

제안 방법

• EPiC 계층을 score-matching 확산과 결합하여 EPiC-JeDi를 통해 빠르고 정확한 제트 구름을 얻는다.
• 직선적이고 최적 운송(groent)과 유사한 궤적을 학습하는 최초의 permutation-equivariant CNF인 EPiC-FM을 도입한다.
• 두 모델 모두에 MLP 입력에 조건 정보를 연결하는 EPiC 네트워크 아키텍처를 사용한다.
• 제트를 상대적 하드로노믹 좌표로 포인트 클라우드로 표현하고 필요 시 pT 및 질량 등의 제트 레벨 관찰값으로 생성에 조건을 걸다.
• 조건 설정에서 질량, pT, 구성 요소 다수성 등의 결합 제트 속성을 생성하기 위해 조건부 정상화 흐름(NormFlows)을 사용한다.
• 샘플링 옵션(ODE 기반 결정적 샘플링과 SDE 기반 확률적 샘플링)을 논의하고 흐름 매칭을 스코어 기반 목표와 연결한다.

Figure 2 : Model schema of the EPiC Network generator architecture used in both EPiC-JeDi and EPiC-FM . Each multi-layer perceptron (MLP) is a two-layer neural network with LeakyReLU activation. The pooling operation is a concatenation of both average and summation pooling.

실험 결과

연구 질문

• RQ1 EPiC-JeDi(diffusion 기반)와 EPiC-FM(flow-matching CNF)가 빠른 샘플링 속도를 유지하면서 top-jet 생성에서 최첨단 재현도를 달성할 수 있는가?
• RQ2 조건부 대 무조건부 생성이 permutation-equivariant 제트 구름 모델의 재현도와 효율성 측면에서 어떻게 다른가?
• RQ3 플로우 매칭과 스코어 매칭 목표를 사용하는 것이 이 제트 물리 설정의 학습 안정성과 샘플링 효율성에 미치는 영향은 무엇인가?
• RQ4 제안된 모델이 서로 다른 JetNet top-jet 구성(예: 구성 요소 수 변화) 및 조건화 방식에 대해 일반화되는가?

주요 결과

• EPiC-FM은 평가된 지표 전반에서 다른 모델에 비해 일관되게 우수한 재현도를 제공합니다.
• 조건부 EPiC-FM은 테스트된 구성 중에서 가장 높은 생성 재현도를 달성합니다.
• EPiC-JeDi는 트랜스포머 기반 확산 모델보다 빠른 샘플링 속도와 경쟁력 있는 성능을 제공합니다.
• 두 모델인 EPiC-JeDi와 EPiC-FM은 top-jet 생성에서 강력한 성능을 보이며, EPiC-FM은 여러 지표에서 종종 다른 기준선보다 우수합니다.
• 연구는 KL-발산 기반 특징 분포와 다중 모델 NLP 기반 분류기 벤치마크라는 두 가지 새로운 입자 구름 지표를 도입합니다.

Figure 6 : Generation time per jet as a function of generated particle multiplicity. The generation time is measured on the same hardware using an NVIDIA ® A100 40 GB graphics card, with per-jet times calculated as the average for all test events generated. The batch size was optimised for optimal g