[论文解读] CaloFlow II: Even Faster and Still Accurate Generation of Calorimeter Showers with Normalizing Flows

Claudius Krause, David Shih|arXiv (Cornell University)|Oct 21, 2021

Computational Physics and Python Applications参考文献 28被引用 30

一句话总结

CaloFlow v2 通过概率密度蒸馏从教师 MAF 学习快速 IAF，以模拟 Geant4 的能量沉积淋浴，获得约 10^4× 的加速，保真度与之前的 CaloFlow 相当，并且达到 GAN 时代的速度同时保持高精度。

ABSTRACT

Recently, we introduced CaloFlow, a high-fidelity generative model for GEANT4 calorimeter shower emulation based on normalizing flows. Here, we present CaloFlow v2, an improvement on our original framework that speeds up shower generation by a further factor of 500 relative to the original. The improvement is based on a technique called Probability Density Distillation, originally developed for speech synthesis in the ML literature, and which we develop further by introducing a set of powerful new loss terms. We demonstrate that CaloFlow v2 preserves the same high fidelity of the original using qualitative (average images, histograms of high level features) and quantitative (classifier metric between GEANT4 and generated samples) measures. The result is a generative model for calorimeter showers that matches the state-of-the-art in speed (a factor of $10^4$ faster than GEANT4) and greatly surpasses the previous state-of-the-art in fidelity.

研究动机与目标

• Motivate fast, accurate calorimeter shower simulations to alleviate Geant4 bottlenecks at the LHC and HL-LHC.
• Develop a fast-sampling generative model that maintains high fidelity to Geant4 shower distributions.

提出的方法

• Use Flow I (small) to model deposited energies conditioned on input energy, and Flow II (large) to model shower shapes conditioned on energies.
• Replace slow MAF-based sampling with a fast IAF via Probability Density Distillation (teacher-student training).
• Train the student IAF to match the teacher MAF through a composite loss that includes x-loss, z-loss, and additional intermediate- and parameter-level matching terms (Lx, Lz, Lx(i), Lz(i), Lkappa).
• Employ a fully-guided training objective combining x- and z-loss components to enforce stepwise agreement between flows (Eq. 17).
• Evaluate fidelity with qualitative visuals, histograms, a classifier-based metric, and timing benchmarks.
• Maintain parity with CaloFlow v1 in architecture while achieving large speedups in sampling.

实验结果

研究问题

• RQ1 Can a fast-sampling IAF be trained to match a slower, high-fidelity MAF for calorimeter shower generation?
• RQ2 Which loss terms and training strategy best align the student IAF with the teacher MAF to achieve Geant4-level fidelity?
• RQ3 Does CaloFlow v2 retain high-fidelity shower features while delivering GAN- and Geant4-comparable generation speed?
• RQ4 How does the classifier-based fidelity metric compare between Geant4, CaloFlow v1, and CaloFlow v2 across particle types?
• RQ5 What are the practical timing gains in generation for large-scale shower sampling tasks?

主要发现

• CaloFlow v2 achieves a near-saturation of teacher NLL by the student across e+, γ, and π+ showers.
• Qualitative average shower images and Flow II histograms from the student closely match Geant4 and teacher outputs with no mode collapse.
• The classifier metric shows Geant4 vs. CaloFlow v2 student remains high-fidelity and materially better than GAN-based baselines across particle types.
• Timing benchmarks show CaloFlow v2 sampling speed comparable to CaloGAN and far faster than Geant4 (up to ~10^4× faster than Geant4 in generation).
• The fully-guided training (Eq. 17) with multiple loss terms yields the best NLL performance among the tested configurations.
• The approach scales to high-dimensional calorimeter data and offers potential extension to more complex calorimeter setups.