[Paper Review] End-To-End Latent Variational Diffusion Models for Inverse Problems in High Energy Physics

Alexander Shmakov, Kevin Thomas Greif+4 more|arXiv (Cornell University)|May 17, 2023

Generative Adversarial Networks and Image Synthesis23 citations

TL;DR

Introduces Variational Latent Diffusion (VLD) for end-to-end unfolding of high-dimensional LHC detector data to truth-level parton distributions, showing improved distribution-level fidelity over baselines.

ABSTRACT

High-energy collisions at the Large Hadron Collider (LHC) provide valuable insights into open questions in particle physics. However, detector effects must be corrected before measurements can be compared to certain theoretical predictions or measurements from other detectors. Methods to solve this extit{inverse problem} of mapping detector observations to theoretical quantities of the underlying collision are essential parts of many physics analyses at the LHC. We investigate and compare various generative deep learning methods to approximate this inverse mapping. We introduce a novel unified architecture, termed latent variation diffusion models, which combines the latent learning of cutting-edge generative art approaches with an end-to-end variational framework. We demonstrate the effectiveness of this approach for reconstructing global distributions of theoretical kinematic quantities, as well as for ensuring the adherence of the learned posterior distributions to known physics constraints. Our unified approach achieves a distribution-free distance to the truth of over 20 times less than non-latent state-of-the-art baseline and 3 times less than traditional latent diffusion models.

Motivation & Objective

• Motivate unfolding (inverse problem) in high energy physics and the need for un-binned, high-dimensional mappings from detector-level to truth-level data.
• Propose a unified end-to-end Variational Latent Diffusion (VLD) framework that combines latent diffusion with variational autoencoders and physics-informed constraints.
• Demonstrate improved global distribution fidelity and physically consistent posteriors for semi-leptonic t tbar events compared to baselines.

Proposed method

• Introduce Variational Latent Diffusion (VLD) that unifies a conditioning encoder, a conditional or unconditional VAE, and a diffusion process into a single objective.
• Adopt continuous-time, variance-preserving diffusion with a learnable noise schedule and a denoising network to predict the original data.
• Incorporate a physics-informed consistency loss to enforce M^2 = E^2 - ||p||^2 relationships among mass, energy, and momentum.
• Explore end-to-end training variants: VLD, UC-VLD (unconditional decoder), and C-VLD (conditional encoder/decoder).
• Condition the latent space on detector observations via a permutation-invariant jet transformer encoder and a latent parton encoder/decoder.
• Evaluate with multiple distance metrics (Wasserstein, Energy, KS, KL with 64/128/256 bins) on semi-leptonic t tbar data.

Experimental results

Research questions

• RQ1 Can end-to-end variational latent diffusion improve unfolding of high-dimensional detector data to truth-level parton distributions?
• RQ2 Does joint training of conditioning encoder, VAE, and diffusion yield better distribution-level fidelity and physical consistency than separate components?
• RQ3 What is the impact of conditioning strategy (unconditional vs conditional decoders) on reconstruction quality and posterior realism?
• RQ4 How do physics-informed constraints affect reconstruction stability and consistency of derived quantities like mass, energy, and momentum?
• RQ5 How well do the proposed models scale to high-dimensional inverse problems in particle physics beyond the semi-leptonic t tbar topology?

Key findings

• The VLD models achieve the best performance across distance metrics, with UC-VLD and VLD outperforming baselines.
• Conditional decoder variants (C-VLD, CVAE) worsen reconstruction in this setup, suggesting unconditional decoders are more robust for inference data.
• Latent diffusion models (VLD/UC-VLD) outperform direct latent approaches (CINN, VDM) and end-to-end training improves over pre-trained LDM.
• The posterior samples from VLD are smoother and closer to true parton configurations than brute-force posteriors, capturing features like bimodal neutrino η distributions.
• Physics-informed consistency loss improves stability and aligns mass-energy-momentum relationships in predictions.
• Across 55 components, total distance metrics show VLD/UC-VLD achieving lower distances than baselines, indicating superior global distribution fidelity.