Contrastive Normalizing Flows for Uncertainty-Aware Parameter Estimation

TL;DR

This work tackles uncertainty-aware parameter estimation under domain shifts induced by nuisance parameters in high-energy physics. It introduces Contrastive Normalizing Flows (CNFs) to learn discriminative, shift-robust density representations for signal and background, paired with a DNN classifier and Neyman construction to produce calibrated confidence intervals on the signal fraction μ. The method demonstrates robust performance on Toy Gaussian tasks and, critically, on the HiggsML Uncertainty Challenge, where it outperforms likelihood-based baselines by achieving tighter intervals with proper coverage under systematic distortions. The approach leverages the CNF's distributional knowledge to stabilize both classification and inference under nuisance perturbations, offering a principled, practical tool for uncertainty quantification in physics analyses.

Abstract

Estimating physical parameters from data is a crucial application of machine learning (ML) in the physical sciences. However, systematic uncertainties, such as detector miscalibration, induce data distribution distortions that can erode statistical precision. In both high-energy physics (HEP) and broader ML contexts, achieving uncertainty-aware parameter estimation under these domain shifts remains an open problem. In this work, we address this challenge of uncertainty-aware parameter estimation for a broad set of tasks critical for HEP. We introduce a novel approach based on Contrastive Normalizing Flows (CNFs), which achieves top performance on the HiggsML Uncertainty Challenge dataset. Building on the insight that a binary classifier can approximate the model parameter likelihood ratio, we address the practical limitations of expressivity and the high cost of simulating high-dimensional parameter grids by embedding data and parameters in a learned CNF mapping. This mapping yields a tunable contrastive distribution that enables robust classification under shifted data distributions. Through a combination of theoretical analysis and empirical evaluations, we demonstrate that CNFs, when coupled with a classifier and established frequentist techniques, provide principled parameter estimation and uncertainty quantification through classification that is robust to data distribution distortions.

Figures (10)

• Figure 1: Randomly sampled signal (blue) and background (red) points with contours overlaid. CNF loss-minimizing distribution contour $p^{(s)}_{\theta^*}(\mathbf{x})$ for $c =1$ (left), learned CNF distribution contour for $c=1$ (center), and DNN classifier score contour (right) for two 2-dimensional Gaussian distributions. The DNN decision boundary $r = 0.5$ (dashed black, right panel) closely matches the zero contour of the CNF distributions.
• Figure 2: A rotation of the data (black arrows) can move points across the DNN decision boundary (dashed black, left), but remain in the $p_{\theta}^{(s)} = 0$ region of the CNF distribution (right).
• Figure 3: Accuracy as a function of nuisance parameter $\phi$ for three different types of nuisance parameter deformations: a large 2D subspace rotation (left), a small 2D subspace rotation (center), and a subspace deformation (right). See text for details.
• Figure 4: Setup of the HiggsML Uncertainty Challenge. The left panel shows an example of the marginal distribution of one of the 28 features (note the log scale, demonstrating a signal fraction $f_s \ll 1$), and the right panel shows an example coverage plot.
• Figure 5: Flow chart for our CNF-based method of estimating $\mu$ and 1$\sigma$ confidence intervals.