[Paper Review] Decorrelation with conditional normalizing flows

Samuel Klein, T. Golling|arXiv (Cornell University)|Nov 4, 2022

Anomaly Detection Techniques and Applications4 citations

TL;DR

This paper proposes using conditional normalizing flows to decorrelate machine learning discriminants from protected attributes—such as invariant mass—in high-energy physics, preserving signal separation power while eliminating mass-dependent background sculpting. The method achieves near-perfect decorrelation with minimal performance loss, significantly improving downstream bump-hunt sensitivity compared to standard approaches.

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 our approach by building supervised jet taggers that produce almost no sculpting in the mass distribution of the background.

Motivation & Objective

• To address the problem of background mass sculpting in high-energy physics analyses, where signal-enriched samples distort the smooth background distribution due to discriminant correlations with invariant mass.
• To develop a post-hoc decorrelation method that preserves the original discriminant's separation power at every value of the protected attribute (e.g., mass).
• To enable robust bump-hunt analyses by ensuring uniform background rejection across the mass spectrum, independent of the discriminant's correlation with mass.
• To provide a flexible, invertible framework that can be applied to any discriminant and protected attribute set, without requiring retraining the original model.

Proposed method

• A conditional normalizing flow is trained to map the discriminant score $ s_{\phi}(\boldsymbol{x}) $ to a new representation $ f_{\theta}(s_{\phi}(\boldsymbol{x}), \boldsymbol{m}) $, where $ \boldsymbol{m} $ is the protected attribute (e.g., mass).
• The flow is trained via maximum log-likelihood under the change of variables formula, ensuring the transformed output follows a base distribution independent of $ \boldsymbol{m} $, thus achieving decorrelation.
• The invertibility of the flow ensures that the signal separation power (e.g., AUC) remains unchanged at every fixed value of $ \boldsymbol{m} $, preserving the discriminant's utility.
• The flow is constrained to be monotonically increasing to maintain the ordering of scores across all mass values, which is critical for maintaining separation power.
• The method can be applied as a post-processing layer on top of any pre-trained classifier, enabling decorrelation without retraining the original model.
• Alternative approaches like conditional quantile regression are also evaluated, showing similar decorrelation performance but with less stable quantile predictions.

Experimental results

Research questions

• RQ1 Can a post-hoc transformation decorrelate a discriminant from protected attributes like invariant mass without degrading its signal separation power?
• RQ2 How effective is a conditional normalizing flow in achieving uniform background rejection across the mass spectrum compared to standard decorrelation methods?
• RQ3 Does applying a conditional flow to the output of a pre-trained classifier improve the robustness of bump-hunt analyses in high-energy physics?
• RQ4 What is the impact of decorrelating the discriminant output versus the input features on classifier performance?
• RQ5 How do conditional normalizing flows compare to quantile regression in terms of stability and decorrelation performance across multiple runs?

Key findings

• The conditional normalizing flow achieves near-perfect decorrelation of the discriminant from the mass variable, as evidenced by a Jensen-Shannon divergence (JSD) of approximately 0.001 between the background mass distribution after selection and a random sample.
• The mass-inclusive AUC drops from 0.90 for a standard vDNN to 0.84 for the flow-decorrelated version, but this reflects the removal of mass sculpting, not performance loss, as the expected AUC over the signal mass distribution remains 0.84 for both models.
• The conditional flow-decorrelated DisCo classifier outperforms both standard DisCo and DisCo with input decorrelation, demonstrating that output decorrelation is more effective and less costly than input-level decorrelation.
• The method preserves signal separation power at every mass value, as the invertible transformation ensures that the ordering of events by score is maintained across all mass bins.
• Quantile regression also achieves good decorrelation performance, but with higher variance across runs due to discontinuous quantile predictions, making it less reliable than the flow-based approach.
• The approach is robust across different levels of background rejection, with both the conditional flow and quantile regression maintaining decorrelation even at high rejection thresholds, as shown in Figure 9.