[논문 리뷰] PC-Droid: Faster diffusion and improved quality for particle cloud generation

Matthew Leigh, Debajyoti Sengupta외 3명|arXiv (Cornell University)|2023. 07. 13.

Advanced Neural Network Applications참고 문헌 59인용 수 5

한 줄 요약

PC-Droid는 고에너지물리학에서 입자 클러스터 생성을 위한 최첨단 확산 모델을 도입하며, 개선된 EDM 노이즈 스케줄러, 고도화된 ODE 솔버, 일致성 디스틸레이션을 활용하여 PC-JeDi 대비 최대 두 배수 빠르고 Delphes 대비 세 배수 더 빠른 생성 속도를 달성하면서도 모든 제트 유형과 메트릭에서 생성 품질을 유지하거나 향상시킨다.

ABSTRACT

Building on the success of PC-JeDi we introduce PC-Droid, a substantially improved diffusion model for the generation of jet particle clouds. By leveraging a new diffusion formulation, studying more recent integration solvers, and training on all jet types simultaneously, we are able to achieve state-of-the-art performance for all types of jets across all evaluation metrics. We study the trade-off between generation speed and quality by comparing two attention based architectures, as well as the potential of consistency distillation to reduce the number of diffusion steps. Both the faster architecture and consistency models demonstrate performance surpassing many competing models, with generation time up to two orders of magnitude faster than PC-JeDi and three orders of magnitude faster than Delphes.

연구 동기 및 목표

• 고에너지물리학 시뮬레이션에서 입자 제트를 점군으로 생성하기 위한 더 빠르고 고해상도의 확산 모델을 개발하기.
• 기존의 서rogate 모델인 PC-JeDi와 Delphes와 비교해 추론 시간을 단축하면서도 생성 품질을 유지하거나 향상시키기.
• 주의 기반 아키텍처와 일치성 모델을 활용해 생성 속도와 품질 간의 상호 상충 관계를 탐색하기.
• 모든 제트 유형에 걸쳐 공동 학습을 통해 일반화 능력 향상과 조건부 없이 생성할 수 있는 능력 확보하기.
• 표준 메트릭과 실제 하드웨어 평가를 활용해 입자 클러스터 생성을 위한 새로운 기준 설정하기.

제안 방법

• 단순화된 SDE와 ODE를 사용한 EDM(노이즈 제거 스코어 기반) 확산 공식화를 채택하며, $\gamma(t) = 1$ 및 $\sigma(t) = t$를 적용해 역방향 궤적을 더 직선으로 만들고 절단 오차를 줄인다.
• 성능 향상이 가장 큰 영향을 미치는 중간 노이즈 수준에 집중하기 위해 노이즈 비율 $\sigma \sim \mathcal{N}(-1.2, 1.2)$의 로그 정규 분포 샘플링을 활용한다.
• 노이즈 수준이 낮을 때 네트워크를 건너뛸 수 있도록 $\sigma$-의존적 스케일링과 스킵 연결($c_{\text{in}}, c_{\text{out}}, c_{\text{skip}}$)을 도입해 단위 분산 유지 및 성능 안정성 확보.
• 노이즈 예측이 아닌 재구성 오차를 직접 최소화하기 위해 손실 함수를 $\mathcal{L} = \|D_\theta(\boldsymbol{x}; \sigma, \boldsymbol{y}) - \boldsymbol{x}_0\|^2_2$로 재구성.
• k-diffusion 라이브러리에서 제공하는 고급 ODE 통합 솔버를 활용하며, 일치성 모델 학습에 DPM2가 뛰어난 성능을 보였다.
• PC-Droid를 테이처로 하여 일치성 디스틸레이션 모델(CD1, CD5)을 학습시켜 한 번의 스텝으로도 수용 가능한 정밀도를 확보하고, 다섯 번째 스텝까지 성능 포화 상태에 도달함.

Figure 1 : PC-Droid network architecture and training setup.

실험 결과

연구 질문

• RQ1 다시 설계된 확산 공식화가 입자 클러스터 생성의 추론 시간을 크게 줄이고 품질을 향상시킬 수 있는가?
• RQ2 현대적 ODE 솔버와 개선된 네트워크 정규화 기법이 제트 입자 클러스터 생성의 정밀도와 속도에 어떤 영향을 미치는가?
• RQ3 일치성 모델이 표준 확산 모델과 비교해 경쟁 가능한 품질을 갖춘 한 번의 스텝 생성을 얼마나 잘 실현할 수 있는가?
• RQ4 다양한 아키텍처와 디스틸레이션 전략 간에 생성 속도와 품질 간의 성능 상충 관계는 어떠한가?
• RQ5 모든 제트 유형에 걸쳐 공동 학습을 통해 일반화 능력 향상과 성능 저하 없이 조건부 없이 생성할 수 있는 능력 향상이 가능한가?

주요 결과

• PC-Droid는 모든 제트 유형과 평가 메트릭에서 최첨단 성능을 달성하며, PC-JeDi와 Delphes를 모두 품질과 속도 면에서 능가한다.
• 가장 빠른 모델(CD1)은 EPiC-GAN의 약 세 배 시간에 최상의 제트를 생성하지만, FPND 점수는 이상적인 Delphes 기준에 한 계단 차이로 더 가까이 접근한다.
• CD5(다섯 스텝 일치성 모델)는 여덟 스텝 FPCD(PD8) 모델보다 2.5배 더 빠르며, $\mathrm{W}_1^{\mathrm{M}}$ 및 $\mathrm{W}_1^{\tau_{32}}$를 포함한 모든 메트릭에서 향상된 성능을 보였다.
• 전체 PC-Droid 모델은 EPiC-GAN보다 250배 느리지만, $\mathrm{W}_1^{\tau_{32}}$를 제외한 모든 메트릭에서 이상적인 성능에 근접한다.
• 일치성 모델은 다섯 번째 스텝 이후 성능 포화 상태에 도달함을 보여주며, CD1은 질량 스며남 현상이 있지만 $\tau_{32}$에 대해 특히 뛰어난 서브스트럭처 재현 능력을 보였다.
• 개선된 노이즈 스케줄러와 네트워크 정규화 기법은 절단 오차를 감소시키고, 더 적은 역방향 스텝 수로도 안정적인 학습을 가능하게 하여 품질을 희생시키지 않은 채 더 빠른 추론을 가능케 하였다.

Figure 2 : A single cross-attention encoder block updating both an input point cloud $x$ and global tokens $g$ , using layer-norm (LN), multi-headed attention (MHA). Converging arrows represent vector addition and contextual information $c$ is injected into the network by concatenating to the inputs