Physically Consistent Parameter Inference: Transparent Machine Learning Emulation in High Energy Physics and Cosmology

Technical Summary: Physically Consistent Parameter Inference via Transparent Machine Learning

Problem Statement
Global fits in high-energy physics (HEP) and cosmology frequently encounter the challenge of exploring high-dimensional parameter spaces constrained by hundreds of correlated measurements. The likelihood functions in these domains often exhibit complex topologies, including strong parameter correlations, non-Gaussian behavior, multiple local minima, and elongated "curved" degeneracies. Furthermore, evaluating the theoretical predictions required for these likelihoods is computationally expensive, often involving renormalization-group evolutions, numerical distance integrals, or the calculation of hundreds of observables. Traditional statistical approaches, such as Markov Chain Monte Carlo (MCMC) or nested sampling, rely on repeated likelihood evaluations, rendering exhaustive parameter scans impractical or prohibitively slow. There is a critical need for fast, interpretable surrogate models that can emulate these complex likelihood landscapes without sacrificing statistical rigor or physical interpretability.

Methodology
The authors propose a unified Machine Learning (ML) framework designed to emulate the log-likelihood function (or $\chi^2$) using XGBoost (eXtreme Gradient Boosting), an ensemble of gradient-boosted regression trees. The framework integrates three core components:

1. Surrogate Construction:

   • Training Data Generation: The training dataset is constructed using a hybrid strategy. It combines a space-filling design (using Latin Hypercube Sampling or Sobol sequences) over the full prior volume with a denser Gaussian cloud concentrated around the best-fit region, derived from the Hessian of the fit. This ensures coverage of both broad parameter spaces and narrow, anisotropic degeneracies.
   • Active Learning: For regions where the likelihood is sparse or exhibits sharp structures, an active learning loop refines the dataset. A Gaussian Process (GP) surrogate predicts uncertainty, and an acquisition function (Expected Improvement) selects new points for exact evaluation, progressively concentrating sampling where it improves the emulated likelihood most.
   • Target Transformation: To handle the vast dynamic range of $\chi^2$ values (from $\sim 10^3$ to $10^6$) and ensure the model resolves the narrow confidence regions near the minimum, the authors employ specific target transformations. These include direct log-likelihood for flavor fits, a two-stage classifier-plus-regressor strategy using a sigmoid transform for axion-like particles (ALPs), and a shifted-log10 transform ($y = \log_{10}(\chi^2 - \chi^2_{\min} + 1)$) for cosmology. This shift maps the minimum to a consistent steep region of the transform regardless of the absolute $\chi^2$ value.

2. Interpretability via SHAP:

   • To ensure the ML predictions remain physically consistent, the framework utilizes SHAP (Shapley Additive exPlanations) values. Based on cooperative game theory, SHAP quantifies the marginal contribution of each feature (model parameter) to the prediction.
   • The TreeSHAP algorithm is used to compute exact SHAP values in polynomial time for tree ensembles. This provides a model-independent ranking of parameter importance and reveals non-linear dependencies, ensuring the surrogate captures the underlying physics rather than numerical noise.

