Contrastive Normalizing Flows for Uncertainty-Aware Parameter Estimation

TL;DR

This paper introduces Contrastive Normalizing Flows (CNFs), a novel method for uncertainty-aware parameter estimation in physical sciences, effectively handling data distribution shifts and improving robustness in high-energy physics applications.

Contribution

The paper proposes CNFs, a new approach that embeds data and parameters in learned flows, enabling robust, uncertainty-aware parameter estimation under domain shifts, with theoretical and empirical validation.

Findings

CNFs achieve top performance on the HiggsML Uncertainty Challenge dataset.

CNFs provide principled uncertainty quantification through classification.

The method is robust to data distribution distortions.

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.