Decorrelation with Conditional Normalizing Flows

Samuel Klein
University of Geneva
samuel.klein@unige.ch

Tobias Golling
University of Geneva
tobias.golling@unige.ch

Abstract
The sensitivity of many physics analyses can be enhanced by constructing discriminants that preferentially select signal events. Such discriminants become much
more useful if they are uncorrelated with a set of protected attributes. In this paper
we show that a normalizing flow conditioned on the protected attributes can be used
to find a decorrelated representation for any discriminant. As a normalizing flow
is invertible the separation power of the resulting discriminant will be unchanged
at any fixed value of the protected attributes. We demonstrate the efficacy of ourapproach by building supervised jet taggers that produce almost no sculpting in themass distribution of the background.

1

Introduction

To study particular signal processes many analyses in high energy physics (HEP) will generate bothsignal and background samples using a simulator and define a discriminant for preferentially selecting
signal [1–6]. Once a discriminant has been constructed it is used in an event selection to create asignal enriched data sample. In a HEP analysis this selection will only ever be one step in a chainof procedures. A common downstream task is the study of the invariant mass spectrum, which is ofinterest because it typically falls smoothly for the background but not the signal. The discriminant
is invariably correlated with the invariant mass and this typically means the smooth backgroundassumption is violated in the signal enriched sample. Therefore the discriminant is required to satisfy
smoothness criteria with respect to the changes it induces in the mass distribution of the background.
In this work we introduce a method that removes the mass correlation from a discriminant after it
has been constructed and leaves its separation power unchanged at every mass value. This method
can be used to decorrelate any number of variables from any number of protected attributes and iscomplementary to other decorrelation methods. We study supervised machine learning models on
a supervised task and demonstrate the proposed method significantly improves the decorrelation
performance of existing approaches.

2

MethodThe method introduced here will be described as a procedure for decorrelating discriminants, but it
can be applied to decorrelate any set of variables from any set of protected attributes, though this may
come with a cost (App. A). When combined with a machine learning model the methods we present
can be seen as additional layers that calibrate the output of the model to enforce uniform backgroundrejection as a function of the protected attributes.
Existing approaches Existing methods often perform decorrelation during the construction of thediscriminant. This can be done by minimizing a measure of the correlation between the discriminants and the protected attributes m while training a machine learning model [7–10]. The constructionof such a discriminant is performed by minimizing the expectation over some objective function of
Machine Learning and the Physical Sciences workshop, NeurIPS 2022.

the form
L = Lclass (sφ (x), y) + αLdecor (sφ (x), m),(1)
where φ parameterizes the discriminant and y is the class label for the sample described by features x.
Minimizing the loss function Lclass maximizes the utility of the discriminant while minimizing Ldecor
reduces the correlation between m and sφ (x). The parameter α controls the relative contributionof each of the loss functions and is the minimal additional parameter that can appear in such loss
functions, though more parameters can appear in Ldecor .
Decorrelation methods that use Eq. 1 require additional work when building the discriminant. Bothmoment decomposition (MoDe) [8] and distance correlation (DisCo) [7] methods have some small
overhead in terms of training, but only introduce the single hyper parameter α into the optimization.
All of these methods can be composed with the decorrelation method we propose. Other approaches
are detailed in App. B.
Decorrelation with conditional normalizing flows A normalizing flow [11] is an invertible map
between two distributions – a data distribution from which we have samples and a base distribution
for which we know the likelihood. A conditional normalizing flow pθ can approximate a data
distribution pD (s(x)|m) by defining a neural network fθ (s(x), m) that is invertible given m and abase distribution p that is independent of m. The model is fit to data by maximizing the log-likelihood
under the change of variables formula,log pθ (s(x)|m) = log p(fθ (s(x), m)) + log |det Jfθ (s(x), m)| ,(2)

where Jfθ is the Jacobian of fθ . If this function is fit perfectly the distribution of fθ (s(x), m) willbe the same as p and therefore independent of the mass by definition. In this case fθ (s(x), m) will
be a decorrelated representation of s(x) with the same signal separation power as s(x) at every valueof m and therefore fθ (s(x), m) can be used in place of s(x).
In the case of discriminants the function fθ is required to be monotonically increasing to ensure
the ordering remains unchanged at all values of m. In one dimension this condition can always be
explicitly enforced post training of a flow using reflections of fθ about the center of p as discussed inApp. C. As fθ is an invertible transformation with positive slope, the discriminatory power of s andfθ will be the same for every value of m.
Decorrelation with quantile regression Quantile regression [12] is the task of finding the quantilesof a given distribution as a function of explanatory variables. Typically we know what quantiles wewant to use for a given discriminant before it is constructed [13], and by learning these quantiles
as a function of m we can use the conditional quantiles to perform event selections such that the
distribution of m does not change. Regressing quantiles is easier than training a conditional flowwhich has to learn the entire conditional cumulative distribution function. The definition of quantiles
was recently generalised to greater than one dimension [14, 15] and for decorrelating discriminants
in more than one dimension this is the formulation that should be used to preserve the ordering thatmakes the discriminant useful. This is discussed further in App. D.

