[Paper Review] Set-Conditional Set Generation for Particle Physics

F. A. Di Bello, E. Dreyer+7 more|arXiv (Cornell University)|Nov 11, 2022

Big Data Technologies and Applications6 citations

TL;DR

This paper proposes a graph neural network with slot-attention (GNN+SA) for set-conditional set generation in particle physics, directly modeling reconstructed particle sets from truth particles. The model outperforms a conditional VAE baseline in capturing detector resolution effects and multi-particle correlations, with significantly lower Hungarian loss and improved MMD² metrics on test data.

ABSTRACT

The simulation of particle physics data is a fundamental but computationally intensive ingredient for physics analysis at the Large Hadron Collider, where observational set-valued data is generated conditional on a set of incoming particles. To accelerate this task, we present a novel generative model based on a graph neural network and slot-attention components, which exceeds the performance of pre-existing baselines.

Motivation & Objective

• To accelerate collider physics simulations by replacing computationally expensive detector and reconstruction steps with a learnable generative model.
• To model the conditional distribution $ p(R|T) $ directly, where $ R $ is the reconstructed particle set and $ T $ the truth particle set, bypassing hit-level simulation.
• To address limitations of prior models that fail to capture non-marginal dependencies and resolution effects at the constituent level.
• To develop a scalable, permutation-invariant architecture capable of handling variable-sized particle sets with complex spatial and momentum dependencies.

Proposed method

• The model splits the task into two stages: cardinality prediction $ q(N_R|T) $ and conditional set generation $ q(R|N_R, T) $, enabling end-to-end training.
• A graph neural network (GNN) with message passing enforces permutation invariance over truth and reconstructed particle sets, modeling spatial and momentum relationships.
• Slot-attention mechanisms are used to attend over learned slots, enabling disentangled representation of reconstructed particles with attention-based set generation.
• The model is trained on a synthetic dataset of truth events and multiple reconstructed replicas per event, using supervised learning for cardinality and conditional generation.
• A Hungarian loss is used to compute set-to-set matching cost between predictions and ground truth, enabling robust evaluation of set-level generation quality.
• Evaluation uses MMD² and visual inspection of event displays to assess fidelity to the true distribution, including resolution and correlation modeling.

Experimental results

Research questions

• RQ1 Can a deep generative model accurately approximate the complex, conditional distribution $ p(R|T) $ in particle physics reconstruction without simulating detector hits?
• RQ2 How well can a GNN+SA architecture model non-trivial dependencies between particles, such as those in dense jets, compared to standard VAE-based approaches?
• RQ3 To what extent does the model capture momentum-dependent resolution effects and particle correlations in reconstructed events?
• RQ4 Does the proposed two-stage approach (cardinality prediction + conditional generation) outperform end-to-end conditional VAE baselines in set-level fidelity?
• RQ5 Can the model generalize to representative events with high particle multiplicity and complex spatial configurations?

Key findings

• The GNN+SA model achieves a mean Hungarian loss of 1.88 on the test set, significantly lower than the cVAE model’s 2.87, indicating superior set-level matching to ground truth.
• The GNN+SA model reduces the MMD² metric to 0.0043, compared to 0.0071 for the cVAE, showing better alignment with the true distribution’s second-order statistics.
• The model successfully captures momentum-dependent resolution effects, avoiding the under-estimation and high variance seen in the cVAE at high momenta.
• Visual inspection of representative events confirms that GNN+SA samples resemble the target distribution more closely than cVAE, especially in resolving overlapping jets.
• The cVAE model fails to model jet dynamics at the constituent level, as evidenced by poor reconstruction of particle correlations and misaligned momentum distributions.
• The GNN+SA model demonstrates robustness in modeling conditional dependencies, particularly in spatially dense particle clusters, while maintaining conditional independence in disconnected regions.