[Paper Review] JetFlow: Generating Jets with Conditioned and Mass Constrained Normalising Flows

Benno Käch, D. Krücker+4 more|arXiv (Cornell University)|Jan 1, 2022

Generative Adversarial Networks and Image Synthesis18 citations

TL;DR

This paper proposes JetFlow, a normalising flow model conditioned on jet mass and particle count to generate high-dimensional jet data with complex correlations. Using rational quadratic spline coupling layers, it achieves state-of-the-art performance on the JetNet dataset by conditioning on invariant mass and constraining the flow, enabling accurate modeling of jet mass distributions and stable, fast training.

ABSTRACT

Fast data generation based on Machine Learning has become a major research topic in particle physics. This is mainly because the Monte Carlo simulation approach is computationally challenging for future colliders, which will have a significantly higher luminosity. The generation of collider data is similar to point cloud generation with complex correlations between the points. In this study, the generation of jets with up to 30 constituents with Normalising Flows using Rational Quadratic Spline coupling layers is investigated. Without conditioning on the jet mass, our Normalising Flows are unable to model all correlations in data correctly, which is evident when comparing the invariant jet mass distributions between ground truth and generated data. Using the invariant mass as a condition for the coupling transformation enhances the performance on all tracked metrics. In addition, we demonstrate how to sample the original mass distribution by interpolating the empirical cumulative distribution function. Similarly, the variable number of constituents is taken care of by introducing an additional condition on the number of constituents in the jet. Furthermore, we study the usefulness of including an additional mass constraint in the loss term. On the exttt{JetNet} dataset, our model shows state-of-the-art performance combined with fast and stable training.

Motivation & Objective

• Address the computational burden of Monte Carlo simulations for future high-luminosity colliders by generating fast, realistic jet data using machine learning.
• Overcome limitations of standard normalising flows in modeling complex correlations between jet constituents.
• Improve jet generation by conditioning the flow on global features such as invariant mass and particle count.
• Ensure generated jet mass distributions match real data by incorporating mass constraints and empirical CDF interpolation.
• Demonstrate stable, fast training and high-fidelity generation on the JetNet dataset with state-of-the-art performance.

Proposed method

• Employ normalising flows with rational quadratic spline coupling layers to model high-dimensional jet data (up to 90 dimensions: 30 particles × 3 features).
• Condition the flow on the invariant jet mass and the number of particles to improve modeling of global correlations.
• Use empirical cumulative distribution function (ECDF) interpolation to sample from the true jet mass distribution during generation.
• Apply a mass constraint in the loss function to further align generated mass distributions with ground truth.
• Chain multiple coupling layers with triangular Jacobians to ensure invertibility and efficient determinant computation.
• Train using maximum likelihood via negative log-likelihood minimization, ensuring stable convergence.

Experimental results

Research questions

• RQ1 Can normalising flows conditioned on jet mass and particle count model complex correlations between jet constituents more effectively than standard flows?
• RQ2 How does conditioning on invariant mass improve the fidelity of generated jet mass distributions compared to unconditioned flows?
• RQ3 To what extent can the model generalize to jets with variable numbers of constituents (e.g., 27–30 particles) when trained on a fixed-size subset?
• RQ4 Does including a mass constraint in the loss function further improve the match between generated and real jet mass distributions?
• RQ5 How well does the model generalize under oversampling, i.e., when generating far more events than in the training set?

Key findings

• Vanilla normalising flows (VNF) fail to model correlations correctly, as shown by significant discrepancies in the invariant mass distribution between generated and ground truth data.
• Conditioning the flow on jet mass (CNF(m)) significantly improves performance, bringing the generated mass distribution into close alignment with the true distribution.
• The fully conditioned and constrained model (CCNF(m,n)) achieves state-of-the-art performance across all 6 evaluated metrics on the JetNet dataset.
• The model successfully generates jets with variable particle counts (down to 27 particles), with only minor deviations likely due to sparse training data in low-parton regimes.
• Oversampling by a factor of 100 does not introduce severe artifacts, indicating robustness and generalization beyond the training data size.
• Correlation heatmaps show no significant differences between generated and real jet data, confirming that the model preserves complex feature correlations.