3

Experiments

In this section we will demonstrate the efficacy of the proposed method on a supervised classification
problem. The task is to separate multijet background events from hadronically decaying boosted
W bosons. The samples we use were provided by Refs [7, 16] to emulate those used in an earlier
ATLAS study on mass decorrelation techniques [17]. The same set of ten jet substructure variables as
used in these studies were used as features to train the classifier. Further details on the dataset can befound in the Refs. [7, 8] and in App. E. For this task the invariant mass is the protected attribute m.
Model settings For MoDe and DisCo models the discriminant is parameterized by a DNN withthree hidden layers with 64 nodes in each hidden layer and S WISH activations [18] with a batchnormalization layer after the first fully connected layer and a sigmoid activation on the output. This is
the same model as used in the experiments of MoDe [8]. We also consider a more standard classifier
with the same structure but with R E LU activations [19] and no batch normalization, this will bereferred to as a vDNN. All conditional normalizing flows are constructed with rational quadratic
splines [20] and trained on the background only. Further training details are provided in App. F.
2

AUC

0.6
0.4
0.2

10

103
102
101

100

vDNN
cf-vDNN

0.6

4

Signal efficiency

Background rejection

0.8

200
Mass [GeV]

100

0.4
0.2
0.0

200

100

Mass [GeV]200
Mass [GeV]

Figure 1: The AUC, background rejection and signal efficiency in bins of mass for a vDNN and thediscriminant constructed by a conditional normalizing flow using the output of the vDNN (cf-vDNN).
The binned rejection and efficiency are calculated using a threshold on the mass inclusive discriminantthat gives 50% signal efficiency. In the case of the vDNN some high mass bins do not have enoughstatistics to calculate background rejection and signal efficiency at an inclusive 50% signal efficiency
threshold.
0.150

10−3
10−4
10−5

0.125
0.100
0.075
0.050

50

100

150
200
Mass [GeV]

250

300

0.000

0.100
0.075
0.050
0.025

0.025

10−6

0<m
m ≤ 7070 < m ≤ 100
100 < m0.125
Normalised counts

Normalized counts

10−2

0<m
m ≤ 7070 < m ≤ 100
100 < m0.150
Normalised counts

Inclusive
BackgroundvDNN
cf-vDNN

0.0

0.2

0.4
0.6
vDNN

0.8

1.0

0.000

0.0

0.2

0.4
0.6
cf-vDNN0.8

1.0

Figure 2: The invariant mass distribution for the background samples without applying any selection
and a threshold rejecting 50% of the signal for different discriminants applied to the background (left)
and the distribution of the discriminant over QCD (background) samples for a vDNN (middle) and aconditional flow trained on the vDNN output (cf-vDNN, right) in bins of the mass.

We stress here that the training of the MoDe and DisCo models is different to the trainings in therespective papers and so the performance of the classifiers also differs. In particular the classifiers we
use are trained with simple standard schemes. However, the decorrelation method used here can beused with any training pipeline, in particular on top of the output of publicly available code such asthat provided by DisCo in Ref. [7] (App. G).
Signal separation A conditional flow decorrelated vDNN model (cf-vDNN) has the same separation power in bins of mass and modulates the background and signal mass distributions less than avDNN with no decorrelation as shown in Fig. 1. The difference in background rejection for mass
inclusive thresholds means that the signal separation using these thresholds is different. This isreflected in the mass inclusive AUC which is 0.90 for a vDNN and 0.84 for the cf-vDNN. This isdiscussed further in App. H.
Decorrelation The output distribution of a vDNN classifier trained without decorrelation changes
significantly in different mass bins as shown in Fig. 2. In contrast the cf-vDNN model has approximately the same distribution in all mass bins. Using the vDNN output directly results in preferentiallyselecting samples in a certain mass range, while the cf-vDNN output exhibits no such preference.
To compare the conditional flow decorrelation to existing methods we study the inverse of the JensenShannon divergence 1/JSD50 and background rejection power (inverse false positive rate) R50 at
50% signal efficiency. The JSD is computed between samples that pass and fail the selection. These
metrics are correlated in the regime of finite statistics with an optimal trade off that can be estimatedfrom data by applying random selections with the same proportions defined by R50 , we will refer tothis optimal tradeoff as ideal.
3

