[Paper Review] 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 Applications28 references30 citations

TL;DR

CaloFlow v2 trains a fast IAF via Probability Density Distillation from a teacher MAF to emulate Geant4 calorimeter showers, achieving ~10^4× speedup with fidelity comparable to the previous CaloFlow, and matching GAN-era speed while preserving high accuracy.

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.

Motivation & Objective

• 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.

Proposed method

• 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.

Experimental results

Research questions

• 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?

Key findings

• 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.