Learning to bin: differentiable and Bayesian optimization for multi-dimensional discriminants in high-energy physics

Abstract

Categorizing events using discriminant observables is central to many high-energy physics analyses. Yet, bin boundaries are often chosen manually. A simple, popular choice in multi-classification tasks is to assign events according to the largest per-class score ("argmax") and to apply equidistant binning to the resulting one-dimensional discriminants. We propose a binning optimization for signal significance directly in multi-dimensional discriminants. We use a Gaussian Mixture Model (GMM) to define flexible regions in the score space, which can be interpreted either as bins or as analysis categories. While this GMM-based strategy is applicable in both one and multiple dimensions, we also study a direct bin-boundary optimization in one dimension as a simpler alternative for binary discriminants. On this binning model, we study two optimization strategies: a differentiable and a Bayesian optimization approach. We study two toy setups: a binary classification and a three-class problem with two signals and backgrounds. In the one-dimensional case, both approaches achieve similar gains in signal sensitivity compared to equidistant binning for a given number of bins, while in the multi-dimensional case the differentiable approach performs best. We show that the GMM-based optimization can outperform argmax classification even after optimized binning is applied to the one-dimensional projections. We further study the performance of our methods on the FAIR Universe $H\rightarrowττ$ dataset, where the GMM-based optimization gives the highest signal significance. Both methods are released as lightweight Python plugins intended for straightforward integration into existing analyses.