IdealDisCo
cf-DisCo
MoDe [0]1041/JSD70−10050

1/JSD50

10

4

103
102

103

cf-MoDe [0]DNN (cf-inputs)
cf-DNN (cf-inputs)

102

vDNN
cf-vDNN

101
5

10

20

30

5

10
70−100R50

R50

Figure 3: The inclusive background rejection at 50% signal efficiency (R50 ) and the inverse of theJensen-Shannon divergence between the mass distribution of events that do and do not pass the cut at
the same threshold (1/JSD50 ) in the left plot. The background rejection in the mass window [70−100]
70−100
GeV at 50% inclusive signal efficiency (R50
) and the inverse Jensen-Shannon divergence in
70−100the same mass bin (1/JSD50). The filled circles denote a discriminant constructed using a
conditional normalizing flow. The shaded bands capture the maximum and minimum across ten
different trainings. Models with the cf prefix have been decorrelated with a conditional normalizingflow.

In comparison to other explicit decorrelation methods, using a conditional normalizing flow performs
significantly better in terms of the trade off between 1/JSD50 and R50 as shown in Fig. 3. The
conditional flow decorrelator applied to a model trained with MoDe [8] and DisCo [7] decorrelation
approaches the ideal limit. While the inclusive background rejection power of the conditional flow
decorrelated classifiers has decreased, Fig. 3 shows that the background rejection in the signal region1
has increased. This occurs because correlated classifiers are biased towards accepting background in
the signal region. With reduced inclusive background rejection power the background estimation
based on sideband fits is expected to improve for conditional flow decorrelated models. In combination
with improved background rejection in the signal region, conditional flow decorrelated classifiers
should significantly improve the performance of resonant searches.
A conditional normalizing flow can also be used to directly decorrelate the input features to finda new representation on which a DNN classifier (DNN cf-inputs) can be trained (App. I). As thisdecorrelation is performed on all of the input variables directly some correlation with the mass
remains in the classifier output, this results in improved performance of training another conditional
normalizing flow on the resulting discriminant as seen in Fig.3.
The Disco and MoDe classifiers are trained with α = 0, 25, 50, ..., 300, with each point in Fig. 3
indexing an α which increases from right to left. The simpler vDNN classifier is easier to decorrelate
than the classifier used with DisCo and Mode when α = 0. The improved performance of the DisCoand MoDe models as α increases demonstrates the utility of combining conditional normalizing flow
decorrelation with explicit decorrelation during training. Without this decorrelation during training it
is possible that the classifier discriminant becomes overly complex and the conditional flow would
need significant resources to decorrelate the discriminant.
Searching the space of architectures for a point that is easy to decorrelate is much more complicatedthan choosing the single parameter α to use in MoDe [8] and DisCo [7]. The conditional flow clearly
improves upon the performance of these models while reducing the need to fine tune the value of α.
It is also possible that further development in training conditional normalizing flows will remove theneed for explicit decorrelation during the training of the classifier.
Quantile regression cannot be added to Fig. 3 because the quantiles are trained on the background
only and so no threshold can be found for 50% signal efficiency. However, an additional benefit ofboth conditional normalizing flows and quantile regression is their improved behaviour at higher
levels of background rejection as shown in Fig. 4. Here, even though the quantile regression does
not accurately predict the quantile thresholds, the resulting discriminant is still decorrelated fromthe mass. Due to the discontinuous nature of the quantiles, and the poor quantile predictions, it is1

Defined to be centered on the signal peak and include 80% of the signal.

4

0.8

0.6
1/JSD

Signal efficiency

104

0.4

0.2

103
DisCo
qr-DisCo
cf-DisCo
Ideal

102
50

60
70
80
90
Background rejection (%)

100

50

60
70
80
90
Background rejection (%)

100

Figure 4: The performance of a DisCo trained classifier with α = 200 and both quantile regression(qr-DisCo) and a flow decorrelated output at 50, 90, 95, 99% background rejection. Error bars capture
the maximum and minimum across ten different runs on the same DisCo trained classifier.
hard to consistently quantify the decorrelation performance of quantile regression across runs. The
JSD is calculated between the mass distribution over the background after applying the classifier
threshold and a random sample from the background with the same number of points as passed theclassifier threshold. An ideal upper bound on the JSD is calculated by randomly sampling from themass distribution without replacement.

4

Conclusion