3. Posterior Exploration:

   • The emulated $\chi^2$ is sampled using a parallel Random-Walk Metropolis–Hastings algorithm. The proposal covariance is set to the Cholesky factor of the fit covariance to improve mixing in correlated spaces.
   • Closure Tests: The framework validates the surrogate by comparing the distribution of $\Delta\chi^2$ over accepted points against the expected $\chi^2$ distribution (Wilks' theorem) and, in the cosmological case, by directly comparing the emulated posterior against the exact likelihood sampled via multiple independent back-ends (Astropy, CPU, GPU).

Key Applications and Results
The framework is validated across three distinct global-fit applications of increasing dimensionality and complexity:

1. Flavor Anomalies in Semileptonic B Meson Decays (SMEFT19):

   • Context: Analyzing deviations in $b \to c\tau\nu$, $b \to s\nu\bar{\nu}$, and $b \to s\ell^+\ell^-$ transitions using the SMEFT19 framework (Scenario III).
   • Performance: The XGBoost surrogate achieved a Pearson correlation coefficient of $r=0.99$ with the exact likelihood.
   • Speed-up: Evaluation speed increased by a factor of $\sim 10^5$ (from $\sim 3$ seconds per point to $\sim 34$ microseconds).
   • Physics Insight: The surrogate successfully reproduced the breaking of the $C_1 = C_3$ degeneracy and the curved ridge in the $(C_1, \beta_q)$ plane. SHAP analysis correctly ranked parameters ($C_3 > C_1 > \beta_q$) consistent with the underlying operator structure. Crucially, it revealed that the strong correlation between $R_D$ and $BR(B^+ \to K^+\nu\bar{\nu})$ found in previous scenarios disappears when $C_1$ and $C_3$ vary independently.

2. Axion-Like Particles (ALPs) and the Belle II Anomaly:

   • Context: Explaining the $2.8\sigma$ excess in $B^+ \to K^+\nu\bar{\nu}$ with a light ALP ($m_a \approx 2$ GeV). The likelihood is non-Gaussian and discontinuous due to kinematic thresholds and lifetime-dependent decay probabilities.
   • Strategy: A two-stage approach was used: a classifier (sigmoid transform) to isolate the relevant low-$\chi^2$ region, followed by a regressor.
   • Performance: The classifier achieved an AUC of 0.974, and the regressor reached $r=0.929$.
   • Speed-up: Evaluation speed increased by $\sim 5.3 \times 10^5$ (from $\sim 1.9$ seconds to $\sim 3.5$ microseconds).
   • Physics Insight: SHAP identified the ALP mass ($m_a$) as the most critical parameter, followed by the PQ charge of the third-generation left-handed quarks ($X^3_{qL}$) and the decay constant ($f_a$). The posterior correctly favored $m_a \approx 1.8$ GeV and $f_a \in [10^6, 10^7]$ GeV.

3. Cosmology ($\Lambda$CDM and $w_0w_a$CDM):

   • Context: Fitting Type Ia supernovae (Pantheon+, DES), Baryon Acoustic Oscillations (DESI), and a compressed CMB prior to constrain dark energy parameters ($w_0, w_a$).
   • Validation: Unlike the other cases, the exact cosmological likelihood is cheap enough to sample directly. The authors ran MCMC on the exact likelihood using four different back-ends (Astropy, CPU, GPU, and the ML surrogate) and confirmed the posteriors were statistically indistinguishable.
   • Speed-up: The surrogate was $3.6$–$22.4\times$ faster than the vectorized CPU implementation and $7.2$–$12.1\times$ faster than the GPU kernel.
   • Physics Insight: The surrogate captured the preference for dynamical dark energy ($w_0 > -1, w_a < 0$) in supernova+BAO data and the reversion to $\Lambda$CDM when the CMB prior was added. SHAP analysis correctly identified the hierarchy of constraints: without CMB, background parameters ($H_0, \Omega_m$) dominated; with CMB, dark energy parameters ($w_0, w_a$) became the most constrained, reflecting the physical tightening of the contours.

Significance and Claims
The paper claims to present a general, interpretable, and physically consistent framework for likelihood-based inference. Its primary significance lies in demonstrating that:

• Efficiency: ML surrogates can accelerate parameter-space exploration by several orders of magnitude (up to $10^5$) while maintaining high accuracy.
• Robustness: The method handles diverse likelihood topologies, from smooth Gaussian-like surfaces to non-Gaussian, discontinuous, and highly structured landscapes, through adaptive training designs and target transformations.
• Transparency: By integrating SHAP values, the framework avoids the "black box" nature of typical ML, ensuring that the inferred parameter hierarchies and correlations are physically interpretable and consistent with the underlying theory.
• Validation: The cosmological application serves as a controlled benchmark, proving that the speed-up does not come at the expense of statistical accuracy, as the emulated posteriors are indistinguishable from those derived from exact calculations.

The authors conclude that this framework is readily transferable to other global-fit problems and offers a foundation for future developments in automated active learning and differentiable emulators.