Physically Consistent Parameter Inference: Transparent Machine Learning Emulation in High Energy Physics and Cosmology Jorge Aldaa,b1 , Jacobo Asoreyb,c2 , Alejandro Mirb,c3 , Siannah Peñarandab,c4 a Dipartimento di Fisica e Astronomia “Galileo Galilei”, Università degli Studi di Padova & INFN
Sezione di Padova, via Marzolo 8, 35131 Padova, Italy b Centro de Astropartı́culas y Fı́sica de Altas Energı́as (CAPA), Pedro Cerbuna 12 50009 Zaragoza,
Spain
arXiv:2607.12726v1 [hep-ph] 14 Jul 2026
c Departamento de Fı́sica Teórica, Facultad de Ciencias,
Universidad de Zaragoza, Pedro Cerbuna 12, E-50009 Zaragoza, Spain
Abstract Global fits in high energy physics and cosmology often face the challenge of exploring highdimensional parameter spaces with computationally expensive or topologically complex likelihood functions. In this work, we present a Machine Learning framework designed to emulate complex, often non-Gaussian likelihood landscapes using gradient-boosted regression trees (XGBoost). We discuss the advantages of the Machine Learning approach in terms of computational efficiency and the resolution of confidence regions, particularly in scenarios with complex correlations or “curved” degeneracies. We validate this methodology by applying it to a recent analysis on flavour anomalies in semileptonic B meson decays and discussing the adaptability of this framework to other phenomenological systems, such as axion-like particles or cosmology global fits. Finally, we utilise SHAP (Shapley Additive exPlanations) values to provide a transparent analysis of feature importance, ensuring that the Machine Learning predictions remain physically interpretable and consistent with the underlying physics.
1
[email protected] [email protected] 3 [email protected] 4 [email protected] 2
1
Introduction
Despite recent progress, the exploration of highly non-Gaussian likelihoods remains computationally challenging, motivating the development of fast and interpretable surrogate models. Contemporary analyses in particle physics and cosmology often involve exploring high-dimensional parameter spaces that are constrained by hundreds of correlated measurements. Examples include flavour anomalies in semileptonic B meson decays and constraints on light new particles. The resulting likelihood functions frequently exhibit intricate structures, such as strong parameter correlations, non-Gaussian behavior, multiple local minima, and elongated degeneracy directions. At the same time, evaluating the theoretical predictions required for these likelihoods often requires computationally intensive calculations, rendering exhaustive parameter scans impractical. Traditional statistical approaches rely on repeated evaluations of the likelihood function using Monte Carlo sampling techniques, such as Markov chain Monte Carlo (MCMC) or nested sampling. However, the computational cost of these approaches increases rapidly with the dimensionality and complexity of the problem. This has motivated increasing interest in Machine Learning (ML) techniques that can emulate expensive likelihood functions while preserving the statistical information necessary for robust inference. Combining global fits with interpretable ML techniques provides a powerful framework for disentangling complex correlations among observables, identifying the most informative measurements, and improving the physical interpretability of New Physics (NP) constraints in the phenomenology of particle physics and cosmology. Gradient-boosted decision trees have emerged as a particularly attractive ML algorithm due to their flexibility, robustness, and ability to accurately capture nonlinear structures. An important challenge in this type of analysis is understanding the relationship between observables and model parameters. SHAP values provide a systematic approach to quantifying the contribution of individual features to model predictions. This enables transparent and physically meaningful interpretations of ML models. In this work, we present a unified and interpretable ML framework for global fits in high energy physics and cosmology, in which an XGBoost gradient-boosted regression ensemble is trained to emulate the (generally non-Gaussian) log-likelihood, while SHAP values keep the resulting inference transparent and physically meaningful. Rather than tailoring the method to a single problem, we demonstrate its generality across three representative global-fit applications of increasing dimensionality and complexity, spanning both particle physics and cosmology. These include the flavour anomalies in semileptonic B meson decays, light and feebly-interacting new states (axion-like particles addressing the Belle II B + → K + ν ν̄ excess), and dynamical dark-energy fits that drive contemporary precision cosmology. In all three cases, the surrogate reproduces the exact likelihood to high accuracy and accelerates parameter-space exploration by several orders of magnitude. Furthermore, SHAP analysis recovers the physically expected hierarchy of observables and parameters, ensuring that the gain in computational efficiency does not come at the expense of interpretability. The paper is organised as follows. Section 2 introduces the ML algorithm at the heart of the framework: the construction of the XGBoost surrogate of the χ2 function, the underlying statistical formalism (likelihood, profile likelihood and confidence regions), the strategies used to build and sample the training set, the choice of training target, the SHAP interpretation of feature importance, and the Markov-Chain Monte Carlo exploration of the emulated posterior together with its closure tests. Section 3 presents three applications of the framework. We first analyse flavour anomalies in semileptonic B meson decays, construct an XGBoost surrogate model of the SMEFT19 global likelihood, and use it to map high-resolution confidence regions. We also rank the model parameters with SHAP, and study the correlations among low-energy observables that characterise the considered scenario. Then, we consider light NP, where the lightness of the new states induces non-Gaussian and even discontinuous likelihoods. As a concrete example, we explore axion-like particles as a candidate explanation for the Belle II B + → K + ν ν̄ excess, training the surrogate through an active-learning programme and a two-stage classifier-plus-regressor strategy to resolve the relevant region of parameter space. Finally, we address cosmology, fitting the ΛCDM and dynamical dark-energy (w0 wa CDM) models to Type Ia supernovae, baryon acoustic oscillations, and a compressed cosmic-microwave1
background prior. Because the exact cosmological likelihood is cheap enough to sample directly, this application doubles as a controlled validation of the method, in which the emulated posteriors are shown to be statistically indistinguishable from the exact ones. Our conclusions and perspectives are presented in Section 4.
2
Machine Learning Algorithm
ML algorithms are powerful tools for global fits both in high energy physics and cosmology. One central feature of the framework is the application of SHAP values to quantify the relative importance of each parameter in constraining the fit. This provides a model-independent ranking of model features based on their impact on specific Wilson coefficients of an effective field theory, and allows a transparent mapping between experimental inputs and the parameters of the effective field theories, going beyond standard χ2 -based sensitivity analyses [1–6]. With this type of analysis, we can identify the regions of parameter space allowed by any model compatible with current experimental data. We adopt a ML approach to approximate the log-likelihood function. Specifically, we use the XGBoost algorithm [7], an ensemble method based on regression trees that can efficiently approximate arbitrary functions. Regression trees divide the parameter space into subsets and assign a numerical value to each leaf. By combining many trees into an ensemble, the model constructs an approximation of the log-likelihood function using supervised learning. The training dataset consists of parameter points and their corresponding precomputed log-likelihood values. Model optimization relies on a loss function that measures prediction accuracy and a regularization term that limits model complexity to avoid overfitting. The ensemble is built iteratively by adding new tree structures while applying shrinkage techniques to improve generalization. Using the machine-learned approximation instead of the exact likelihood significantly reduces computational costs, since Monte Carlo methods require a large number of likelihood evaluations. SHAP values [8,9] are used to interpret the influence of individual parameters in the approximation. Based on concepts from cooperative game theory [10], SHAP values quantify the contribution of each feature to the model prediction while satisfying properties such as local accuracy, missingness, and consistency. SHAP analysis evaluates the marginal impact of including each parameter in the model and provides insight into the relative importance of the different fit parameters. The statistical framework summarised in this section follows Ref. [6]. The central statistical object in a global fit is the likelihood function L(θ), defined as the conditional probability of observing the experimental data given the model parameters θ, L(θ) = p(data | θ) .
(1)
It is convenient to work with χ2 (θ) ≡ −2 log L(θ). For N statistically independent observables with Gaussian uncertainties, N Y 1 (yi − xi (θ))2 √ L(θ) = exp − , (2) 2 σi2 2π σi i=1 where yi are the experimental measurements, xi (θ) are the theoretical predictions, and σi2 = σy2i + σx2i is the sum of experimental and theoretical variances. When the observables are correlated, the P 2 product is replaced by the multivariate Gaussian χ (θ) = ij (yi − xi (θ)) (C −1 )ij (yj − xj (θ)), with C the combined experimental and theoretical covariance matrix. The best-fit point is the Maximum Likelihood Estimator (MLE), θ̂ = arg max L(θ) = arg min χ2 (θ) , θ
θ
(3)
and the goodness of a new-physics hypothesis is quantified by its pull with respect to the Standard Model (SM), ∆χ2SM = χ2 (θSM ) − χ2 (θ̂). Confidence regions are obtained via the profile likelihood. The profile likelihood for the i-th parameter is defined as LP (θi ) = max L(θ1 , . . . , θn ) , (4) θj̸=i
2
and the statistic test λP (θi ) = −2 log
LP (θi ) L(θ̂)
= ∆χ2 (θi )
(5)
follows asymptotically a χ2 distribution with one degree of freedom (Wilks’ theorem [11]). The k σ confidence interval on θi is therefore the set of values for which λP (θi ) ≤ k 2 . For two-dimensional confidence regions one uses the profile likelihood over the remaining parameters, and the 1 σ and 2 σ contours correspond to ∆χ2 = 2.30 and 6.18 respectively. In practice, however, the likelihood in Eq. (2) is often only a simplified picture. Some applications considered in this work involve likelihoods that are non-Gaussian, depend on a large number of correlated observables, and require expensive theoretical computations for each prediction xi (θ). Mapping the full χ2 surface therefore requires a large number of evaluations, making a brute-force scan computationally prohibitive. We address this by training a ML surrogate to approximate the χ2 function [6]. 2 The training dataset D = {(θi , yi )}N i=1 is built by evaluating the exact χ at a set of parameterspace points. To distribute these points efficiently we combine a space-filling design over the full prior volume with a denser draw concentrated around the best-fit region, where the likelihood is largest and the confidence contours are most sensitive. For the space-filling component we use, depending on the application, plain uniform random sampling, Latin Hypercube Sampling (LHS) [12], or a lowdiscrepancy Sobol sequence [13]. The latter two guarantee substantially more uniform coverage than a purely random sample of the same size, while uniform sampling is adequate whenever a large point budget and the dense best-fit cloud already cover the relevant region. The denser best-fit draw is most efficiently realised as a Gaussian cloud whose covariance is derived from the Hessian of the fit, so that the sampling follows the actual likelihood ellipsoid and automatically resolves narrow, anisotropic degeneracies without a hand-tuned grid. When the relevant region is not known a priori, or when the likelihood develops sharp or narrow structures that a static design samples too coarsely, this base design is refined iteratively through active learning. At each iteration an inexpensive auxiliary model is trained on the data accumulated so far, typically a Gaussian Process (GP), the infinite-dimensional limit of a multivariate Gaussian, which returns both a prediction and an uncertainty at every point. The next points to be evaluated exactly are then chosen by maximising an acquisition function such as the Expected Improvement, which balances exploration of the regions where the surrogate is most uncertain against exploitation of the regions of lowest χ2 . The exact χ2 is computed at the selected points, they are appended to the training set, and the cycle repeats, progressively concentrating the sampling effort where it most improves the emulated likelihood. The surrogate model is an XGBoost gradient-boosted ensemble of regression trees [7]. A single regression tree partitions the parameter space into T disjoint regions (“leaves”) through a sequence of binary cuts and assigns a constant weight to each one, defining a function f (θ) = wq(θ) , where q(θ) ∈ {1, . . . , T } is the leaf into which θ falls and wq(θ) its weight. Since a single tree is too rigid to reproduce an arbitrary function, the surrogate is built as an ensemble of K trees, ϕ(θ) =
K X
f
(k)
(θ) =
k=1
K X
(k)
wq(θ) ,
(6)
k=1
whose output ϕ(θ) is the approximation of the training target (the χ2 or the log-likelihood). The trees are obtained by supervised learning on the dataset D, by minimising the objective X X Obj[ϕ] = ℓ ϕ(θi ), yi + Ω f (k) , (7) i
k
which balances a loss function against a regularisation term. The loss ℓ measures the discrepancy between the target yi and its approximation ϕ(θi ); we use the mean absolute error, ℓ(ϕ(θi ), yi ) = |ϕ(θi ) − yi |, which is more robust to the large-χ2 outliers than a quadratic loss. The regularisation Ω penalises the complexity of each tree, Ω(f ) = γ T + 12 λ ∥w∥2 , 3
(8)
with T the number of leaves and w the vector of leaf weights, so that the hyperparameters γ and λ introduced below directly control the trade-off between flexibility and overfitting. XGBoost additionally provides an optional L1 penalty α ∥w∥1 on the leaf weights (the hyperparameter α below), which we employ in the axion-like-particle fit. The ensemble is built iteratively: starting from a single tree that captures the overall shape of the χ2 surface, each subsequent tree provides an incremental correction whose newly added weights are scaled (shrinkage) by a learning rate η < 1. Training is stopped early as soon as the validation loss stagnates for a fixed number of consecutive rounds. Because regression trees partition the parameter space adaptively rather than assuming global smoothness, they are particularly well suited to likelihoods with narrow ridges, sharp transitions and plateau-like regions, where neural-network surrogates would require significantly larger training sets and careful regularisation to avoid over-smoothing the fine structure. The behavior of the model is controlled by a set of hyperparameters: • Number of estimators: number of trees in the ensemble. • Maximum depth: maximum number of levels per tree. Large values increase expressiveness but also the tendency to overfit. • Learning rate (η): scales the contribution of each tree. Smaller values require more trees but reduce overfitting. • Minimum split loss (γ): minimum loss reduction required to create a new partition in a leaf. Larger values yield more conservative trees. • Regularisation parameters (α, λ): penalise large leaf weights, further preventing overfitting. • Subsampling fractions (subsample, colsample bytree): proportion of training rows and feature columns used per tree, introducing randomness and reducing variance. These hyperparameters can be tuned automatically using Optuna [14], a gradient-free optimisation library based on a tree-structured Parzen estimator, although for low-dimensional, smooth landscapes a fixed, conservative configuration (moderate depth, small learning rate, aggressive early stopping) is often more robust. The choice of training target (direct χ2 , sigmoid transform, or shifted-log10 transform) depends on the structure of the likelihood landscape and is discussed in each application section. Because tree ensembles do not extrapolate reliably beyond their training range, the surrogate is trusted only within the sampled domain: the space-filling component of the training design covers the full prior volume, and the sampling density is increased wherever the landscape is most structured, so that the emulated likelihood remains faithful across the whole parameter space. After training, SHAP values [8, 9] are computed to quantify the individual contribution of each parameter to the model output at every point in the parameter space. Denoting by F the full set of n features (the fit parameters), the SHAP value of the j-th feature at a point θ is its average marginal contribution over all subsets S of the remaining features [8, 9], ϕj (θ) =
X S⊆F \{j}
|S|! (|F | − |S| − 1)! fθ (S ∪ {j}) − fθ (S) , |F |!
(9)
where fθ (S) is the prediction obtained using only the features in the subset S. These attributions satisfy three key properties: local accuracy, missingness, and consistency. Local accuracy means that they decompose the model output exactly, ϕ(θ) = ϕ0 +
n X
ϕj (θ) ,
(10)
j=1
with ϕ0 the base value, equal to the average prediction over the dataset. Missingness dictates that a feature missing from the input space or irrelevant is assigned a null contribution (ϕj = 0). Consistency guarantees that if the model changes such that the marginal contribution of a feature increases or remains the same regardless of the other features, its corresponding SHAP value will not decrease. 4
Evaluating Eq. (9) directly would require retraining the model on every feature subset, an operation whose computational cost grows exponentially with the number of features. For tree ensembles, however, the exact SHAP values are recovered in polynomial time by the TreeSHAP algorithm [9]. Aggregated over the training dataset, SHAP values provide a global, model-independent ranking of parameter importance that complements the standard ∆χ2 sensitivity analysis and reveals non-linear dependencies and correlations between parameters. To map the posterior distribution, the emulated χ2 is sampled with a Markov-Chain Monte Carlo (MCMC) algorithm. A tentative point θ ′ is accepted over the current one with the Metropolis criterion log u < log L(θ ′ ) − log L(θ) ,
u ∼ U(0, 1) ,
(11)
with u a random number uniformly distributed between 0 and 1; the proposal uses a Gaussian kernel whose covariance is, after a burn-in phase, set to the Cholesky factor of the fit covariance, substantially improving mixing in correlated parameter spaces. A parallel Random-Walk Metropolis–Hastings variant running nchains independent chains can be used, with sampling stopped automatically once the effective sample size exceeds a predefined target. A useful closure test is the distribution of ∆χ2 = 2 [log L(θ̂) − log L(θ)] over the accepted points: in the Gaussian limit, where the log-likelihood is locally quadratic, it follows a χ2 distribution with a number of degrees of freedom equal to the number of fit parameters (Wilks’ theorem [11]), so that any departure from this reference is itself a quantitative diagnostic of the non-Gaussianity of the likelihood. Because each forward pass through the trained ensemble replaces a full physics-library evaluation, the surrogate reduces the wall-clock time of the posterior exploration by several orders of magnitude.
3
Applications
3.1
B Physics
Flavour physics is an incredibly sensitive testing ground for the Standard Model. Because flavourchanging neutral currents (FCNCs) are highly suppressed within the SM framework, they are uniquely susceptible to the virtual contributions of heavy particles. In recent years, a series of persistent discrepancies between SM predictions and experimental data, known as flavour anomalies, have intensified interest in this sector [15–17]. To analyze these anomalies without committing to a single specific ultraviolet theory, effective field theories are a very useful tool [1–3]. By systematically integrating out heavy degrees of freedom, such as hypothetical NP mediators, we are left with an effective theory of contact interactions, governed by operators and their associated coupling strengths, known as Wilson Coefficients. One of the most powerful features of flavour physics is that a single NP mediator cannot alter an isolated observable without triggering a cascade of predictable shifts in others. Global fits exploit these correlations [1]. Semileptonic B meson decays provide some of the most stringent low-energy tests of the SM. Over the last decade a coherent set of deviations, the flavour anomalies, has accumulated in three classes of transitions: the charged-current b → cτ ν ratios RD , RD∗ and RJ/ψ ; the neutral-current b → sν ν̄ branching ratios, most notably the 2.8 σ excess in BR(B + → K + ν ν̄) reported by Belle II [18]; and the b → sℓ+ ℓ− lepton-flavour-universality ratios RK , RK ∗ . Following our previous global analyses [1–4], the short-distance NP is parameterised through the Wilson coefficients of the SMEFT, defined at the scale Λ = 1 TeV in the Warsaw basis and evolved down to the scale of each observable, and the fit is performed with the SMEFT19 package [19], which builds the global likelihood from more than a hundred observables by interfacing flavio [20], smelli [21] and wilson [22]. We work in Scenario III of the mass-rotation framework introduced in Ref. [4], in which the singlet and triplet four-fermion operators carry independent Wilson coefficients C1 ̸= C3 and no mixing is allowed in the lepton sector (αℓ = β ℓ = 0). The fit therefore depends on three parameters (C1 , C3 , β q ), with β q the quark mixing angle between the second and third generations. The relevant transitions probe different combinations of the coefficients: b → sν ν̄ receives tree-level contributions proportional to C1 − C3 , b → cℓν is governed by C3 , and b → sℓ+ ℓ− depends on C1 + C3 . Treating C1 and C3 as 5
40
r = 0.992 MAE = 1.85
0 0.5
1.0
1.5
2.0
σ
2.5
3.0
0.25
−20
Normalized frequency
Predicted ∆χ2SM
20
−40 −60 −80
validation set perfect prediction
−80
−60
−40
−20
0
20
Actual ∆χ2SM
Predicted histogram χ2 distribution
0.20 0.15 0.10 0.05 0.00 0.0
40
3.5
2.5
5.0
7.5
10.0
12.5
15.0
17.5
20.0
χ2bf − χ2
Figure 1: Validation of the B-physics emulator in Scenario III. (a) Predicted versus exact ∆χ2SM on the held-out validation set. (b) Distribution of ∆χ2 = χ2bf − χ2 for a Monte Carlo sample drawn from the emulated likelihood (blue), compared with the χ2 distribution for three degrees of freedom expected in the Gaussian (Wilks) limit (orange); the vertical lines mark the 1 σ and 2 σ levels. independent, unlike Scenarios I and II where C1 = C3 , is precisely what allows the fit to accommodate the enhanced B + → K + ν ν̄ rate together with the RD(∗) anomalies. The parameters are scanned over C1 , C3 ∈ [−0.30, 0.00] and β q ∈ [0.00, 3.20]. With respect to Ref. [4] we update the dominant charged- and neutral-current inputs to their most recent determinations. For the b → cτ ν ratios we use the HFLAV CKM 2025 world averages [17], RD = 0.358 ± 0.024 and RD∗ = 0.281 ± 0.011 with a correlation coefficient of −0.37, supplemented by the two RJ/ψ measurements, 0.71 ± 0.17 ± 0.18 from LHCb [15] and 0.49 ± 0.26 from CMS [23], which we include as independent constraints. For b → sν ν̄ we include the Belle II 362 fb−1 result BR(B + → K + ν ν̄) = (2.3 ± 0.7) × 10−5 [18]. The b → sℓ+ ℓ− ratios RK , RK ∗ , now in agreement with the SM, are kept at their reference values and do not move the fit, since lepton-flavour universality is exact in Scenario III [4]. We train an XGBoost surrogate of the SMEFT19 global log-likelihood over the three-dimensional space (C1 , C3 , β q ), following the methodology of Section 2 and first applied to this problem in Refs. [1, 4]. The training set comprises a base design of ∼ 1.4 × 104 points: a Sobol space-filling sample over the full domain, a Gaussian cloud drawn with the fit covariance around the best-fit point (which resolves the narrow C1 direction automatically), and three two-dimensional grids used to validate the likelihood maps. This base design is supplemented by a further ∼ 1.2 × 104 exact evaluations concentrated, through an active-learning refinement, in the weakly-constrained high-β q region where it is sparse, bringing the total to ∼ 2.6 × 104 points. Densifying this region improves the coverage of the extended β q ridge and mitigates the over-extrapolation of the surrogate where the training data would otherwise be scarce. Trained directly on the log-likelihood with a small learning rate and aggressive early stopping, the surrogate reaches a Pearson correlation coefficient r = 0.99 (R2 = 0.98) on a held-out validation set (Figure 1a) while evaluating each point far faster than the exact likelihood, as quantified below. The decisive advantage of the surrogate is its evaluation speed, quantified in Table 1. All times quoted here are measured on a single CPU core, so the comparison is reproducible and does not depend on the number of cores or on the specific machine. Once trained, the XGBoost ensemble predicts the log-likelihood in ∼ 34 µs per point, to be compared with ∼ 3 s for a single evaluation of the exact SMEFT19 global likelihood, which requires running the SMEFT-to-WET renormalization-group evolution and recomputing more than a hundred observables at each point. This is a speed-up of about 1.0 × 105 per evaluation. The gain is most tangible in the construction of the training set: 6
Per evaluation Training set (2.6 × 104 points)
XGBoost emulator 34 µs 0.88 s
Exact SMEFT19 3.3 s 24 core-h
Speed-up 1.0 × 105 1.0 × 105
Table 1: Evaluation cost of the XGBoost emulator against the exact SMEFT19 global likelihood in Scenario III, all on a single CPU core. Per-evaluation times use 104 surrogate predictions (batch) and a serial run of distinct exact evaluations; the second row is the cost of generating the full ∼ 2.6 × 104 point training set. Single-core figures make the speed-up reproducible and independent of the core count. generating the ∼ 2.6 × 104 exact points that define the emulator costs about 24 core-hours of SMEFT19 computation, whereas the trained surrogate re-evaluates the same number of points in ∼ 0.9 s. It is precisely this gain that makes the high-resolution confidence contours and the SHAP sensitivity analysis computationally feasible, while the agreement of the emulated likelihood maps with the exact ones (Figure 2) guarantees that the speed-up does not come at the price of accuracy. The fit locates the best-fit point at C1 = −0.185 ,
C3 = −0.110 ,
β q = 0.83 ,
(12)
with ∆χ2SM = 48.8, corresponding to a pull above 6 σ with respect to the SM. The fit clearly breaks the C1 = C3 degeneracy, confirming the qualitative picture of Scenario III with the updated data. Beyond the pointwise accuracy of Figure 1a, Figure 1b shows the distribution of ∆χ2 = χ2bf − χ2 for a Monte Carlo sample drawn from the emulated likelihood, together with the χ2 distribution for three degrees of freedom expected in the Gaussian (Wilks) limit. The sample follows the overall shape of this reference but is displaced towards somewhat larger ∆χ2 . This is not a deficiency of the emulator or the sampler but a direct imprint of the non-Gaussian structure of the likelihood: the quark-mixing angle β q is only weakly constrained and remains nearly degenerate over a broad range. This is visible in the two-dimensional log-likelihood maps of Figure 2, where β q develops an extended, curved ridge in the (C1 , β q ) and (C3 , β q ) planes, far from the compact Gaussian ellipse that a quadratic log-likelihood would produce. The χ23 curve should therefore be read as the Gaussian reference recovered in the asymptotic limit, not as an expectation the distribution is required to match; its controlled departure is precisely the non-Gaussian behavior the emulator is designed to capture. The same maps show excellent agreement between the exact SMEFT19 evaluation and the emulator across the (C1 , C3 ), (C1 , β q ) and (C3 , β q ) planes, confirming that the surrogate reproduces this structure faithfully. The SHAP analysis (Figure 3) ranks the parameters as C3 > C1 > β q . This hierarchy aligns with the underlying physics: the dominant sensitivity of the global fit comes from the charged-current RD(∗) observables (controlled by C3 ), followed by the neutral-current B → K (∗) ν ν̄ modes (controlled by C1 ), while the quark-mixing angle β q modulates both sectors without dominating. The agreement between the data-driven SHAP ranking and the expected operator structure confirms that the surrogate captures the correct EFT dynamics rather than overfitting numerical noise. Using a sample of 1.5 × 104 points generated with the emulator, we run the SMEFT-to-WET renormalization-group evolution at each point and compute the Pearson correlation coefficients between the resulting low-energy observables (Figure 4). Because no NP enters the lepton sector in µ Scenario III (αℓ = β ℓ = 0, so that C10 = 0 and RK (∗) = 1 exactly), the b → sℓ+ ℓ− ratios and the + − B(s) → µ µ branching ratios remain at their SM values and show no variation across the sample; the matrix therefore retains only the observables that genuinely respond to (C1 , C3 , β q ). The b → cτ ν ratios RD and RJ/ψ are perfectly correlated, reflecting their common origin in the single coefficient CVL ∝ C3 , while the neutral-current modes BR(B + → K + ν ν̄) and BR(K + → π + ν ν̄) respond to the orthogonal combination C1 − C3 . The central result concerns the interplay between RD and BR(B + → K + ν ν̄): the almost-perfect correlation between these two observables found in Scenario II [1] disappears (it drops to ≈ 0.06) once C1 and C3 are allowed to vary independently, since BR(B + → K + ν ν̄) ∝ (C1 − C3 ) while RD ∝ C3 . The lack of a fixed correlation between the enhanced B + → K + ν ν̄ rate and the b → cτ ν ratios is a distinctive prediction of Scenario III and a target for future experimental scrutiny. 7
20
0.00
20
−0.05
10
−0.05
10
−0.25 −0.2
C1
−0.1
0.0
2.5
−30
−0.30 −0.3
1.5
−0.2
C1
−0.1
0.0
−30
20
2.5
10
2.0
βq
0
−20
3.0
log L
0
1.5
1.0
−10
1.0
−10
0.5
−20
0.5
−20
−30
0.0 −0.3
0.0 −0.3
−0.2
C1
−0.1
0.0
20
3.0 2.5
0
1.5
C1
−0.1
0.0
−30
20
2.5
10
2.0
−0.2
3.0
10
2.0
βq
βq
−0.25
10
2.0
βq
−20
20
3.0
−10
−0.20
0
1.5
1.0
−10
1.0
−10
0.5
−20
0.5
−20
−30
0.0 −0.3
0.0 −0.3
−0.2
C3
−0.1
0.0
log L
−10
−0.20
0
−0.15
−0.2
C3
−0.1
0.0
log L
C3
−0.15
−0.30 −0.3
log L
0
−0.10
log L
C3
−0.10
log L
0.00
−30
Figure 2: Two-dimensional log-likelihood maps in the (C1 , C3 ), (C1 , β q ) and (C3 , β q ) planes (rows). Columns show the exact SMEFT19 evaluation (left) and the XGBoost emulator (right), both on the same 40 × 40 grid. The red star marks the best-fit point.
3.2
Light New Physics
If the NP candidate is light and feebly interacting, there is a variety of new challenges, both in the experimental and in the theoretical arenas, arising from the fact that the new particles are physical states detectable in experiments (as opposed to the heavy case, where the new particles appear only as virtual particles in internal lines of Feynman diagrams). In general, we can distinguish several kinds of experimental signatures [24–26]: • Off-shell processes where the corresponding on-shell process is forbidden by Lorentz invariance (e.g. leptonic and radiative meson decays M → X ∗ → ℓ+ ℓ− , γγ and neutral meson mixing M → X ∗ → M ). These processes are able to impose bounds for all values of mass and lifetime of the new particle, although they are most relevant in the resonant region mX = mM . • Prompt decays: the particle is produced in a reaction A → BX and subsequently decays into X → C (where A, B, C are states that may contain multiple physical particles), and the experiment is not able to resolve the production and decay vertices. The experimental signature is therefore an excess in the rate of A → BC as compared to the SM prediction. If the experiment is able to resolve a distance Lmin , then the probability for the new particle to travel less than
8
+7.64
C3 +5.73
C1 βq
+4.44 0
1
2
3
4
5
6
7
8
mean(|SHAP value|)
Figure 3: SHAP summary plot for Scenario III. Each point is a sample from the parameter space; its horizontal position gives the SHAP value (contribution to the log-likelihood) and the colour encodes the feature value. Features are ordered by mean absolute SHAP value. Lmin and decay promptly is Pprompt = 1 − exp
−Lmin cτX βX γX
,
(13)
where cτX = cℏ/ΓX is the proper decay length of the particle, and βX and γX its velocity and Lorentz boost in the lab frame. If the Narrow Width Approximation (NWA) is valid ΓX ≪ mX , then the NP prediction for this process is given by σ(A → BC)NWA = σ(A → BX) × Γ(X → C) × Pprompt .
(14)
• Displaced vertex decays: the particle is produced in A → BX and decays as X → C, but the experiment is able to resolve the production and decay vertices. This is the case of experiments with excellent vertex resolution, like LHCb, Belle II and BESIII. The dedicated displaced vertex analyses of these experiments provide bounds directly on the product σ(A → BX) × Γ(X → C) as a function of both mX and cτX . • Invisible decays: The new particle is produced in A → BX processes inside the detector, but it is very long-lived and decays outside the detector. The probability for a decay outside the detector is −Lmax Pout = exp , (15) cτX βX γX where Lmax is the maximum distance that the experiment is able to resolve. The experimental signature is a process with missing energy and momentum, which usually is interpreted as an excess in a process with neutrinos in the final state, σ(A → B + inv)X = σ(A → BX) × Pout .
(16)
Another possible way to obtain an invisible decay, that we will not explore in this work, happens if the new particle acts as a portal to an invisible sector and decays predominantly to noninteracting dark matter particles. From a purely computational point of view, the first difficulty that we encounter is the dependence of the predictions on cτX , that forces us to calculate every decay channel of X even for channels that are not directly observed. The exponential functions also induce non-Gaussianities in the likelihood function. Moreover, since some observables are only defined for certain values of the mass and/or the proper decay length of the particle, the likelihood function will not even be continuous at some points of the parameter space. All of this, together with the large number of model parameters and experimental observations at our disposal, make ML-based exploration particularly compelling.
9
1.00
BR(K + → π + νν̄)
-0.18
R`D
0.07
-0.18
-0.10
1.00
1.00
-0.19
-0.19
-0.24
0.50 0.25
1.00
-0.00
−0.25 −0.50 −0.75 −1.00
∆χ
J/
1.00
0.00
2 SM
-0.00
ψ
-0.00
R
π+ νν̄
-0.00
BR
(K
+
→
K+ νν̄
)
-0.24
)
-0.10
1.00
1.00
D
-0.19
1.00
R`
-0.19
0.07
∆χ2SM
(B + →
0.07
0.75
R J/ψ
BR
0.07
Pearson correlation coefficient
BR( B+ → K + νν̄)
Figure 4: Matrix of Pearson correlation coefficients between the flavour observables that vary in Scenario III, computed on the 1.5 × 104 -point Monte Carlo sample generated with the emulator. Observables that remain SM-like in Scenario III (b → sℓ+ ℓ− and B(s) → µ+ µ− ) carry no variation and are omitted. 3.2.1
Axion-like particles
One of the most popular and theoretically motivated examples of SM extensions featuring new light, feebly interacting particles are axions and axion-like particles (ALPs). Originally, axions were proposed as an elegant solution to the strong CP problem [27–29], where a pseudo-scalar particle a of mass ma and decay constant fa , acting as the pseudo-Nambu Goldstone boson (pNGB) of the Peccei-Quinn (PQ) U (1)PQ symmetry, dynamically drives the value of the electric dipole moment of the neutron to zero. However, the axion solution requires a correlation between its mass and decay constant, namely [28, 30] mu md m2a fa2 = m2 f 2 , (17) (mu + md )2 π π which severely constrains the available parameter space. Alternatively, if one chooses to consider ma and fa as independent parameters, we are instead in the ALP paradigm. While an ALP might not necessarily solve the strong CP problem, it has other interesting applications [31], like serving as a dark matter candidate [32–37], explaining the flavour puzzle (axiflavon/flaxion [38, 39]), neutrino masses (majoron [40–43]), or being generic predictions of string theory (axiverse [44, 45]). From a purely bottom-up perspective, we can study ALPs in the framework of EFTs [46], as an expansion in powers of 1/fa . The first terms of this expansion are the renormalisable operators for the kinetic and mass terms, and dimension-5 operators describing the lowest-order interactions with SM particles: derivative interactions with fermions that respect the shift-symmetry a → a + constant, and anomalous interactions with gauge bosons: X αV cV a 1 2 2 1 µ Vµν Ve µν Ld≤5 ALP−EFT =LSM + ∂µ a ∂ a − ma a + 2 2 4π fa V =B,W,G i ∂µ a h ij ij ¯ µ ij µ µ ij µ + cqL q̄i γ µ qj + cij ū γ u + c d γ d + c ℓ̄ γ ℓ + c ē γ e j j j j , uR i eR i dR i ℓR i fa 10
(18)
where Ve µν = 12 ϵµνρσ Vρσ is the dual tensor in the convention ϵ0123 = +1, q and ℓ are the SU (2)L doublets of left-handed quarks and leptons, and u, d and e are the SU (2)L singlets of right-handed quarks and charged leptons. In general, the couplings to fermions cij f are hermitian 3 × 3 matrices in flavour space. From the model building perspective, we can distinguish two very broad categories of UV-complete models, depending on which fermions are charged under the PQ symmetry. If only new heavy fermions are charged under PQ, as is the case in the KSVZ model [47, 48] and its generalizations [49–51], the ALP at the matching scale Λ will only interact with the SM gauge bosons as dictated by the gauge anomalies, unless the heavy and SM fermions mix. On the other hand, if the SM fermions fi have PQ charges Xfi , then the ALP at the matching scale will feature diagonal (but not necessarily universal) i couplings to fermions cii f (Λ) = −Xf , and couplings to gauge bosons fixed by the anomalies [49], 1X (2XqiL − Xui R − XdiR ) , 2 i 1X (3XqiL + XℓiL ) , cW (Λ) = − 2 i 1X i (XqL − 8Xui R − 2XdiR + 3XℓiL − 6XeiR ) . cB (Λ) = − 6 cG (Λ) = −
(19)
i
Note that in specific realizations, like DFSZ model [52, 53] and its generalizations [49, 54, 55], or flaxion models, the scalar sector will impose constraints to combinations of the PQ charges. Once the ALP couplings have been determined at the matching scale Λ ∼ 4πfa , we must solve the Renormalization Group Equations (RGEs) and integrate out the heavy SM particles (t, h, Z, W ± ) in order to obtain the values of the coefficients at the physical scale of each observable [56–62]. One consequence of the running and matching is that in models with UV couplings to top quarks and/or W ± , sizable off-diagonal couplings in the down sector will be generated. 3.2.2
ALPs as a solution to the Belle II anomaly
The Belle II experiment has found an excess in the branching ratio of B + → K + ν ν̄ of 2.8 σ with respect to the SM prediction [18]. Several analyses of the q 2 distribution [63–67] have shown that the excess is compatible with the production of one single particle of mass around 2 GeV that does not decay inside the detector. Furthermore, by combining also the upper limit for the process B → K ∗ ν ν̄ at BaBar [68], with different detector size and Lorentz boost, Ref. [26] found a lower bound for the proper decay length of the particle, cτ ≥ 80 cm (95% C.L.). The identification of this presumptive 2 GeV particle with an ALP proves to be challenging when we consider scenarios motivated by UV-complete models. On the one hand, the flavour-violating coupling |cVsb |/fa that enters in the ALP production B + → K + a needs to be large enough to explain the excess, but on the other hand, ALP decay channels such as a → µ+ µ− and a → ηππ need to be suppressed in order to respect the constraint on the ALP lifetime. Ref. [26] identified a model that fulfills all the requirements, based on the astrophobic axions of Refs. [69–71]. The PQ charges of the SM fermions are XqL = XℓL = (0, 0, −1),
XuR = − sin2 β ,
XdR = XeR = − cos2 β .
(20)
In this construction, the coupling cVsb is generated by top loops, and it is largely independent of β. A(0) Meanwhile, the coupling to muons receives a tree level contribution cµµ = cos2 β which is partially 2 2 canceled by the running of the coupling to τ leptons cA τ τ = cos β − 1 = − sin β, and consequently it is minimized for intermediate β angles. 3.2.3
Fit parameters and results
The model in [26] is rather predictive, as it only depends on two free parameters, β and fa . Here we aim to extend it in a more descriptive direction, exploring a larger region of the parameter space by 11
7.5 14.00 13.75
7.0
6.5
13.25 13.00
6.0
χ2
log10 f a [GeV]
13.50
12.75 12.50
5.5 12.25
5.0
12.00 1.7
1.8
1.9
2.0
2.1
2.2
m a [GeV]
Figure 5: Points of the training dataset in the plane ma -log10 fa . The active learning algorithm ensures that the density of points is larger in the areas of small χ2 (darker colours), while still covering the whole parameter space. lifting the dependence relations between the PQ charges. The free parameters of our new model will be the PQ charges of the third-generation left-handed fermions Xq3L and Xℓ3L and the universal PQ charges of the right-handed fermions XuR , XdR and XeR , all of them allowed in the range Xf ∈ [−1, 1], as well as the energy scale fa in the range fa ∈ [105 , 108 ] GeV (while fixing Λ = 4πfa ) and the ALP mass in the range ma ∈ [1.7, 2.2] GeV. We perform the fit of this model using the likelihood function implemented by the ALP-aca library [72]. This likelihood function includes more than 100 measurements, but not all of them are operative depending on the ALP mass and proper lifetime. In our setup, in addition to B → K (∗) a(→ inv) processes, we have to consider off-shell ALP processes (leptonic and radiative meson decays, neutral meson mixing), and in the limit of short lifetime, prompt/displaced vertex processes like B → K (∗) a(→ µ+ µ− ) decays. The training dataset is generated following an active learning programme in order to improve the coverage of the region of large likelihood. A small sample (800 points) generated by LHS is used as a starting seed. Then, at each iteration of the active learning algorithm, a Gaussian Process (GP), the infinite-dimensional limit of a multivariate Gaussian, is trained on the dataset up to that point. A set of candidate points are proposed, for which the GP predicts its χ2 , and more importantly, the uncertainty of said prediction. The candidate points are scored according to their modified Expected Improvement 2 2 χ∗ − µ − ξ χ∗ − µ − ξ 2 EI = (χ∗ − µ − ξ) Φ +σ ϕ , (21) σ σ where µ and σ are the predicted central value and uncertainty of the χ2 for the point, χ2∗ is the minimum χ2 observed so far, and Φ and ϕ are the cumulative distribution function and probability distribution function of the Gaussian distribution, respectively. The parameter ξ regulates the tradeoff between exploration and exploitation: if ξ is large, the best score will be achieved by points where the GP is very uncertain, while for small ξ, points with small predicted χ2 are selected instead. At each iteration of the algorithm, the point with the best score is selected, its true χ2 is computed and added to the training dataset, and ξ is slightly decreased. The total size of the dataset, including both the initial seeds and the points selected by the GP, is 2500 points. The training dataset is depicted in Figure 5, showing that the highest density is achieved in the region of low χ2 . In order to fully capture the details of the χ2 function in the region around its minimum, we propose a two-step strategy: we first train a XGBoost instance on the sigmoid of the displaced χ2 function, 1 2 2 , (22) σ(χ2 ) = −χc 1 + exp χ 0.4 12
22
1.0 1.0
r = 0.929
20
0.8
0.6
0.4
0.4
18
0.2
Predicted χ2
0.6
σ(χ2 )
True Positive Rate
0.8
16
0.2
14
0.0
12
validation dataset perfect prediction
0.0 0.0
0.2
0.4 0.6 False Positive Rate
0.8
1.0
12
14
16
18
20
22
Actual χ2
Figure 6: Left: Receiver Operating Characteristic (ROC) curve of the first XGBoost trained on the ALP dataset. Right: Predicted versus exact χ2 on the held-out validation dataset. Per evaluation Training set (2.5 × 103 points)
XGBoost emulator 3.5 µs 0.009 s
Exact ALP-aca 1.85 s 1.3 core-h
Speed-up 5.3 × 105 5.3 × 105
Table 2: Evaluation cost of the XGBoost surrogate against the exact ALP-aca likelihood, on a single CPU core. As in the B-physics benchmark, per-evaluation times use 104 surrogate predictions (batch) and a serial run of distinct exact evaluations; the second row is the cost of generating the training set. where χ2c = χ2b.f. + 10 is the cutoff: this first stage acts as a classifier, retaining only the points below the cutoff. The second stage is a XGBoost instance trained only on those points below the cutoff, thus focusing on the relevant region of the parameter space. The performance of the classifier stage is characterized by its Receiver Operating Characteristic (ROC) curve in Figure 6. It presents the True Positive Rate (proportion of the points below the cutoff that are correctly identified by the classifier) and False Positive Rate (proportion of the points above the cutoff that are misidentified) for various thresholds of σ(χ2 ). A random classifier would lie on the diagonal of the ROC, while a perfect classifier would be situated in the upper left corner. By setting the threshold for the classifier at σ(χ2 ) = 0.5, we achieve TPR = 0.96 and FPR = 0.16 on the validation dataset. The Area Under the Curve (AUC) of the ROC plot is AUC = 0.974, also indicating a great classifying performance. In the regressor stage, the surrogate achieves a Pearson correlation coefficient r = 0.929, as shown in Figure 6 (Right). As in the B-physics case, the decisive advantage of the surrogate is its evaluation speed, quantified in Table 2. On a single CPU core, one evaluation of the exact ALP-aca likelihood over the seven-dimensional parameter space takes ∼ 1.9 s (it must solve the RGEs, perform the matching and recompute more than a hundred observables), whereas the trained XGBoost ensemble predicts the same point in ∼ 3.5 µs, a speed-up of about 5.3 × 105 . The gain is most tangible in the construction of the training set: generating the 8 × 103 exact points costs about 4.1 core-hours of ALP-aca computation, which the emulator re-evaluates in ∼ 0.03 s. As in the flavour benchmark, all times are measured on a single core, so the comparison is reproducible and does not depend on the number of cores or on the specific machine. The SHAP values help us again to interpret the results of the fit. In Figure 7 (left), we have displayed the mean of absolute values of SHAPs in the dataset, which corresponds to the expected importance of each parameter. The most important parameter is ma , since it determines if a solution to the Belle II anomaly is even kinematically allowed in the first place. After that, the two most important parameters are Xq3L and fa , which control the top loops and the overall impact of ALP effects. The impact of the PQ charges of right-handed fermions is more limited. In the case of ma , we can see in Figure 7 (right) that larger ALP masses, which are less compatible with the q 2 distribution of the recast to two-body decay kinematics, result in larger SHAP values, that is, larger contribution
13
3.0 2.5
Xq33L
2.0
SHAP value for m a
ma
log10 f a
X`33 L XeR Xu R
1.5 1.0 0.5 0.0
−0.5
Xd R
−1.0 0.0
0.2
0.4
0.6
1.6
0.8
1.7
1.8
1.9
2.0
2.1
2.2
ma
mean(|SHAP value|)
Figure 7: (left) Global SHAP feature importance (mean absolute SHAP value, i.e. each parameter’s average contribution to the emulated χ2 ) for the parameters of the ALP fit. (right): Point-wise SHAP values for ma . to the χ2 function. Finally, using the surrogate of the likelihood function and a Markov chain Monte Carlo (MCMC) implemented by the emcee library, we can efficiently sample the posterior distributions of both the parameters of the model and of derived quantities. In Figure 8, we have obtained the posterior distributions for the input parameters ma and fa , for the low-energy couplings |cVsb |, |cA µµ | and |cG | (all of them at the physical scale µ ∼ ma ) and the relevant parameters for the Belle II anomaly: cτa and BR(B → Ka). The posterior distribution correctly reflects functional dependencies, such as cτa ∝ fa2 and BR(B + → K + a) ∝ |cVsb |2 . Clearly, results favor an ALP mass around 1.8 GeV, in agreement with previous analyses of the Belle II excess based on the kinematic reconstruction of the missing-energy signal. The preferred values of fa lie in the 106 −107 GeV range, balancing the need for sufficiently large flavour-changing interactions to explain the anomaly against the suppression required by constraints included in the fit. The posterior distributions of the effective couplings show that |cVsb |2 must be large enough to generate an observable BR(B + → K + a), while |cA µµ | remains comparatively small in order to suppress visible ALP decays and ensure compatibility with searches for displaced and prompt dimuon resonances. In Figure 9, we have shown again the posterior obtained by the surrogate and MCMC, this time fixing ma = 1.8 GeV, for cτa and BR(B + → K + a) (in purple), compared to the fit using only the B → K + inv data alone. The BaBar and Belle II data was only able to set a lower bound for the decay length of the ALP, but the combination with other observables in the fit provides complementary information that further restricts the proper decay length. As it was the case in the SMEFT analysis, we once more underline the importance of global explorations of the parameter space, which hugely benefit from the ML techniques like the ones we have described.
3.3
Cosmology
The last decades have seen cosmology become precision physics. And in this path, the methods used in order to do parameter inference and data analysis have become more and more similar to to the ones used in high energy physics. There are different types of cosmological probes when doing cosmological parameter inference. The space-time geometry can be determined through the distanceredshift relation through time-domain probes or through the study of the large-scale structure of the Universe. The use of these probes has allowed us to define a baseline cosmological standard model, the flat ΛCDM, but some tensions between observations may indicate the need for NP [73]. In all cases, the likelihood is again built from a χ2 comparing a handful of cosmological parameters against hundreds of measurements through theoretical predictions that require numerical distance integrals, so the same obstacles encountered in the flavour fits (expensive forward evaluations and curved parameter degeneracies) reappear here and make the ML emulation strategy of Section 2 equally attractive. The main cosmological probes used in this work are the distance–redshift relation, measured with Type Ia supernovae as standard candles and with baryon acoustic oscillations as a standard ruler, 14
0.46 log10 f a [GeV] = 6.92+ −0.36
95 1. 65
1.
80
m a [GeV]
2.
10
0.13 m a [GeV] = 1.81+ −0.08
2 1
100 × |cV sb |
3
4
1.
+1.42 100 × |cV sb | = 0.78−0.60
A | |cµµ
0.
3
0. 6
0.
9
1.
2
A | = 0.20+0.35 |cµµ −0.14
0. 8
1. 6
|cG |
2.
4
3. 2
0.95 |cG | = 0.95+ −0.68
3. 0 1. 5 0.
0
log10 (cτ ) [m]
0.92 log10 (cτ ) [m] = 1.17+ −0.82
− 5 7.
− .0 10
− .5 12
m a [GeV]
log10 f a [GeV]
100 × |cV sb |
A | |cµµ
log10 (cτ ) [m]
0
5
5.
−
.0
7.
−
10
−
.5 12
−
.0 15
−
0
5
3.
0
|cG |
1.
0.
2
4
3.
6
2.
8
1.
0.
2
9
1.
6
0.
3
0.
0.
4
3
2
1
10
95
2.
80
1.
1.
0
65 1.
5
8.
0
7.
7.
5 6.
−
15
.0
−
log10 BR( B+ → K + a)
5. 0
+1.24 log10 BR( B+ → K + a) = −6.94− 1.27
log10 BR( B+ → K + a)
Figure 8: Posterior distribution of ALP parameters and derived quantities.
10−2
m a = 1.8 GeV Belle II
10−3
BR( B → K a)
10−4
10−5
10−6 BaBar 10−7 ML fit 10−8
101
102
103 cτa [cm]
104
105
Figure 9: Allowed region in the cτa − BR(B → Ka) plane determined by the surrogate-based fit including all observables (purple) compared to the allowed regions by BR(B → K + inv) alone (BaBar in green, Belle II in blue). The ALP mass is fixed to its best fit value.
15
complemented by the Cosmic Microwave Background (CMB) as a prior on cosmological parameters. 3.3.1
Cosmological probes and the likelihood
The baseline model for cosmological analysis is the flat ΛCDM model where the accelerated expansion of the Universe is explained by the cosmological Λ with constant density and spatial flatness at each moment of cosmological time. The expansion history of the Universe is encoded in the dimensionless Friedmann equation: E(z) ≡ H(z)/H0 , h i1/2 , E(z) = Ωm (1 + z)3 + (1 − Ωm ) fDE (z)
(23)
where spatial flatness fixes the dark-energy density parameter to 1 − Ωm and radiation is negligible at the redshifts probed. The dynamics are driven by the parameters: H0 is the Hubble constant, Ωm = Ωb + Ωcdm is the density parameter of matter in the Universe, comprised by ordinary matter (Ωb ) and cold dark matter (Ωcdm ). The dynamical nature of dark energy is enclosed in fDE (z). For the ΛCDM model, as the energy density of the cosmological constant component does not evolve with time (ρΛ ∼ const.), fDE = 1. In the last decade, different cosmological probes, old and new, seem to converge in that the most plausible model for cosmology was ΛCDM. But recently, supernovae and standard ruler measurements have indicated a possible dynamical dark energy component. The evidence for dynamical dark energy is based upon the flat w0 wa CDM dark-energy model, in which the dark-energy equation of state is allowed to evolve with redshift according to the Chevallier–Polarski–Linder (CPL) parameterisation [74, 75], w(z) = w0 + wa
z , 1+z
(24)
which reduces to a cosmological constant (ΛCDM) for w0 = −1, wa = 0. In this particular case, the energy density of dynamical dark energy component evolves according to: fDE (z) = (1 + z)3(1+w0 +wa ) e−3wa z/(1+z) .
(25)
All geometric observables follow from common definitions. The transverse distance between two objects at a given redshift is the comoving (transverse) distance, DM (z), while the radial geometric distance relation with redshift is the Hubble distance, DH (z), and finally we can capture the isotropic nature of a spherical ruler with the spherically averaged distance, DV (z), [76] Z z 1/3 c dz ′ c c 2 DM (z) = , DH (z) = = , DV (z) = z DM (z) DH (z) , (26) ′ H0 0 E(z ) H(z) H0 E(z) so that the model depends on the four parameters θ = (Ωm , H0 , w0 , wa ) (reduced to (Ωm , H0 ) in the ΛCDM baseline below). The three probes enter the global likelihood as follows: • Type Ia Supernovae are standardisable candles: after light-curve correction their peak brightness is fixed, so the observed apparent magnitude mB measures the luminosity distance through the distance modulus µ(z) ≡ mB − M = 5 log10 dL (z)/10 pc = 5 log10 dL (z)/Mpc + 25 , dL (z) = (1 + z) DM (z) ,
(27)
with M the absolute magnitude. Because M is fully degenerate with H0 , it is analytically marginalised, so the supernovae constrain the shape of the distance–redshift relation rather than its absolute scale. The χ2 uses the full statistical+systematic covariance matrix C, χ2SN = ∆T C −1 ∆ with ∆i = µobs − µ(zi ). We use the Pantheon+ / SH0ES compilation [77, 78] (1657 i supernovae, with the Cepheid calibrators excluded so that M is marginalised) and, as an independent sample, the Dark Energy Survey five-year SN Ia data set (DES-SN5YR), in both its 2024 release (1829 supernovae) [79] and its 2025 update [80] (1820 supernovae); for DES the magnitude offset is marginalised analytically, including the corresponding normalisation term. 16
• Baryon Acoustic Oscillations imprint a fixed comoving scale, the sound horizon at the baryon drag epoch rd , that as a standard ruler in the galaxy distribution. The sound R ∞acts s (z) horizon is given by: rd = zd cH(z) dz, where zd is the redshift at decoupling and the speed of sound prior to recombination, cs (z), is set by the baryon, ρb and radiation, ργ , densities as r −1 3ρb (z) cs (z) = c 3 1 + 4ργ (z) . Surveys measure the ratios DM /rd , DH /rd and the isotropic combination DV /rd at several effective redshifts. We use the 13 measurements of the DESI −1 Data Release 2 [81] with their Gaussian likelihood χ2BAO = ∆T CBAO ∆, fixing rd to its Planck fiducial value. • Cosmic Microwave Background. Rather than recomputing the full CMB power spectra at every point, we summarise the Planck constraint as a Gaussian prior on the late-time parameters. The Planck PR4+lensing chain in the CPL model obtained in the re-analysis of the latest Planck data [82] (superseding the official PR3 [83]) made by [84] provides a covariance matrix for its full set of about 25 parameters; of these we keep only the five that map onto our fit, the Hubble constant H0 , the two physical densities of baryons ωb and cold dark matter ωcdm (with ωx ≡ Ωx h2 and h = H0 /100), and the dark-energy parameters w0 and wa , discarding the remaining nuisance columns. Because our sampling variable is the total matter density rather than its two components, we combine them into ωm = ωb + ωcdm by linear error propagation, which collapses the 3 × 3 block of (H0 , ωb , ωcdm ) into a correlated 2 × 2 covariance for (H0 , ωm ); the dark-energy parameters are carried with their Planck variances. The result is a 4 × 4 covariance ΣCMB for (H0 , ωm , w0 , wa ) in which H0 and ωm stay correlated, while w0 and wa enter as independent Gaussian constraints. Centred on the Planck ΛCDM best fit (w0 = −1, wa = 0), it contributes ⋆ ∆ = H0 − H0⋆ , ωm − ωm , w0 − w0⋆ , wa − wa⋆ , (28) χ2CMB = ∆T Σ−1 CMB ∆ ,
where ωm = Ωm (H0 /100)2 is evaluated from the sampled parameters and the starred quantities are the Planck central values. Combining the two densities by propagating their covariance is the exact marginalisation of the joint Gaussian onto (H0 , ωm ), so for observables that depend only on the total matter density it is equivalent to carrying ωb and ωcdm as explicit parameters; the correlated (H0 , ωm ) block is retained, while in this particular covariance w0 and wa happen to be uncorrelated with the rest and enter as independent Gaussians of width σ(w0 ) = 0.02, σ(wa ) = 0.05. The same prior is used in every CMB combination of Sections 3.3.4 and 3.3.5. P Combined data sets are fit by summing the individual χ2 functions, χ2tot (θ) = k χ2k (θ), exactly as in the flavour global fit. 3.3.2
Surrogate construction and the choice of training target
For each data combination we train an XGBoost surrogate of χ2tot (θ) following the methodology of Section 2. The training design combines two-dimensional slices through the best-fit point, a uniform space-filling draw over the full prior box (Ωm ∈ [0.1, 0.9], H0 ∈ [20, 100], w0 ∈ [−3, 0.2], wa ∈ [−3, 2]), and a Gaussian cloud whose covariance is the Hessian of the fit, so that the narrow, anisotropic w0 –wa degeneracy is sampled densely where it matters most. Each design contains 2.5–5 × 105 points, with the exact distance integrals of Eqs. (25)–(26) evaluated in parallel across CPU cores. A central, problem-specific ingredient is the choice of training target. The cosmological χ2 spans an enormous dynamic range, from χ2min ∼ 103 at the best fit to ≳ 106 at the edges of the prior box, while the physically relevant structure, the ∆χ2 ∼ 2–6 band that defines the 1 σ and 2 σ confidence regions, lives in a tiny sliver near the minimum. Crucially, the absolute scale of χ2min differs by two orders of magnitude between the data combinations, from χ2min ≈ 1.7 × 103 for the supernova+BAO fits (more than 1600 data points) down to χ2min ≈ 5.6 for the BAO-only fit (13 points), and this is what makes the two naive targets inadequate in a data-dependent way: • Training directly on χ2 lets the loss be dominated by the high-χ2 tail: the surrogate fits the outer wall of the likelihood well but is blind to the O(few) variations near the minimum. This happens for every combination, irrespective of χ2min , and makes the recovered contours unusable. 17
• Training on a plain logarithm, log10 χ2 , compresses the tail, but whether it works now depends on where the minimum falls on the logarithmic curve, which is precisely why it succeeds for some data sets and fails for others. For the supernova+BAO and +CMB fits the minimum is large, so it sits on the locally flat part of the logarithm and the decisive 1 σ–2 σ band maps onto a negligible change of the target (log10 1706 − log10 1700 ≈ 1.5 × 10−3 ): the contours remain unresolved. For the BAO-only fit the minimum is small, so it sits on the steep part of the logarithm (log10 12 − log10 5.6 ≈ 0.33) and the plain transform does resolve the contours. No single un-shifted logarithm can therefore serve all combinations at once. Both problems are solved by a shifted-log10 transform, which is the target adopted throughout this work: y(θ) = log10 χ2 (θ) − χ2min + 1 , χ2 (θ) = 10 y(θ) − 1 + χ2min , (29)
where χ2min is the minimum over the training set. The shift maps the best fit to y = 0 regardless of the absolute value of χ2min , so that the minimum of every likelihood, whether χ2min ≈ 5 or ≈ 1700, is placed on the same steep, well-resolved part of the transform; this is what removes the data-dependence of the plain logarithm discussed above. The unit offset additionally guarantees a strictly positive argument (≥ 1), removing the log(0) singularity that a naive shift χ2 − χ2min would introduce exactly at the densely sampled best-fit point. The transform compresses the 106 tail and stretches precisely the near-minimum band that encodes the confidence regions, amplifying it by about three orders of magnitude in target space. The surrogate is trained on y and its predictions are mapped back to linear χ2 through the exact inverse in Eq. (29), so that downstream contour and SHAP computations operate on physical χ2 values. With this target the emulators reach a coefficient of determination R2 = 0.997–0.9999 on a held-out validation set across all data combinations. Posterior distributions are then mapped by running the parallel Random-Walk Metropolis–Hastings sampler of Section 2 (1024 chains, with the proposal covariance set to the Hessian of the fit) directly on the emulated likelihood, and rendered as marginalised contours. Because each likelihood call is a single forward pass through the tree ensemble rather than a set of numerical distance integrals, the cost of the posterior exploration is reduced by more than an order of magnitude, as quantified in Section 3.3.5. 3.3.3
Baseline: the ΛCDM fit
As a simple, fully controlled example we first emulate the two-parameter ΛCDM likelihood of Pantheon+ combined with DESI BAO, fitting only (Ωm , H0 ). The surrogate reproduces the exact χ2 with R2 = 0.99988, and the emulator MCMC recovers Ωm = 0.310 ± 0.008 ,
H0 = 68.4 ± 0.5 km s−1 Mpc−1 ,
(30)
at χ2min = 1703.6. Figure 10 shows the resulting Ωm –H0 posterior: a single, closed, mildly anticorrelated ellipse with no secondary modes or open directions. These values reproduce the standard concordance picture, Ωm = 0.310 matching the Planck and DESI determinations to within 1 σ and the supernova-only Pantheon+ value Ωm ≈ 0.33. Because the supernova absolute magnitude is marginalised, the fit carries no SH0ES distance-ladder anchor, so H0 is set by the BAO scale with the Planck sound horizon rd ; it therefore lands on the Planck-like value H0 ≈ 68 rather than the higher local-distance-ladder measurement, exactly as this dataset combination behaves in the DESI and Pantheon+ analyses. This baseline confirms that the pipeline reproduces the established result before turning to the more demanding, and more interesting, four-parameter dark-energy case. 3.3.4
Dynamical dark energy: w0 wa CDM with and without CMB
We now turn to the full w0 wa CDM model and study how the dark-energy constraints depend on the supernova sample and on the inclusion of the CMB prior. We consider three supernova compilations (Pantheon+, DES-SN5YR 2024 and 2025), each combined with DESI BAO, and contrast the fit without and with the Planck prior. The best-fit parameters of the six emulators are collected in Table 3, and the corresponding posteriors are compared pairwise in Figure 11. 18
Figure 10: Marginalised Ωm –H0 posterior of the ΛCDM baseline (Pantheon+ & DESI BAO), obtained by sampling the XGBoost surrogate (R2 = 0.99988) with the Random-Walk Metropolis–Hastings sampler. The inner and outer contours are the 1 σ and 2 σ credible regions. Data set Pantheon+ & BAO Pantheon+ & BAO & CMB DES-2024 & BAO DES-2024 & BAO & CMB DES-2025 & BAO DES-2025 & BAO & CMB
Ωm 0.323 0.310 0.322 0.309 0.314 0.307
H0 66.9 67.8 67.0 67.8 67.6 68.0
w0 −0.766 −0.963 −0.782 −0.964 −0.843 −0.970
wa −0.785 0.017 −0.726 0.020 −0.536 0.018
χ2min 1686.9 1702.3 1654.6 1667.3 1647.0 1655.3
R2 0.9978 0.9990 0.9976 0.9991 0.9976 0.9991
σ(wa ) 0.42 0.045 0.45 0.045 0.45 0.045
Table 3: Best-fit parameters of the w0 wa CDM emulators for the three supernova samples combined with DESI BAO, without and with the Planck CMB prior. Each row is a separate XGBoost surrogate; R2 is the validation score and σ(wa ) the marginalised 1 σ uncertainty on wa . Adding the CMB prior pulls every fit back towards ΛCDM (w0 = −1, wa = 0) and shrinks the dark-energy uncertainties by roughly an order of magnitude. A consistent pattern emerges. Without the CMB prior, all three supernova+BAO combinations prefer a dynamical dark energy with w0 > −1 and wa < 0, i.e. an equation of state that crosses the phantom divide as it evolves, the same hint of evolving dark energy recently reported by DESI. The associated contours are broad, with σ(wa ) ≈ 0.4, reflecting the well-known w0 –wa geometric degeneracy that the surrogate captures as an elongated, curved ridge. The deviation is strongest for Pantheon+ and DES-2024 and milder for the DES-2025 update, which sits between the two. Adding the Planck prior pulls every fit back towards ΛCDM (w0 ≈ −0.96 to −0.97, wa ≈ 0.02) and tightens the constraints dramatically: the uncertainty on wa shrinks by about a factor of ten (from ∼ 0.42 to ∼ 0.045), on w0 by a factor of five, and on Ωm by a factor of three. Figure 11 makes this contraction visually explicit for the three samples. Reading the Pantheon+ panel directly (the DES panels are qualitatively identical), the supernova+BAO posterior (blue, filled) is a wide, tilted ellipse in the w0 –wa plane, elongated along the anti-correlation direction and displaced from the ΛCDM point (dashed lines) towards w0 > −1, wa < 0, whereas the +CMB posterior (red) collapses to a small blob centred on (w0 , wa ) = (−1, 0); the Ωm and H0 marginals, already well determined, barely move. The supernova+BAO best fits, for instance 19
Figure 11: w0 wa CDM posteriors without and with the Planck CMB prior, for (a) Pantheon+, (b) DES-SN5YR 2024 and DES-SN5YR 2025, each combined with DESI BAO. Dashed lines mark the ΛCDM point (w0 , wa ) = (−1, 0). (w0 , wa ) ≈ (−0.78, −0.73) for DES-2024, fall in the same region and direction of the w0 –wa plane as the dynamical-dark-energy preference reported by the DESI DR2 analysis for the analogous DESI+supernova combinations, so the emulator reproduces that published result. The CMB enters here not as a forward computation of the temperature and polarisation spectra, but as a compressed Gaussian constraint built from the Planck PR4+lensing covariance matrix (Eq. 28), a standard and economical way to fold the CMB into a late-time fit. It carries the Planck constraints on all four parameters, with the H0 –ωm correlation retained, and in particular the tight bounds on the dark-energy sector that the supernovae and BAO leave nearly free; this is what allows it to tighten w0 and wa by up to an order of magnitude and to draw the joint fit into agreement with ΛCDM. The relevant point for the present work is that the surrogate reproduces this compressed-CMB likelihood as faithfully as the supernova+BAO one, so the contraction of the contours is a genuine output of the emulated fit and not an artefact of the emulation. 3.3.5
Computational benchmark: surrogate versus exact likelihood
The cosmological application also offers a clean setting to quantify the computational gain of the ML emulation, because here the exact likelihood is cheap enough to be sampled directly for comparison, unlike the flavour case, where a brute-force scan is prohibitive. We benchmark three nested data combinations of increasing constraining power, BAO only, BAO & Pantheon+ and CMB & BAO & Pantheon+, which illustrate how each probe closes the w0 –wa degeneracy: with only the 13 BAO points the four-parameter fit is essentially unconstrained (χ2min = 5.6, with wa railing against the prior edge), adding the supernovae closes the contour, and the CMB prior tightens it onto ΛCDM (Figure 12). We run the same Random-Walk Metropolis–Hastings sampler (1024 chains, identical settings and seed) with four interchangeable likelihood back-ends: the XGBoost surrogate evaluated on the GPU, and the exact χ2 computed with (i) the astropy-based reference parallelised over a CPU process pool (“Astropy”), (ii) a vectorized NumPy reimplementation of the distance integrals on the CPU (“CPU”), and (iii) a JAX vmap kernel on the GPU (“GPU”) [85]. All four back-ends are run on the same workstation, an 18-core/36-thread Intel Core i9-10980XE CPU (3.0 GHz base, 4.8 GHz turbo) with 188 GB of RAM and an NVIDIA RTX A6000 GPU (48 GB). The Astropy reference must be evaluated one cosmology at a time and cannot be vectorized across the sampled points; its only 20
Figure 12: w0 wa CDM posteriors for the three nested data combinations, BAO only (blue), BAO & Pantheon+ (red) and CMB & BAO & Pantheon+ (green), sampled from the XGBoost surrogate (left) and from the exact likelihood (right). The two are statistically indistinguishable, as are the four likelihood back-ends of Table 4.
Data set BAO only BAO & Pantheon+ CMB & BAO & Pantheon+
ML 7.4 s 6.8 s 4.8 s
Exact χ2 Astropy CPU 98.5 s 26.9 s n/a 151.4 s n/a 98.6 s
GPU 89.9 s 82.1 s 34.3 s
Speed-up ×CPU ×GPU 3.6 12.1 22.4 12.1 20.8 7.2
Table 4: Wall-clock time of the posterior exploration (1024 chains, identical sampler settings) for the three nested data combinations, comparing the XGBoost surrogate against the exact likelihood evaluated with the astropy reference (Astropy, process pool), with a vectorized NumPy kernel (CPU), and on the GPU with JAX (GPU). The last two columns give the speedup of the surrogate over the CPU and the GPU exact likelihoods. The Astropy process pool is reported only for the BAO-only case (it was impractical for the large supernova likelihoods). parallelism is to spread those points over a process pool (up to 16 workers), which scales poorly and becomes impractical for the large supernova likelihoods (the n/a entries in Table 4). This is precisely why the distance integrals are reimplemented directly as the vectorized NumPy (CPU) and JAX (GPU) back-ends, which evaluate the whole batch of points in a single call. Crucially, all four back-ends produce statistically identical posteriors, confirming that the surrogate is accurate and not merely fast. The wall-clock times are summarised in Table 4: the surrogate completes each run in 4.8–7.4 s, 3.6–22.4× faster than the CPU (vectorized NumPy) likelihood and 7.2–12.1× faster than the exact GPU kernel, with the advantage growing for the more expensive Pantheon+ likelihoods. The structure of Table 4 reflects two competing effects. The emulator’s per-evaluation cost is essentially independent of the dataset, since a forward pass through the tree ensemble does not depend on the number of observables, whereas the exact likelihood scales with the data; this is why the speed-up over the CPU back-end grows from ∼ 4× for the 13 BAO points to ≳ 20× for the ∼ 1700-supernova Pantheon+ fits. The GPU exact kernel behaves oppositely at small data volume: for BAO only its fixed kernel-launch and host–device-transfer overhead dominates the tiny per-step compute, so it is slower than the CPU (vectorized NumPy) code (90 vs 27 s), whereas for the much heavier Pantheon+ likelihood that overhead is amortised and the GPU wins. Finally, because every chain is run to a
21
fixed effective-sample-size target rather than a fixed number of steps, the wall-clock time also tracks how fast each posterior mixes: the tightly CMB-constrained fit converges in fewer steps, so BAO & Pantheon+ & CMB is cheaper than BAO & Pantheon+ despite containing more data. The same acceleration benefits the up-front cost of building the training set. Evaluating the 2 × 105 exact χ2 points needed to train a supervised surrogate takes about 4 × 103 s (more than an hour) with the astropy reference on 16 CPU cores for the Pantheon+ combinations, whereas the batched JAX/GPU implementation of the distance integrals produces the same set in a few seconds. The GPU and Astropy evaluations agree to better than one unit in χ2 near the minimum, where the training signal matters; they differ only in the extreme tails (χ2 ∼ 106 ), a harmless discrepancy since those points sit in the compressed part of the shifted-log transform of Eq. (29). Thus both the construction and the sampling of the emulated likelihood are accelerated, and the one-time training cost is amortised over the whole posterior analysis, after which the surrogate replaces the physics computation entirely. 3.3.6
Feature importance
As in the flavour analysis, SHAP values provide a transparent, model-independent ranking of the parameters by their mean absolute contribution to the emulated χ2 , that is, by how steeply the likelihood responds to each one across the sampled volume, which is a direct proxy for how tightly each parameter is constrained. The rankings (Figure 13) follow the expected physics and, tellingly, reorganise when the CMB prior is added. In the ΛCDM baseline the two parameters are almost tied, Ωm (0.67) just above H0 (0.63): both are comparably and tightly determined by the joint supernova+BAO geometry, with Ωm marginally ahead because it enters every distance through E(z). In the w0 wa CDM supernova+BAO fit the background parameters dominate, H0 (0.47) and Ωm (0.44) well above the dark-energy shape parameters w0 (0.24) and wa (0.16). Within the dark-energy sector w0 always ranks above wa , and this ordering has a clear physical origin: w0 is the value of the equation of state today, so it controls w(z) at low redshift, exactly where dark energy dominates the energy budget and where the supernovae and BAO are most numerous and constraining, whereas wa only sets the evolution through the z/(1 + z) factor, which vanishes at z = 0 and grows only towards higher redshift, where dark energy is subdominant to matter and the data are sparser. w0 therefore has far more leverage on the observables than wa , as the marginalised errors confirm (σ(w0 ) ≈ 0.066 versus σ(wa ) ≈ 0.42, i.e. w0 about six times better determined), and wa , the flattest and most degenerate direction, ranks last; this is the featureimportance counterpart of the elongated w0 –wa ridge of Figure 11. Adding the compressed CMB prior inverts the background-versus-dark-energy hierarchy: w0 jumps to the top (0.53) and wa to second (0.27), while Ωm (0.25) and H0 (0.24) drop. The CMB prior adds tight, direct constraints on w0 and wa , which the supernovae and BAO left almost free, so the χ2 now rises steeply along w0 and wa and they take most of the SHAP weight, whereas Ωm and H0 , already well measured and consistent with the prior, contribute proportionally less. Note that w0 stays ahead of wa throughout (the same ∼ 2 : 1 ratio with and without the prior), for the same low-redshift-leverage reason. The SHAP reordering is thus the feature-level fingerprint of the contour contraction seen in Figure 11, and the agreement between this data-driven ranking and the expected constraining structure confirms that the surrogate has learned the correct physical dependencies rather than numerical artefacts. Beyond this global ranking, SHAP also resolves how each parameter acts point by point across the sampled volume. Figure 14 shows this for w0 in the Pantheon+ & DESI BAO fit: its SHAP value, the signed contribution of w0 to the emulated χ2 at each sampled cosmology, plotted against w0 and coloured by wa . The points trace a clear parabola, the data-driven image of the χ2 well along w0 : the contribution is least, and slightly negative, in the densely sampled best-fit region near w0 ≈ −0.8, where w0 sits at its optimum and helps lower the χ2 , and grows steeply and positive as w0 departs in either direction and drives the χ2 up. It is the magnitude of this contribution, not its sign, that defines importance: the mean absolute value reported in Figure 13 is dominated by these steep flanks, so the parabola itself is the fingerprint of a tightly constrained parameter, whereas a flat, weakly constrained direction such as wa shows small SHAP values throughout and no comparable well. The small contribution in the best-fit cloud therefore does not mean that w0 is unimportant there; it 22
Figure 13: Global SHAP feature importance (mean absolute SHAP value, i.e. each parameter’s average contribution to the emulated χ2 ) for (a) the ΛCDM baseline, (b) w0 wa CDM with Pantheon+ & DESI BAO, and (c) the same fit with the CMB prior added.
Figure 14: Point-wise SHAP values for w0 in the w0 wa CDM fit to Pantheon+ & DESI BAO. Each point is a sampled cosmology: its horizontal position is the value of w0 , its vertical position the SHAP value of w0 (its signed contribution to the emulated χ2 at that point), and its colour the value of wa . simply reflects that, at the bottom of the well, the likelihood is locally flat in w0 , exactly as expected for a parameter pinned at its optimum. The colour finally encodes the residual w0 –wa interaction: at fixed w0 the contribution still varies with wa , the feature-level counterpart of the elongated w0 – wa degeneracy of Figure 11. Analogous per-parameter plots for the remaining parameters and data combinations are collected in the public PhysicsML repository (see the Data Availability Statement). 3.3.7
What the emulation brings
The cosmological application plays a double role in this work. Methodologically it is a validation of the ML strategy: unlike the flavour and ALP likelihoods, the exact cosmological χ2 is cheap enough to be sampled directly, so we could run the same MCMC with four independent back-ends (Section 3.3.5) and verify that the surrogate posteriors are statistically identical to the exact ones. This is a check that cannot be carried out where the exact likelihood is prohibitive, and it certifies that the order-ofmagnitude speedup of the emulator does not come at the price of accuracy. Beyond this cross-check, the cosmological fits make the practical advantages of the approach concrete. First, speed: once trained, the emulator samples each posterior in a few seconds against tens to 23
hundreds of seconds for the exact likelihood, and even the up-front construction of the training set is accelerated on the GPU, so the cost is paid once and amortised over the whole analysis; the advantage grows with the cost of the underlying likelihood, which is precisely why it becomes decisive in the flavour and ALP applications. Second, robustness and adaptability: a single pipeline handles the two-parameter ΛCDM fit and the four-parameter w0 wa CDM one, three different supernova samples, the presence or absence of the CMB prior, and likelihoods that are tightly (χ2min ≈ 1700) or barely (χ2min ≈ 5) constrained, with only the shifted-log target of Eq. (29) needed to make the surrogate resolve all of them. Third, resolution: the emulator returns smooth, closed confidence regions even for the strongly curved w0 –wa degeneracy and the nearly flat BAO-only landscape, where a direct grid would be either too coarse or too expensive. Fourth, transparency: the SHAP analysis recovers the physically expected constraining hierarchy and its reorganisation when the CMB is added, so the gain in speed does not turn the inference into a black box. Taken together, these results show that the framework of Section 2 transfers cleanly from high-energy physics to cosmology, the cosmological case providing the controlled benchmark that certifies the method while the more expensive applications reap the largest computational gains.
4
Conclusions
We have presented a general and interpretable Machine Learning framework for accelerating likelihoodbased inference in high-energy physics and cosmology. Using XGBoost surrogate models combined with SHAP interpretability and parallel Markov-Chain Monte Carlo sampling, the framework efficiently explores high-dimensional, non-Gaussian parameter spaces. This demonstrates that the approach is general rather than tied to an underlying physical system. The surrogate’s robustness across diverse likelihoods stems from three core methodological choices. First, the training design combines a global space-filling prior sample with a local, Hessian-covariancescaled best-fit cloud to automatically resolve narrow, anisotropic degeneracies. This design is augmented by an active-learning Gaussian process expected improvement in sparse or sharply structured regions. Second, training targets are dynamically adapted to the local landscape, using a direct log-likelihood for the flavour fit, a two-stage classifier-plus-regressor via a sigmoid transform for the axion-like-particle fit, and a shifted-log10 transform to resolve cosmological confidence regions across a χ2min spanning two orders of magnitude. Finally, regression trees partition the parameter space adaptively to capture narrow ridges, sharp transitions, and plateau-like regions without neural-network oversmoothing. The space-filling baseline ensures that emulator evaluations remain strictly within the validated training domain. Across the three cases studied, the framework delivered substantial computational gains while preserving statistical accuracy. Likelihood evaluations were accelerated by speed-ups of up to 105 in flavour and axion-like particle analyses, while posterior sampling in the cosmological application achieved speed-ups ranging from 4 to 22 over the CPU (vectorized NumPy) implementation and from 7 to 12 over the JAX/GPU kernel. The advantage grows as the underlying likelihood becomes more expensive. Besides, a key advantage of the framework is its interpretability. SHAP analyses consistently recovered the expected hierarchy of parameter and observable importance, providing a transparent physical interpretation of the surrogate predictions. The cosmological case doubles as a controlled validation: because the exact likelihood could be sampled with four independent back-ends, we could verify directly that the surrogate posteriors are statistically identical to the exact ones. The three applications show that the framework reproduces the relevant phenomenology at a fraction of the computational cost. It demonstrates that the almost-perfect RD –BR(B + → K + ν ν̄) correlation found in the previous Scenario II [1] fit disappears when C1 and C3 vary independently, identifies the axion-like-particle parameter region favored by the Belle II B + → K + ν ν̄ excess, and recovers the preference for dynamical dark energy in supernova+BAO data, which reverts to ΛCDM once the compressed CMB prior is included. In all cases, SHAP values provide a transparent identification of the measurements driving the results. Overall, our results demonstrate that interpretable Machine Learning surrogates provide a fast, accurate, and scalable solution for likelihood-based inference. The framework is readily transferable
24
to other global-fit problems and offers a promising foundation for future developments, including automated active learning, differentiable emulators, and applications to higher-dimensional parameter spaces.
Acknowledgments This work is partially supported by Spanish MINECO/FEDER Grants PGC2022-126078NB-C21 and PID2024-160228NB-I00, funded by MCIN/AEI/10.13039/ 501100011033 and “ERDF A way of making Europe”, and Grant E21-23R funded by Aragón Government and the European Union, and the NextGenerationEU Recovery and Resilience Program on Astrofı́sica y Fı́sica de Altas Energı́as CEFCA-CAPA-ITAINNOVA. Jorge Alda thanks the warm hospitality of University of Zaragoza/CAPA and of the Kavli Institute for the Physics and Mathematics of the Universe in Tokyo during the completion of this work. Data availability: The code, Jupyter notebooks and data used to produce the results of this paper are publicly available in the PhysicsML repository, which gathers the three applications presented here (semileptonic B meson anomalies, axion-like particles, and cosmology). Each subproject provides the full pipeline used to build the training sets, train the XGBoost surrogate, compute the SHAP analysis and sample the posterior with MCMC, together with the input experimental data required to reproduce the figures and tables.
References [1] J. Alda, J. Guasch, and S. Penaranda, “Using Machine Learning techniques in phenomenological studies on flavour physics,” JHEP 07 (2022) 115, arXiv:2109.07405 [hep-ph]. [2] J. Alda, J. Guasch, and S. Penaranda, “Anomalies in B mesons decays: a phenomenological approach,” Eur. Phys. J. Plus 137 no. 2, (2022) 217, arXiv:2012.14799 [hep-ph]. [3] J. Alda Gallo, J. Guasch, and S. Penaranda, “Exploring B-physics anomalies at colliders,” PoS EPS-HEP2021 (2022) 494, arXiv:2110.12240 [hep-ph]. [4] S. Penaranda, J. Alda, and A. Mir, “Flavour Anomalies: A comparative Analysis using a Machine Learning Algorithm,” Int. J. Theor. Phys. 65 no. 2, (2026) 46, arXiv:2412.15830 [hep-ph]. [5] A. Mir Ramos, J. Alda, and S. Penaranda, “B-Meson Anomalies: Effective Field Theory Meets Machine Learning,” PoS EPS-HEP2025 (2026) 349, arXiv:2510.17742 [hep-ph]. [6] J. Alda, “Lecture notes on Machine Learning applications for global fits,” 4, 2026. arXiv:2604.07520 [hep-ph]. [7] T. Chen and C. Guestrin, “XGBoost: A Scalable Tree Boosting System,” arXiv:1603.02754 [cs.LG]. [8] S. Lundberg and S.-I. Lee, “A Unified Approach to Interpreting Model Predictions,” arXiv:1705.07874 [cs.AI]. [9] S. M. Lundberg, G. Erion-Barner, and S.-I. Lee, “Consistent individualized feature attribution for tree ensembles,” arXiv:1802.03888 [cs.LG] (2018) . [10] L. S. Shapley, “A Value for n-Person Games,” in Contributions to the Theory of Games (AM-28), Volume II, H. W. Kuhn and A. W. Tucker, eds., pp. 307–318. Princeton University Press, 1953. [11] S. S. Wilks, “The Large-Sample Distribution of the Likelihood Ratio for Testing Composite Hypotheses,” Annals Math. Statist. 9 no. 1, (1938) 60–62.
25
[12] M. D. McKay, R. J. Beckman, and W. J. Conover, “A Comparison of Three Methods for Selecting Values of Input Variables in the Analysis of Output from a Computer Code,” Technometrics 21 no. 2, (1979) 239–245. [13] I. M. Sobol’, “On the distribution of points in a cube and the approximate evaluation of integrals,” USSR Comput. Math. Math. Phys. 7 no. 4, (1967) 86–112. [14] T. Akiba, S. Sano, T. Yanase, T. Ohta, and M. Koyama, “Optuna: A Next-generation Hyperparameter Optimization Framework,” in Proc. 25th ACM SIGKDD Int. Conf. on Knowledge Discovery and Data Mining. 2019. arXiv:1907.10902 [cs.LG]. [15] LHCb Collaboration, R. Aaij et al., “Measurement of the ratio of branching fractions B(Bc+ → J/ψτ + ντ )/B(Bc+ → J/ψµ+ νµ ),” Phys. Rev. Lett. 120 no. 12, (2018) 121801, arXiv:1711.05623 [hep-ex]. [16] LHCb Collaboration, R. Aaij et al., “Test of lepton universality in beauty-quark decays,” Nature Phys. 18 no. 3, (2022) 277–282, arXiv:2103.11769 [hep-ex]. [Addendum: Nature Phys. 19, (2023)]. [17] Heavy Flavor Averaging Group (HFLAV), “Average of R(D) and R(D∗ ) for CKM 2025.” https://hflav-eos.web.cern.ch/hflav-eos/semi/ckm25/html/RDsDsstar/RDRDs.html, 2025. [18] Belle-II Collaboration, I. Adachi et al., “Evidence for B+→K+νν¯ decays,” Phys. Rev. D 109 no. 11, (2024) 112006, arXiv:2311.14647 [hep-ex]. [19] J. Alda, “SMEFT19: global SMEFT fits of B-physics anomalies.” https://github.com/Jorge-Alda/SMEFT19, 2024. Methodology in Refs. [1, 2, 4]. [20] D. M. Straub, “flavio: a Python package for flavour and precision phenomenology in the Standard Model and beyond,” arXiv:1810.08132 [hep-ph]. [21] J. Aebischer, J. Kumar, P. Stangl, and D. M. Straub, “A Global Likelihood for Precision Constraints and Flavour Anomalies,” Eur. Phys. J. C 79 no. 6, (2019) 509, arXiv:1810.07698 [hep-ph]. [22] J. Aebischer, J. Kumar, and D. M. Straub, “Wilson: a Python package for the running and matching of Wilson coefficients above and below the electroweak scale,” Eur. Phys. J. C 78 no. 12, (2018) 1026, arXiv:1804.05033 [hep-ph]. [23] CMS Collaboration, V. Chekhovsky et al., “Measurement of the ratio of the Bc+→J/ψτ +ντ and Bc+→J/ψµ+νµ branching fractions using three-prong τ lepton decays,” Phys. Rev. D 113 no. 11, (2026) L111101, arXiv:2510.21559 [hep-ex]. [24] T. Ferber, A. Filimonova, R. Schäfer, and S. Westhoff, “Displaced or invisible? ALPs from B decays at Belle II,” JHEP 04 (2023) 131, arXiv:2201.06580 [hep-ph]. [25] S. Bruggisser, L. Grabitz, and S. Westhoff, “Global analysis of the ALP effective theory,” JHEP 01 (2024) 092, arXiv:2308.11703 [hep-ph]. [26] J. Alda, M. Fuentes Zamoro, L. Merlo, X. Ponce Dı́az, and S. Rigolin, “Comprehensive ALP Searches in Meson Decays,” arXiv:2507.19578 [hep-ph]. [27] R. D. Peccei and H. R. Quinn, “CP Conservation in the Presence of Instantons,” Phys. Rev. Lett. 38 (1977) 1440–1443. [28] S. Weinberg, “A New Light Boson?,” Phys. Rev. Lett. 40 (1978) 223–226. [29] F. Wilczek, “Problem of Strong P and T Invariance in the Presence of Instantons,” Phys. Rev. Lett. 40 (1978) 279–282. 26
[30] G. Grilli di Cortona, E. Hardy, J. Pardo Vega, and G. Villadoro, “The QCD axion, precisely,” JHEP 01 (2016) 034, arXiv:1511.02867 [hep-ph]. [31] C. Albertus et al., “WISPedia – the WISPs Encyclopedia,” arXiv:2602.09089 [hep-ph]. [32] L. F. Abbott and P. Sikivie, “A Cosmological Bound on the Invisible Axion,” Phys. Lett. B 120 (1983) 133–136. [33] M. Dine and W. Fischler, “The Not So Harmless Axion,” Phys. Lett. B 120 (1983) 137–141. [34] J. Preskill, M. B. Wise, and F. Wilczek, “Cosmology of the Invisible Axion,” Phys. Lett. B 120 (1983) 127–132. [35] R. T. Co, L. J. Hall, and K. Harigaya, “Axion Kinetic Misalignment Mechanism,” Phys. Rev. Lett. 124 no. 25, (2020) 251802, arXiv:1910.14152 [hep-ph]. [36] C.-F. Chang and Y. Cui, “New Perspectives on Axion Misalignment Mechanism,” Phys. Rev. D 102 no. 1, (2020) 015003, arXiv:1911.11885 [hep-ph]. [37] R. T. Co, L. J. Hall, K. Harigaya, K. A. Olive, and S. Verner, “Axion Kinetic Misalignment and Parametric Resonance from Inflation,” JCAP 08 (2020) 036, arXiv:2004.00629 [hep-ph]. [38] Y. Ema, K. Hamaguchi, T. Moroi, and K. Nakayama, “Flaxion: a minimal extension to solve puzzles in the standard model,” JHEP 01 (2017) 096, arXiv:1612.05492 [hep-ph]. [39] L. Calibbi, F. Goertz, D. Redigolo, R. Ziegler, and J. Zupan, “Minimal axion model from flavor,” Phys. Rev. D 95 no. 9, (2017) 095009, arXiv:1612.08040 [hep-ph]. [40] Y. Chikashige, R. N. Mohapatra, and R. D. Peccei, “Spontaneously Broken Lepton Number and Cosmological Constraints on the Neutrino Mass Spectrum,” Phys. Rev. Lett. 45 (1980) 1926. [41] Y. Chikashige, R. N. Mohapatra, and R. D. Peccei, “Are There Real Goldstone Bosons Associated with Broken Lepton Number?,” Phys. Lett. B 98 (1981) 265–268. [42] G. B. Gelmini and M. Roncadelli, “Left-Handed Neutrino Mass Scale and Spontaneously Broken Lepton Number,” Phys. Lett. B 99 (1981) 411–415. [43] I. Z. Rothstein, K. S. Babu, and D. Seckel, “Planck scale symmetry breaking and majoron physics,” Nucl. Phys. B 403 (1993) 725–748, arXiv:hep-ph/9301213. [44] A. Arvanitaki, S. Dimopoulos, S. Dubovsky, N. Kaloper, and J. March-Russell, “String Axiverse,” Phys. Rev. D 81 (2010) 123530, arXiv:0905.4720 [hep-th]. [45] N. Gendler, D. J. E. Marsh, L. McAllister, and J. Moritz, “Glimmers from the axiverse,” JCAP 09 (2024) 071, arXiv:2309.13145 [hep-th]. [46] H. Georgi, D. B. Kaplan, and L. Randall, “Manifesting the Invisible Axion at Low-energies,” Phys. Lett. B 169 (1986) 73–78. [47] J. E. Kim, “Weak Interaction Singlet and Strong CP Invariance,” Phys. Rev. Lett. 43 (1979) 103. [48] M. A. Shifman, A. I. Vainshtein, and V. I. Zakharov, “Can Confinement Ensure Natural CP Invariance of Strong Interactions?,” Nucl. Phys. B 166 (1980) 493–506. [49] L. Di Luzio, M. Giannotti, E. Nardi, and L. Visinelli, “The landscape of QCD axion models,” Phys. Rept. 870 (2020) 1–117, arXiv:2003.01100 [hep-ph]. [50] V. Plakkot and S. Hoof, “Anomaly ratio distributions of hadronic axion models with multiple heavy quarks,” Phys. Rev. D 104 no. 7, (2021) 075017, arXiv:2107.12378 [hep-ph].
27
[51] L. Di Luzio, S. Hoof, C. Marinissen, and V. Plakkot, “Catalogues of cosmologically self-consistent hadronic QCD axion models,” JCAP 04 (2025) 072, arXiv:2412.17896 [hep-ph]. [52] A. R. Zhitnitsky, “On Possible Suppression of the Axion Hadron Interactions. (In Russian),” Sov. J. Nucl. Phys. 31 (1980) 260. [53] M. Dine, W. Fischler, and M. Srednicki, “A Simple Solution to the Strong CP Problem with a Harmless Axion,” Phys. Lett. B 104 (1981) 199–202. [54] J. Diehl and E. Koutsangelas, “Dine-Fischler-Srednicki-Zhitnitsky-type axions and where to find them,” Phys. Rev. D 107 no. 9, (2023) 095020, arXiv:2302.04667 [hep-ph]. [55] P. Cox, M. J. Dolan, M. Hayat, A. Thamm, and R. R. Volkas, “Classification of three-family flavoured DFSZ axion models that have no domain wall problem,” JHEP 02 (2024) 011, arXiv:2310.16348 [hep-ph]. [56] K. Choi, S. H. Im, C. B. Park, and S. Yun, “Minimal Flavor Violation with Axion-like Particles,” JHEP 11 (2017) 070, arXiv:1708.00021 [hep-ph]. [57] J. Martin Camalich, M. Pospelov, P. N. H. Vuong, R. Ziegler, and J. Zupan, “Quark Flavor Phenomenology of the QCD Axion,” Phys. Rev. D 102 no. 1, (2020) 015023, arXiv:2002.04623 [hep-ph]. [58] M. Chala, G. Guedes, M. Ramos, and J. Santiago, “Running in the ALPs,” Eur. Phys. J. C 81 no. 2, (2021) 181, arXiv:2012.09017 [hep-ph]. [59] M. Bauer, M. Neubert, S. Renner, M. Schnubel, and A. Thamm, “The Low-Energy Effective Theory of Axions and ALPs,” JHEP 04 (2021) 063, arXiv:2012.12272 [hep-ph]. [60] J. Bonilla, I. Brivio, M. B. Gavela, and V. Sanz, “One-loop corrections to ALP couplings,” JHEP 11 (2021) 168, arXiv:2107.11392 [hep-ph]. [61] S. Das Bakshi, J. Machado-Rodrı́guez, and M. Ramos, “Running beyond ALPs: shift-breaking and CP-violating effects,” JHEP 11 (2023) 133, arXiv:2306.08036 [hep-ph]. [62] L. C. Bresciani, G. Brunello, G. Levati, P. Mastrolia, and P. Paradisi, “Renormalization of effective field theories via on-shell methods: the case of axion-like particles,” JHEP 10 (2025) 190, arXiv:2412.04160 [hep-ph]. [63] W. Altmannshofer, A. Crivellin, H. Haigh, G. Inguglia, and J. Martin Camalich, “Light new physics in B→K(*)νν¯?,” Phys. Rev. D 109 no. 7, (2024) 075008, arXiv:2311.14629 [hep-ph]. [64] K. Fridell, M. Ghosh, T. Okui, and K. Tobioka, “Decoding the B→Kνν excess at Belle II: Kinematics, operators, and masses,” Phys. Rev. D 109 no. 11, (2024) 115006, arXiv:2312.12507 [hep-ph]. [65] P. D. Bolton, S. Fajfer, J. F. Kamenik, and M. Novoa-Brunet, “Impact of new invisible particles on B→K(*)Emiss observables,” Phys. Rev. D 112 no. 3, (2025) 035010, arXiv:2503.19025 [hep-ph]. [66] M. Abumusabh, G. Dujany, D. Guadagnoli, A. Iohner, and C. Toni, “Using B+→K+νν¯ to search for the QCD axion,” Phys. Rev. Res. 8 no. 2, (2026) 023129, arXiv:2510.18953 [hep-ph]. [67] L. Gärtner, N. Krug, T. Kuhr, M. A. Schmidt, S. Stefkova, and B. Yabsley, “Constraints on invisible B + → K + X decays from the Belle II B + → K + ν ν̄ measurement,” arXiv:2602.09666 [hep-ph]. 28
[68] BaBar Collaboration, J. P. Lees et al., “Search for B → K (∗) νν and invisible quarkonium decays,” Phys. Rev. D 87 no. 11, (2013) 112005, arXiv:1303.7465 [hep-ex]. [69] L. Di Luzio, F. Mescia, E. Nardi, P. Panci, and R. Ziegler, “Astrophobic Axions,” Phys. Rev. Lett. 120 no. 26, (2018) 261803, arXiv:1712.04940 [hep-ph]. [70] F. Björkeroth, L. Di Luzio, F. Mescia, E. Nardi, P. Panci, and R. Ziegler, “Axion-electron decoupling in nucleophobic axion models,” Phys. Rev. D 101 no. 3, (2020) 035027, arXiv:1907.06575 [hep-ph]. [71] L. Di Luzio, F. Mescia, E. Nardi, and S. Okawa, “Renormalization group effects in astrophobic axion models,” Phys. Rev. D 106 no. 5, (2022) 055016, arXiv:2205.15326 [hep-ph]. [72] J. Alda, M. Fuentes Zamoro, L. Merlo, X. Ponce Dı́az, and S. Rigolin, “ALPaca: The ALP Automatic Computing Algorithm,” arXiv:2508.08354 [hep-ph]. [73] CosmoVerse Network Collaboration, E. Di Valentino et al., “The CosmoVerse White Paper: Addressing observational tensions in cosmology with systematics and fundamental physics,” Phys. Dark Univ. 49 (2025) 101965, arXiv:2504.01669 [astro-ph.CO]. [74] M. Chevallier and D. Polarski, “Accelerating universes with scaling dark matter,” Int. J. Mod. Phys. D 10 (2001) 213–224, arXiv:gr-qc/0009008. [75] E. V. Linder, “Exploring the expansion history of the universe,” Phys. Rev. Lett. 90 (2003) 091301, arXiv:astro-ph/0208512. [76] SDSS Collaboration, D. J. Eisenstein et al., “Detection of the Baryon Acoustic Peak in the Large-Scale Correlation Function of SDSS Luminous Red Galaxies,” Astrophys. J. 633 (2005) 560–574, arXiv:astro-ph/0501171. [77] D. Scolnic et al., “The Pantheon+ Analysis: The Full Data Set and Light-Curve Release,” Astrophys. J. 938 no. 2, (2022) 113, arXiv:2112.03863 [astro-ph.CO]. [78] D. Brout et al., “The Pantheon+ Analysis: Cosmological Constraints,” Astrophys. J. 938 no. 2, (2022) 110, arXiv:2202.04077 [astro-ph.CO]. [79] DES Collaboration, T. M. C. Abbott et al., “The Dark Energy Survey: Cosmology Results with ∼1500 New High-redshift Type Ia Supernovae Using the Full 5 yr Data Set,” Astron. Astrophys. 690 (2024) A12, arXiv:2401.02929 [astro-ph.CO]. [80] DES Collaboration, B. Popovic et al., “The Dark Energy Survey Supernova Program: A Reanalysis Of Cosmology Results And Evidence For Evolving Dark Energy With An Updated Type Ia Supernova Calibration,” Mon. Not. Roy. Astron. Soc. 548 (2026) stag632, arXiv:2511.07517 [astro-ph.CO]. [81] DESI Collaboration, M. Abdul Karim et al., “DESI DR2 results. II. Measurements of baryon acoustic oscillations and cosmological constraints,” Phys. Rev. D 112 no. 8, (2025) 083515, arXiv:2503.14738 [astro-ph.CO]. [82] M. Tristram et al., “Cosmological parameters derived from the final Planck data release (PR4),” Astron. Astrophys. 682 (2024) A37, arXiv:2309.10034 [astro-ph.CO]. [83] Planck Collaboration, N. Aghanim et al., “Planck 2018 results. VI. Cosmological parameters,” Astron. Astrophys. 641 (2020) A6, arXiv:1807.06209 [astro-ph.CO]. [Erratum: Astron.Astrophys. 652, C4 (2021)]. [84] J. Asorey and J. de Cruz Pérez, “Revisiting ΛCDM extensions in light of re-analyzed CMB data,” Physics of the Dark Universe 52 (2026) 102283, arXiv:2506.13537 [astro-ph.CO].
29
[85] J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang, “JAX: composable transformations of Python+NumPy programs.” http://github.com/jax-ml/jax, 2018.
30