In this paper we presented a novel method for performing decorrelation using conditional normalizing
flows. The approaches we developed were benchmarked against other state of the art methods in asupervised setting and can be seen to perform well in conjunction with these approaches where theconditional flow reduces the hyper parameter optimization burden while improving the decorrelation
performance. The approach we develop can also be applied to unsupervised settings and it was shown
to be successful for decorrelating variables other than discriminants by building mass decorrelated
feature representations (cf-inputs). The approach we develop could be particularly useful in anomaly
detection settings which rely on decorrelation [21, 22]. The approach we develop can also be applied
outside of the high energy physics domain to data anonymity and fairness.
While the method we propose can be seen to be effective here, there is no explicit smoothness
enforced on the distribution of the mass after applying a conditional normalizing flow. There isevidence that conditional normalizing flows do implicitly impose some kind of smoothness criteria [23,24], but explicitly enforcing this could significantly improve the performance of conditional flow
decorrelation and further work in this direction would be worthwhile. Conditional normalizing flows
for decorrelation is a promising direction for allowing state of the art machine learning methods to be
applied to high energy physics analyses.
Both quantile regression and conditional normalizing flows can be used to decorrelate discriminants,
with quantile regression providing a fast alternative that still significantly reduces the amount of mass
sculpting, even in the situation where it does not learn the correct quantiles. Both of these approaches
can in theory perform perfect decorrelation, however this can be difficult to achieve in practice if thecorrelations between the variable of interest and the protected attributes is very strong. In practice it
appears that some level of decorrelation during the construction of the discriminant is desireable toensure the conditional distribution can be accurately modelled by the flow.
All code is publicly available at https://github.com/sambklein/dequantile.

5

Acknowledgements
The authors would like to acknowledge funding through the SNSF Sinergia grant called Robust Deep
Density Models for High-Energy Particle Physics and Solar Flare Analysis (RODEM) with funding
number CRSII5_193716. We would also like to thank Benjamin Nachman, David Shih and Gregor
Kasieczka for useful feedback on the manuscript and John Andrew Raine for useful discussions andfeedback on the manuscript.

6

References[1] Pierre Baldi, Peter Sadowski, and Daniel Whiteson. Searching for exotic particles in high-energy
physics with deep learning. Nature communications, 5(1):1–9, 2014.
[2] Luke de Oliveira, Michael Kagan, Lester Mackey, Benjamin Nachman, and Ariel Schwartzman.
Jet-images—deep learning edition. Journal of High Energy Physics, 2016(7):1–32, 2016.
[3] CMS collaboration. Identification of heavy-flavour jets with the cms detector in pp collisions at
13 tev. arXiv preprint arXiv:1712.07158, 2017.
[4] Anja Butter, Gregor Kasieczka, Tilman Plehn, and Michael Russell. Deep-learned top tagging
with a lorentz layer. SciPost Physics, 5(3):028, 2018.
[5] Eric A Moreno, Thong Q Nguyen, Jean-Roch Vlimant, Olmo Cerri, Harvey B Newman, Avikar
Periwal, Maria Spiropulu, Javier M Duarte, and Maurizio Pierini. Interaction networks for theidentification of boosted h→ b b decays. Physical Review D, 102(1):012010, 2020.
[6] Quark versus Gluon Jet Tagging Using Jet Images with the ATLAS Detector. Technical report,
CERN, Geneva, Jul 2017. URL http://cds.cern.ch/record/2275641.
[7] Gregor Kasieczka and David Shih. Robust jet classifiers through distance correlation. PhysicalReview Letters, 125(12), sep 2020. doi: 10.1103/physrevlett.125.122001. URL https://doi.
org/10.1103%2Fphysrevlett.125.122001.
[8] Ouail Kitouni, Benjamin Nachman, Constantin Weisser, and Mike Williams. Enhancing searches
for resonances with machine learning and moment decomposition. Journal of High Energy
Physics, 2021(4), apr 2021. doi: 10.1007/jhep04(2021)070. URL https://doi.org/10.
1007%2Fjhep04%282021%29070.
[9] Chase Shimmin, Peter Sadowski, Pierre Baldi, Edison Weik, Daniel Whiteson, Edward Goul,and Andreas Søgaard. Decorrelated jet substructure tagging using adversarial neural networks.
Physical Review D, 96(7), oct 2017. doi: 10.1103/physrevd.96.074034. URL https://doi.
org/10.1103%2Fphysrevd.96.074034.
[10] Gilles Louppe, Michael Kagan, and Kyle Cranmer. Learning to pivot with adversarial networks.
2016. doi: 10.48550/ARXIV.1611.01046. URL https://arxiv.org/abs/1611.01046.
[11] E. G. Tabak and Cristina V. Turner. A family of nonparametric density estimation algorithms.
Communications on Pure and Applied Mathematics, 66(2):145–164, 2013. doi: https://doi.org/
10.1002/cpa.21423. URL https://onlinelibrary.wiley.com/doi/abs/10.1002/cpa.
21423.
[12] Roger Koenker and Gilbert Bassett. Regression quantiles. Econometrica, 46(1):33–50, 1978.
ISSN 00129682, 14680262. URL http://www.jstor.org/stable/1913643.
√[13] ATLAS collaboration. Dijet resonance search with weak supervision using s = 13 tev pp
collisions in the atlas detector. Physical review letters, 125(13):131801, 2020.
[14] Guillaume Carlier, Victor Chernozhukov, Gwendoline De Bie, and Alfred Galichon. Vector
quantile regression and optimal transport, from theory to numerics. Empirical Economics, 62(1):35–62, aug 2020. doi: 10.1007/s00181-020-01919-y. URL https://doi.org/10.1007%
2Fs00181-020-01919-y.
[15] Victor Chernozhukov, Alfred Galichon, Marc Hallin, and Marc Henry. Monge–kantorovich
depth, quantiles, ranks and signs. The Annals of Statistics, 45(1):223–256, 2017.
[16] Gregor Kasieczka and David Shih. Datasets for boosted w tagging, January 2020. URL
https://doi.org/10.5281/zenodo.3606767.
[17] Performance of mass-decorrelated jet substructure observables for hadronic two-body decay
tagging in ATLAS. Technical report, CERN, Geneva, Jul 2018. URL https://cds.cern.
ch/record/2630973.
7

