OmniFold: A Method to Simultaneously Unfold All Observables

Collider data must be corrected for detector effects ( “ unfolded ” ) to be compared with many theoretical calculations and measurements from other experiments. Unfolding is traditionally done for individual, binned observables without including all information relevant for characterizing the detector response. We introduce O MNI F OLD , an unfolding method that iteratively reweights a simulated dataset, using machine learning to capitalize on all available information. Our approach is unbinned, works for arbitrarily high-dimensional data, and naturally incorporates information from the full phase space. We illustrate this technique on a realistic jet substructure example from the Large Hadron Collider and compare it to standard binned unfolding methods. This new paradigm enables the simultaneous measurement of all observables, including those not yet invented at the time of the analysis.

Measuring properties of particle collisions is a central goal of particle physics experiments, such as those at the Large Hadron Collider (LHC). After correcting for detector effects, distributions of collider observables at "truth level" can be compared with semi-inclusive theoretical predictions as well as with measurements from other experiments. These comparisons are widely used to enhance our understanding of the Standard Model, tune parameters of Monte Carlo event generators, and enable precision searches for new physics. "Unfolding" is the process of obtaining these truth distributions (particle-level) from measured information recorded by a detector (detectorlevel). The unfolding process ensures that measurements are independent of the specific experimental context, allowing for comparisons across different experiments and usage with the latest theoretical tools [1], even long after the original analysis is completed. Many unfolding methods have been proposed and are currently used by experiments. See Refs. [2][3][4][5] for reviews and Refs. [6][7][8] for the most widely used unfolding algorithms.
Current unfolding methods face three key challenges. First, all of the widely used methods require the measured observables to be binned into histograms. This binning must be determined ahead of time and is often chosen manually. Second, because the measurements are binned, one can only unfold a small number of observables simultaneously. Multidifferential cross section measurements beyond two or three dimensions are simply not feasible. Finally, unfolding corrections for detector effects often do not take into account all possible auxiliary features that control the detector response. Even though the inputs to the unfolding can be calibrated, if the detector response depends on features that are not used directly in the unfolding, then the results will be suboptimal and potentially biased.
This Letter introduces OMNIFOLD, a new approach that solves all three of these unfolding challenges. Detectorlevel quantities are iteratively unfolded, using machine learning to handle phase space of any dimensionality without requiring binning. Utilizing the full phase space information mitigates the problem of auxiliary features controlling the detector response. There have been previous proposals to use machine learning methods for unfolding [9][10][11] as well as proposals to perform unfolding without binning [10][11][12][13]. These proposals, however, are untenable in high dimensions and do not reduce to standard methods in the binned case. OMNIFOLD naturally processes highdimensional features, in the spirit of previous machinelearning-based reweighting strategies [14][15][16][17][18][19], and it reduces to well-established methods [6] in the binned case. We also introduce simpler versions of the procedure, using single or multiple observables, named UNIFOLD and MULTIFOLD, respectively.
All unfolding methods require a trustable detector simulation to estimate the detector response. In the binned formulation, the folding equation can be written as m ¼ Rt, where m and t are vectors of the measured detector-level  (7) 182001-1 Published by the American Physical Society and true particle-level histograms, respectively. R is the "response matrix" In general, R is not invertible, so the unfolding problem has no unique solution, and methods attempt to achieve a useful solution in various ways. One of the most widely used methods is iterative Bayesian unfolding (IBU) [6], also known as Richardson-Lucy deconvolution [20,21]. Given a measured spectrum m i ¼ Prðmeasure iÞ and a prior spectrum t ð0Þ j ¼ Pr 0 ðtruth is jÞ, IBU proceeds iteratively according to the equation: where n is the iteration number. OMNIFOLD uses machine learning to generalize Eq. (2) to the unbinned, full phase space. A key concept for this approach is the likelihood ratio L½ðw; XÞ; ðw 0 ; X 0 ÞðxÞ ¼ p ðw;XÞ ðxÞ where p ðw;XÞ is the probability density of x estimated from empirical weights w and samples X. The function L½ðw; XÞ; ðw 0 ; X 0 ÞðxÞ can be approximated using a classifier trained to distinguish ðw; XÞ from ðw 0 ; X 0 Þ. This property has been successfully exploited using neural networks for full phase-space Monte Carlo reweighting and parameter estimation [19,[22][23][24][25][26]. Here, we use neural network classifiers to iteratively reweight the particle-and detector-level Monte Carlo weights, resulting in an unfolding procedure. The OMNIFOLD technique is illustrated in Fig. 1. Intuitively, synthetic detector-level events ("simulation") are reweighted to match experimental data ("data"), and then the reweighted synthetic events, now evaluated at particle-level ("generation"), are further reweighted to estimate the true particle-level information ("truth"). The starting point is a synthetic Monte Carlo dataset composed of pairs ðt; mÞ, where each particle-level event t is pushed through the detector simulation to obtain a detector-level event m. Particle-level events have initial weights ν 0 ðtÞ, and when t is pushed to m, these become detector-level weights ν push 0 ðmÞ ¼ ν 0 ðtÞ. OMNIFOLD iterates the following steps: 1. ω n ðmÞ¼ν push n−1 ðmÞL½ð1;DataÞ;ðν push n−1 ;Sim:ÞðmÞ, 2. ν n ðtÞ¼ ν n−1 ðtÞL½ðω pull n ;Gen:Þ;ðν n−1 ;Gen:ÞðtÞ. The first step yields new detector-level weights ω n ðmÞ, which are pulled back to particle-level weights ω pull n ðtÞ ¼ ω n ðmÞ using the same synthetic pairs ðt; mÞ. Note that ν push and ω pull are not, strictly speaking, functions because of the multivalued nature of the detector simulation. The second step ensures that ν n is a valid weighting function of the particle-level quantities. Assuming ν 0 ðtÞ ¼ 1, in the first iteration step 1 learns ω 1 ðmÞ ¼ p Data ðmÞ=p Sim ðmÞ, which is pulled back to the particle-level weights ω pull 1 ðtÞ.
Step 2 simply converts the per-instance weights ω pull 1 ðtÞ to a valid particle-level weighting function ν 1 ðtÞ. After one iteration, the new induced truth is This is a continuous version of IBU from Eq. (2), where the sum has been promoted to a full phase-space integral. In fact, OMNIFOLD (and IBU) are iterative strategies that converge to the maximum likelihood estimate of the true particle-level distribution [27][28][29][30][31], which we discuss in detail in the Supplemental Material [32]. After n iterations, the unfolded distribution is p ðnÞ unfolded ðtÞ ¼ ν n ðtÞp Gen ðtÞ: The unfolded result can be presented either as a set of generated events ftg with weights fν n ðtÞg (and 1. An illustration of OMNIFOLD, applied to a set of synthetic and natural data. As a first step, starting from prior weights ν 0 , the detector-level synthetic data ("simulation") is reweighted to match the detector-level natural data (simply "data"). These weights ω 1 are pulled back to induce weights on the particle-level synthetic data ("generation"). As a second step, the initial generation is reweighted to match the new weighted generation. The resulting weights ν 1 are pushed forward to induce a new simulation, and the process is iterated.
uncertainties) or, more compactly, as the learned weighting function ν n and instructions for sampling from p Gen . To demonstrate the versatility and power of OMNIFOLD, we perform a proof-of-concept study relevant for the LHC. Specifically, we unfold the full radiation pattern (i.e., full phase space) of jets, which are collimated sprays of particles arising from the fragmentation and hadronization of high-energy quarks and gluons. Jets are an ideal environment in which to benchmark unfolding techniques, since detector effects often account for a significant portion of the experimental measurement uncertainties for many jet substructure observables [33]. With the radiation pattern unfolded, one can obtain the unfolded distribution of any observable using Eq. (5). Hence, this procedure can be viewed as simultaneously unfolding all observables.
Our study is based on proton-proton collisions generated at ffiffi ffi s p ¼ 14 TeV with the default tune of HERWIG 7.1.5 [34][35][36] and Tune 26 [37] of PYTHIA 8.243 [38][39][40] in order to study a challenging setting where the "natural" and "synthetic" distributions are substantially different. As a proxy for detector effects and a full detector simulation, we use the DELPHES 3.4.2 [41] fast simulation of the CMS detector, which uses particle flow reconstruction. Jets with radius parameter R ¼ 0.4 are clustered using either all particle flow objects (detector-level) or stable non-neutrino truth particles (particle-level) with the anti-k T algorithm [42] implemented in FastJet 3.3.2 [43,44]. One of the simulations (HERWIG) plays the role of data and truth, while the other (PYTHIA) is used to derive the unfolding corrections. To reduce acceptance effects, the leading jets are studied in events with a Z boson with transverse momentum p Z T > 200 GeV. After applying the selections, we obtain approximately 1.6 × 10 6 events from each generator.
Any suitable machine learning architecture can be used for OMNIFOLD. For this study, we use particle flow networks (PFNs) [45,46] to process jets in their natural representation as sets of particles. Intuitively, PFNs learn and processes a set of additive observables via PFNðfp i g M i¼1 Þ ¼ Fð P M i¼1 Φðp i ÞÞ for an event with M particles p i , where F and Φ are parametrized by fully connected networks. We specify the particles by their transverse momentum p T , rapidity y, azimuthal angle ϕ, and particle identification code [47], restricted to the experimentally accessible information (PFN-Ex [45]) at detector level. To define separate models for step 1 and step 2, we use the PFN architecture and training parameters of Ref. [45] with latent space dimension l ¼ 256, implemented in the ENERGYFLOW PYTHON package [48]. Neural networks are trained with Keras [49] and TensorFlow [50] using the Adam [51] optimization algorithm. The models are randomly initialized in the first iteration and subsequently warm started using the model from the previous iteration. 20% of the events are reserved as a validation set during training.
The unfolding performance of OMNIFOLD is shown in Fig. 2 and compared to IBU, both with n ¼ 5 iterations. We found little difference between n ¼ 3 and n ¼ 5, though OMNIFOLD exhibits a slight preference for more iterations. OMNIFOLD succeeds in simultaneously unfolding all of these observables, achieving performance comparable to or better than IBU applied to each observable individually. The mass is challenging for all methods as particle-type information is relevant at the particle level but is not fully known at the detector level, introducing additional prior dependence. Though OMNIFOLD is unbinned, the data are only able to constrain energy and angular scales comparable to the detector resolution.
Statistical uncertainties from the prior distribution are shown in the bottom panels of Fig. 2, holding the unfolding procedure (i.e., response matrix and reweighting) fixed. For this proof-of-concept study, we do not show systematic uncertainties, though the procedure for deriving them is the same as for IBU. Nonclosure and modeling uncertainties can be derived in the standard way by testing the procedure on different Monte Carlo samples and comparing the results to the known truth distributions. (We checked that OMNIFOLD satisfies technical closure when PYTHIA is unfolded to itself.) Experimental systematic uncertainties can be obtained by varying the relevant effects and repeating the unfolding procedure. Like other unfolding procedures, OMNIFOLD cannot improve the results in phase-space regions that are unconstrained by observed quantities. It can, however, improve the performance if the full phase space contains auxiliary features relevant for the detector response. To capitalize on this full phase-space approach, it is essential that the detector simulation properly describes these features and that systematic uncertainties are estimated using a highdimensional approach [61,62].
To highlight the flexibility of our unfolding framework, we study variations of OMNIFOLD, where the available information is varied by controlling the inputs: (i) UNIFOLD: A single observable as input. This is an unbinned version of IBU. (ii) MULTIFOLD: Many observables as input. Here, we use the six jet substructure observables in Fig. 2 to derive the detector response. (iii) OMNIFOLD: The full event (or jet) as input, using the full phase space information. The unfolding performance of each method on our six substructure observables is tabulated in Table. I and compared to IBU. The UNIFOLD and MULTIFOLD implementations both use dense PHYSICAL REVIEW LETTERS 124, 182001 (2020) 182001-3 networks with three layers of one hundred nodes each and a two-node output layer. We see good unfolding performance across all methods, and even though OMNIFOLD is not directly trained on these six observables, it performs comparably to or better than MULTIFOLD. While the detector response depends on the jet rapidity, we checked that MULTIFOLD did not significantly benefit from including the rapidity, though doing so could be important in a real experimental context. In general, additional information can be included and the unfolding procedure can be repeated, with the final model chosen as the one with the best detector-level agreement with the data.
Since OMNIFOLD unfolds the full radiation pattern, it can be used to probe new, physically interesting quantities that are challenging to unfold with existing methods. One example is the recently proposed correlation (fractal) dimension of jets [66,67], which is a function of the energy scale Q. This complicated statistic is defined by pairwise metric distances between jet radiation patterns, falling outside of the purview of single-observable unfolding techniques. Within our jet samples, we restrict to energetic jets with p jet T > 500 GeV, boosted to the origin of the rapidity-azimuth plane, and with constituents rescaled to have p T summing to 500 GeV. The correlation dimensions of these jets, both before and after applying OMNIFOLD, are shown in Fig. 3. The unfolded results match the true distribution over a wide range of Q values, with residual prior dependence seen at low Q (i.e., the infrared) where jets have a higher dimensionality and detector effects have a larger impact, thus making the unfolding problem more difficult. More broadly, OMNIFOLD opens the door to going beyond per-event collider observables towards more nuanced or intricate measurements of the data.
In conclusion, we have presented a potentially transformative unfolding paradigm based on iteratively reweighting a set of simulated events with machine learning. Our OMNIFOLD approach allows an entire dataset to be unfolded using all of the available information, avoiding the need for binning and restricting to single observables. We have demonstrated the power of this method in a (simulated) case of interest by unfolding the full radiation pattern of jets, paving the way for significant advances in jet substructure at the LHC. Our unfolding framework allowed us to go beyond per-event observables towards unfolding more complex dataset statistics, such as fractal dimensions of the space of jets. Going even further, (unsupervised) machine learning models may be trained directly at the particle level by using the unfolded and weighted dataset, which is a fascinating avenue for further exploration. These advances have broad applicability beyond particle physics in domains where deconvolution or unfolding is used, such as image-based measurements and quantum computation [68]. To enable future unfolding studies and developments, we have made our code and jet datasets publicly available [69,70], including two additional tunes of PYTHIA beyond those presented here. Finally, our reweighting-based unfolding strategy allows for new observables to be measured long after the unfolding is carried out, which can significantly empower future public and archival collider data analyses [71]. FIG. 3. The correlation dimension of the space of jets, unfolded with OMNIFOLD. The unfolded results closely match the truthlevel dimension over most of the energy range, tending toward the prior in the more difficult phase space region at low Q. Unfolding a complicated statistic such as the correlation dimension is challenging with standard methods.