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

Technical Summary: Learning to Bin in High-Energy Physics

Problem Statement
In high-energy physics (HEP) analyses, statistical inference often relies on binned likelihoods where events are categorized using discriminant observables. The choice of bin boundaries significantly impacts analysis sensitivity. Current practices typically involve manual selection of bin boundaries, often using equidistant or quantile-based binning in one dimension, or projecting multi-dimensional classifier scores into one dimension via "argmax" (assigning events to the class with the highest score) before binning. These approaches are not guaranteed to be optimal for the analysis figure of merit, particularly in multi-dimensional spaces where projection can discard information. While some existing methods address analysis-aware optimization (e.g., ThickBrick, INFERNO, neos), they often focus on density estimation, fix bin edges a priori, or rely on specific neural network architectures that may not integrate easily into existing workflows.

Methodology
The authors propose a framework to directly optimize bin boundaries with respect to signal significance in both one- and multi-dimensional discriminant spaces without requiring one-dimensional projections.

1. Binning Model (Gaussian Mixture Model):
   The core of the approach is a Gaussian Mixture Model (GMM) where each bin is represented by a Gaussian component.

   • Events are assigned to the component $k$ that maximizes the unnormalized logarithmic score $s_k(x) = \log \mathcal{N}(x | \mu_k, \Sigma_k) + \log \pi_k$.
   • This formulation allows for flexible, non-linear bin boundaries in multi-dimensional space, as boundaries arise from the comparison of multiple Gaussian densities rather than fixed hyperplanes.
   • Parameters (means $\mu_k$, covariances $\Sigma_k$, and weights $\pi_k$) are constrained to ensure valid probability distributions (e.g., using Cholesky factorization for positive-definite covariances and log-softmax for weights).
   • The objective function is the Asimov significance ($Z$), calculated from the binned signal ($S$) and background ($B$) yields.
   • To ensure robustness, the objective includes penalty terms for bins with expected background yields below a minimum threshold ($B_{min}$) or relative statistical uncertainties exceeding a limit ($r$).

2. Optimization Strategies:
   Two complementary strategies are introduced to optimize the GMM parameters:

   • GATO (Gradient-based cATegorization Optimizer): A differentiable approach. For multi-dimensional cases, the discrete bin assignment is approximated using a temperature-scaled softmax over the component scores, allowing gradient descent. For one-dimensional cases, a simpler alternative represents bin boundaries directly as trainable cut values approximated by sigmoid functions.
   • BOBR (Bayesian Optimization of Bin boundaRies): A Bayesian optimization approach using a Tree-structured Parzen Estimator (TPE). This method treats the bin boundaries (or GMM parameters) as hyperparameters to be optimized sequentially, balancing exploration and exploitation. It does not require differentiable approximations of the bin edges.

Key Contributions

• Direct Multi-dimensional Optimization: The method defines bin boundaries directly in the full discriminant space, avoiding information loss associated with projecting multi-class scores to one dimension.
• Flexible Bin Shapes: By using GMMs, the resulting bins can take complex, non-elliptical shapes determined by the data distribution, rather than being restricted to rectangular or equidistant regions.
• Dual Optimization Framework: The paper provides two distinct optimization engines (GATO and BOBR) implemented as lightweight Python plugins (gato-hep and bobr-hep) designed for straightforward integration into existing HEP analysis workflows.
• Constraint Handling: The framework incorporates penalty terms to prevent the formation of statistically unstable bins (e.g., those with zero background or excessive uncertainty), a critical requirement for realistic physics analyses.

Results
The methods were benchmarked on toy datasets representing typical HEP discriminants: a one-dimensional binary classification and a three-class problem (two signals, one background) with varying degrees of separability.

• One-Dimensional Case: Both GATO and BOBR achieved similar gains in signal sensitivity compared to equidistant binning for a fixed number of bins. The optimized boundaries adapted to the underlying discriminant structure, outperforming fixed binning schemes.
• Multi-Dimensional Case:
  • In scenarios with moderate separability between signals (ROC AUC $\approx$ 0.82), the GATO-based GMM approach significantly outperformed (approx. 20% higher geometric mean significance) both standard argmax classification and k-means clustering. It achieved high sensitivity with fewer bins than baseline methods.
  • In scenarios with high separability (ROC AUC $\approx$ 0.98), the performance of the multi-dimensional GMM approach became comparable to argmax classification followed by one-dimensional optimization, as the problem effectively factorized.
  • Scalability: As the number of bins (and thus free parameters) increased, the differentiable GATO approach consistently outperformed the Bayesian BOBR approach. The authors attribute this to the increasing dimensionality of the optimization problem, which challenges the efficiency of Bayesian optimization in moderately high-dimensional spaces.
• Penalty Efficacy: The study confirmed that the penalty terms effectively guided the optimization away from configurations with low background yields or high statistical uncertainties, ensuring the resulting binning schemes were robust and analysis-friendly.

Significance and Claims
The paper claims that the proposed GMM-based binning model, coupled with GATO and BOBR, offers a flexible and "analysis-aware" alternative to manual binning. The primary significance lies in the ability to learn optimal categorization directly in multi-dimensional spaces, potentially extracting more sensitivity from complex classifier outputs than traditional projection-based methods. The authors emphasize that their tools are designed for immediate integration into existing workflows. They suggest that a natural future direction is embedding this GMM categorization into fully differentiable pipelines for the joint optimization of classifiers and binning, though they note that keeping these as explicit modules may offer advantages in interpretability and validation. The work does not claim to replace standard classifiers but rather to optimize the final step of event categorization for statistical inference.