[18] Prajit Ramachandran, Barret Zoph, and Quoc V Le. Searching for activation functions. arXiv
preprint arXiv:1710.05941, 2017.
[19] Kunihiko Fukushima. Cognitron: A self-organizing multilayered neural network. Biological
Cybernetics, 20(3):121–136, 1975. doi: 10.1007/BF00342633. URL https://doi.org/10.
1007/BF00342633.
[20] Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios. Neural spline flows.
Advances in neural information processing systems, 32, 2019.
[21] Jernej F Kamenik and Manuel Szewc. Null hypothesis test for anomaly detection. arXiv preprint
arXiv:2210.02226, 2022.
[22] Gregor Kasieczka, Benjamin Nachman, Matthew D. Schwartz, and David Shih. Automating the abcd method with machine learning. Phys. Rev. D, 103:035021, Feb 2021. doi:
10.1103/PhysRevD.103.035021. URL https://link.aps.org/doi/10.1103/PhysRevD.
103.035021.
[23] Anna Hallin, Joshua Isaacson, Gregor Kasieczka, Claudius Krause, Benjamin Nachman, Tobias
Quadfasel, Matthias Schlaffer, David Shih, and Manuel Sommerhalder. Classifying anomalies
through outer density estimation (cathode). arXiv preprint arXiv:2109.00546, 2021.
[24] John Andrew Raine, Samuel Klein, Debajyoti Sengupta, and Tobias Golling. Curtains for yoursliding window: Constructing unobserved regions by transforming adjacent intervals, 2022.
URL https://arxiv.org/abs/2203.09470.
[25] James Dolen, Philip Harris, Simone Marzani, Salvatore Rappoccio, and Nhan Tran. Thinking
outside the ROCs: Designing decorrelated taggers (DDT) for jet substructure. Journal ofHigh Energy Physics, 2016(5), may 2016. doi: 10.1007/jhep05(2016)156. URL https://doi.org/10.1007%2Fjhep05%282016%29156.
[26] Sahibsingh A. Dudani. The distance-weighted k-nearest-neighbor rule. IEEE Transactions on
Systems, Man, and Cybernetics, SMC-6(4):325–327, 1976. doi: 10.1109/TSMC.1976.5408784.
[27] Chin-Wei Huang, Ricky T. Q. Chen, Christos Tsirigotis, and Aaron Courville. Convex potential
flows: Universal probability distributions with optimal transport and convex optimization, 2020.
URL https://arxiv.org/abs/2012.05942.
[28] Matteo Cacciari, Gavin P Salam, and Gregory Soyez. Fastjet user manual. The European
Physical Journal C, 72(3):1–54, 2012.
[29] Matteo Cacciari and Gavin P Salam. Dispelling the n3 myth for the kt jet-finder. Physics Letters
B, 641(1):57–61, 2006.
[30] Matteo Cacciari, Gavin P Salam, and Gregory Soyez. The anti-kt jet clustering algorithm.
Journal of High Energy Physics, 2008(04):063, 2008.
[31] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint
arXiv:1412.6980, 2014.
[32] Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv
preprint arXiv:1608.03983, 2016.
[33] Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios. nflows: normalizing
flows in PyTorch, November 2020. URL https://doi.org/10.5281/zenodo.4296287.
[34] Roger Koenker. Quantile Regression. Econometric Society Monographs. Cambridge University
Press, 2005. doi: 10.1017/CBO9780511754098.
[35] Roger Koenker and Kevin F. Hallock. Quantile regression. Journal of Economic Perspectives,
15(4):143–156, December 2001. doi: 10.1257/jep.15.4.143. URL https://www.aeaweb.
org/articles?id=10.1257/jep.15.4.143.
[36] Anna Hallin, Gregor Kasieczka, Tobias Quadfasel, David Shih, and Manuel Sommerhalder.
Resonant anomaly detection without background sculpting, 2022. URL https://arxiv.org/
abs/2210.14924.
8

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope? [Yes](b) Did you describe the limitations of your work? [Yes] See conclusion section.
(c) Did you discuss any potential negative societal impacts of your work? [No] We do notforsee any.
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem? [Yes]2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A](b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [Yes] Code will bemade publicly available with the camera ready version.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they
were chosen)? [Yes] Given the limited space in the main body this is given in App. F.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [Yes] Yes, ten random seeds were run for every experiment.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)? [Yes]4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes](b) Did you mention the license of the assets? [No] Everything is public.
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes]
The training utilities will be made public.
(d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating? [N/A] All data was publicly available.
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content? [N/A] Physics data.
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable? [N/A](b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable? [N/A](c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation? [N/A]9

Background
Signal

0.6

0.5

0.0

1.0
p(x)

p(x|m = 1)

p(x | m = 0)

1.0

0.5

0.2

0.0
−15

−10

−5
x

0

5

0.4

0.0
−15

−10

−5
x

0

5

−15

−10

−5
x

0

5

Figure 5: A dummy dataset for illustrating possible pathologies with task agnostic decorrelation.

A

Task agnostic decorrelation

Normalizing flows can be used to find a decorrelated representation of generic variables from genericconditions. However, this can come at a price in terms of performance on a given task, in particular
if the decorrelation is task agnostic. To illustrate this consider a dataset with a binary conditionalvariable m = {0, 1} and a single scalar x that is intended to be used to be used to discriminate signal
from background. In this dummy example we takepsignal (x|m = 0) = N (2, 0.3)
pbackground (x|m = 0) = N (0, 1)

psignal (x|m = 1) = N (−12, 0.3),
pbackground (x|m = 1) = N (−10, 1),

as illustrated in Fig. 5.
There are two invertible functions that could decorrelate the background in this setting. These
functions can be defined piecewise asx
m=0
f1 (x, m) =x + 10 m = 1
x
m=0
f2 (x, m) =−x − 10 m = 1.
Where the two functions both produce a representation of the data such that the background isperfectly perfectly decorrelated from the feature m. However, the two representations are notequivalent for separating signal from background. Given that we know the likelihoods of the data we
can directly construct the likelihood ratio between the signal and background which is the optimal
test statistic in this setting as stated by the Neyman-Pearson lemma. The separation of each of thepossible representations is shown in Fig. 6 where the area under the receiver operating characteristic
is also calculated to quantify the separation performance of each of the representations of the data.
We observe that the function f1 has a significantly reduced signal discrimination power due to theincreased mixing of signal and background at different mass values.
In general task agnostic decorrelation leads to many degenerate solutions with different performanceon the proposed task. A simple approach to this dilemma is to first construct a network to perform agiven task and then decorrelate. An alternative would be to define a class conditional base distribution
such that the representation fθ(m) (x) is still class independent but the signal is always clustered in
the representation and the previous degeneracy is removed. Such a representation could be used by
maximizing
log pθ (x|m, c) = log p(fθ(m) (x)|c) + log det Jfθ(m) (x) .

10

id

f1

0.6

0.6

0.4

0.4

0.2

0.2

0.0

0.0

f2
Background
Signal

1.25

p(x)

1.00
0.75
0.50
0.25

psig (x)/pbkg (x)

30

0.00

30

AUC: 0.98

AUC: 0.95

30

20

20

20

10

10

10

0

−15

−10

−5
x

0

5

0

−5.0

−2.5

0.0
x

2.5

5.0

0

AUC: 0.98

−5.0

−2.5

0.0
x

2.5

5.0

Figure 6: The joint distributions of each of the possible decorrelated representations.

11

B

Existing approaches

Designed decorrelated taggers (DDTs) [25] apply a handcrafted linear transformation to the discriminant to ensure decorrelation. A generalization of this to non-linear dependence on m is given by fixed
efficiency regression where the expected response for background examples is subtracted from eachvalue of the discriminant [17] with the expected response estimated using k-nearest neighbors [26].
By subtracting the expected value of s this method does not account for the different shape of s asa function of m and is therefore only decorrelated when rejecting 50% of the background. These
approaches are similar to what we propose in that they are invertible, but we consider decorrelationfrom a more general perspective and account for the shape of the distribution of s at all values of mas well as providing a procedure that can be extended to higher dimensions, many variables in m andhighly non-linear dependence between s and m.

12

C

Monotonically increasing flows

Normalizing flows are by definition monotonic and in one dimension it is easy to test if they havepositive slope at any value of the condition m by encoding two points x0 , x1 such that x0 < x1 andensuring fθ (x0 , m) < fθ (x1 , m). If a base density that is symmetric about its mean c is used to train
the flow then a point x∗ that does not satisfy this condition can be forced to transform under somefunction with positive slope by applying the transformation x∗ = −fθ (x∗ , m) + 2c. As normalizing
flows typically use a base density that is symmetric about the mean this restriction does not impose
strong constraints on the flows flexibility.

13

D

Decorrelation with quantile regression.

Quantile regression [12] is the task of finding the quantiles of a given distribution as a functionof explanatory variables. When the base distribution p of a normalizing flow is taken to be theuniform distribution on [0, 1] and fθ is restricted to being monotonically increasing then fθ−1 will be
the quantile function. However, a normalizing flow must learn all quantiles simultaneously, whilequantile regression targets the simpler task of learning prescribed intervals. Typically we know whatquantiles we want to use for a given discriminant before it is constructed, and by learning thesequantiles as a function of m we can use the conditional quantiles to perform event selections suchthat the distribution of m does not change. Quantile regression is useful in the case of discriminants,
but it is not a general tool for decorrelation as only the predefined quantiles will be decorrelated
from the data. Given the conditional quantiles a spline approximation to the quantile function can beconstructed, but this is not explored here as a normalizing flow learns this function directly.
The definition of quantiles was recently generalised to greater than one dimension [14, 15] with a fast
non-linear method now available [15]. A normalizing flow can be used to learn the higher dimensional
quantile function if fθ is defined to be the gradient of a complex potential as has been developed byconvex potential flows [27]. For decorrelating discriminants in more than one dimension this is theformulation that should be used to preserve the ordering that makes the discriminant useful.

14

QCD
W

Normalised counts

10−1
10−2
10−3
10−4
10−5
50

100

150
200
Mass [GeV]250

Figure 7: The invariant mass distribution for the background samples without applying any selection.

E

Dataset details√
Both signal and background were generated by P YTHIA at s = 13 TeV with a detector simulatedby D ELPHES. Jets are reconstructed using FAST J ET [28, 29] and clustered using the anti-kt algorithm [30] with R = 1.0. Each jet is required to have transverse momentum pT ∈ [300, 400] GeV
and mass m ∈ [50, 300]. For each jet ten substructure variables are calculated and used as input tothe machine learning classifiers, these variables are the same as those used in previous studies ofdecorrelation [17, 8, 7]. The invariant mass distribution is shown in Fig. 7

15

FModel settings

The classifiers are trained for 100 epochs using the A DAM optimizer [31] with an initial learning rateof 0.001 annealed to zero following a cosine schedule [32]. All classifiers are trained using binary
cross entropy for Lclass . Following the prescriptions of the respective papers, models trained with
MoDe decorrelation use a batch size of 214 and 2048 for DisCo models. The vDNN uses a batch sizeof 256.
The conditional normalizing flow used to decorrelate discriminants are parameterized by three rational
quadratic spline layers [20] with 8 bins in each layer and the knot placements parameterized by
two residual blocks with 64 hidden nodes as implemented in the nflows library [33]. The basedistribution is chosen to be uniform with support on [0, 1]. The quantiles are regressed using thepinball loss [34, 35] where the quantiles are predicted using a DNN with three hidden layers with 64
nodes. Both quantile regressors and normalizing flows are trained using the same set up as the vDNN.
The conditional normalizing flow decorrelator is trained on the background only.
A conditional normalizing flow will also be used to decorrelate the input features directly. This flow
is constructed from four autoregressive rational quadratic spline flow layers [20] with 10 bins in eachlayer and two residual blocks with 64 nodes per layer to learn the knot placements. The base density
is standard normal distribution and a tail bound of 3 was used.
The hyperparameters specified here were not tuned in any way and classifier architectures were
copied from existing papers. The predefined training, testing and validation data sets were used.
For all experiments ten models with different random seeds were trained. Training a single modelcomprises training a classifier, a flow and a quantile regressor. All experiments ran for < 2.5 hours
on a single Nvidia RTX 3090 GPU card.

16

IdealDisCo
cf-DisCo
DNN (cf-inputs)

104

1/JSD70−10050

1/JSD50

104
103
102

103

cf-DNN (cf-inputs)
vDNN
cf-vDNN

102

101
5

10

20

30

5

10
70−100R50

R50

Figure 8: The performance of a DisCo classifier using the training pipeline of Ref. [7] and a flow
decorrelated output trained on top of this model’s prediction (cf-DisCo).

G

DisCo

The training procedure used in the original DisCo paper is more optimised than the classifiers trainedwith DisCo shown in the body of this work. The code used to produce the results of the DisCoclassifier as trained in Ref. [7] is publicly available at https://github.com/davidshih17/DisCo.
Training classifiers using this framework and saving the predictions on the training, validation andtest sets allows a conditional normalizing flow to be trained on these predictions. The result of thisprocedure, evaluated in the same manner as in the body of this work, can be seen in Fig. 8. This
demonstrates the utility of optimising the classifier training further to improve the performance ofthe conditional flow decorrelation. While the classifier trained on a decorrelated representation of
the features (DNN (cf-inputs)) slightly outperforms the DisCo trained classifiers in this setting theconditional flow decorrelated DisCo is significantly better and doesn’t rely on a complex flow modelfor decorrelating the input features.

17

H

Signal separation

The separation performance of a given discriminant can be quantified by various different metrics fs .
In all cases the mass unconditional values of these metrics is easy to calculate, but as we are assumingthat resonant physics will be localised in the resonant variable, this does not measure how useful thediscriminant will be for a bump hunt. An alternative approach would be to consider the metrics asa function of the mass fs (m) and calculate the expected value of the metric pfs = Ep(m) [fs (m)]for some mass distribution p(m). This expectation can be taken with respect to the signal mass
distribution and produce a score that is weighted towards mass regions with large signal densities.
A particular example of this is shown in Fig. 1 for the area under the receiver operator curve (AUC). In
bins of mass the AUC is the same for a vDNN and a conditional flow deccorelated vDNN (cf-vDNN),
while the mass inclusive AUC for a vDNN is 0.90 and 0.84 for the cf-vDNN. The difference inthese AUC values is due to the different levels of mass sculpting. Small thresholds on the vDNN
discriminant will reject both background and signal events at high mass values where there is very
little signal, while the cf-vDNN rejects background uniformly across the mass spectrum. In contrast
to the mass unaware AUCs the expected AUC over the signal mass distribution is 0.84 for both the
DNN and cf-DNN. The uneven background rejection of the DNN is clearly reflected in the binned
background rejection power R50 and signal rejection at 50% signal rejection.

18

I

Decorrelating the input features

Given a set of features x that are conditionally dependent on m a decorrelated representationfφ (x|m) of x can be found by replacing s(x) in Eq. 2 with x. As described in App. A this will
inevitably lead to a reduction of the performance of the classifier, which is empirically reinforced
in Fig. 3 where the cf-inputs model has been trained on a mass decorrelated representation of the
input features. This approach provides a generic perspective on Ref. [36] which focusses on anomaly
detection and also observes a drop in performance when using a decorrelated representation.

19

Background rejection: 0.531%

Background rejection: 0.920%
104

10

DisCo
qr-DisCo
cf-DisCo
Random4

Normalized counts

103
103
102
102
101

101

100

100
50

100

150
Mass [GeV]

200

250

50

Background rejection: 0.962%

100

150
Mass [GeV]

200

250

Background rejection: 0.994%
103

Normalized counts

103
102
10

2

101

101

100

100
50

100

150
Mass [GeV]

200

250

50

100

150
Mass [GeV]

200

250

Figure 9: The mass distribution for different levels of rejection as defined by the thresholds ofthe quantile regression, which regressed the 50, 90, 95, 99th quantiles. The mass distribution after
applying cuts from a classifier trained with DisCo decorrelation (α = 200) and corrected with aconditional normalizing flow (cf-DisCo) and conditional quantile regression (qr-DisCo). Randombounds are defined by sampling the same number of statistics randomly 100 times per threshold.

J

Extended discussion

An additional benefit of both conditional normalizing flows and quantile regression is their improved
behaviour at higher levels of background rejection as shown in Fig. 9. As can be seen in this plot,even though the quantile regression does not correctly predict the quantile thresholds, the resulting
discriminant is still decorrelated from the mass. Due to the discontinuous nature of the quantiles,
and the poor quantile predictions, it is hard to consistently quantify the decorrelation performance ofquantile regression across runs.
For the model in Fig. 9 we can see that both the conditional flow decorrelated output and the quantile
regressor perform very well at decorrelating a classifier trained with DisCo decorrelation acrossseveral cuts in Fig. 4. The JSD is calculated between the mass distribution over the background after
applying the classifier threshold and a random sample from the background with the same number ofpoints as passed the classifier threshold. An ideal upper bound on the JSD is calculated by randomly
sampling from the mass distribution without replacement. Errors are estimated by calculating the
JSD 10 times.

20