Semiparametric Inference and Equation Discovery with the Bayesian Machine Scientist

Authors: kai-hendrik  cohrs
gherardo  varando
marta  sales-pardo
roger  guimerà
gustau  camps-valls
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-09 00:01:13

Abstract: Hybrid modeling, combining machine learning with physical equations, is promising in many fields of science, in particular for climate and Earth Sciences, but faces challenges like interpretability, inconsistent extrapolation, lack of speed, and robust inference. Here we show that the Bayesian machine scientist, a Bayesian approach to symbolic regression is an ideal choice for the challenges in the hybrid modeling task. We formulate the hybrid Bayesian machine scientist and showcase its potential in the example of modeling ecosystem respiration with the $Q_{10}$ model. Specifically, we show that our proposed hybrid equation discovery method (i) extracts the correct equations, (ii) extrapolates better in different scenarios than the non-hybrid and deep-learning-based baselines, and (iii) is able to infer more accurately parameters of interest, even in the presence of equifinality. We anticipate a spur of development of hybrid equation discovery algorithms in the sciences to approach fully interpretable data-driven models.

Main paper: 
 introduction: 
 Earth system models (ESMs) play a major role in understanding human impact on the Earth's climate (IPCC, 2021) , crucial for addressing climate change. They are complex mathematical representations of the Earth's climate system involving coupled nonlinear ordinary and partial differential equations (Scholze et al., 2012) . To enhance these models, the community pursues integrations with machine learning (ML), made possible by abundant Earth observation data (Grundner et al., 2022; Reichstein et al., 2019; Camps-Valls et al., 2021) . However, this approach is not without its challenges, which include violating physical laws, poor generalization to out-of-domain (OOD) data, and lack of interpretability in black-box ML models (Reichstein et al., 2019; Marcus, 2018; Quiñonero-Candela et al., 2009; Sugiyama & Kawanabe, 2012; Rudin & Radin, 2019) .Hybrid modeling, integrating ML with physical equations, has emerged as a popular approach to harness these models (Karpatne et al., 2022; Koppa et al., 2022; Zhao et al., 2019; Tramontana et al., 2020) . However, challenges persist in speed, interpretability, and addressing equifinality (Oberpriller et al., 2021; Reichstein et al., 2022) , that is the existence of multiple (parametric and non-parametric) models and sets of parameters that describe the data similarly well.Symbolic regression is gaining attention as a promising solution to address some of the problems of ML (Camps-Valls et al., 2023) . It has been shown to be effective in supporting scientific discovery in different areas of science (Abdellaoui & Mehrkanoon, 2021; Martinez-Gil & Chaves-Gonzalez, 2020; Weng et al., 2020; Wang et al., 2019; Batra et al., 2021; Lemos et al., 2022; Reichardt et al., 2020) , including the Earth sciences (Grundner et al., 2023) . Unlike black-box ML models, symbolic regression yields mathematical equations that are white-box models, i.e., interpretable to humans. This enables direct insight into its extrapolation behavior, providing clarity and transparency in predicting outcomes beyond the observed data range.Because in AI-enhanced ESMs there are several sources of uncertainty, it is crucial to sample from distributions over parameter values, ML models, and, ideally, discovered equations to capture the uncertainty of its projections. A Bayesian approach is the natural choice to obtain distributions over equations (and eventually parameters). The Bayesian machine scientist (BMS) is a Bayesian formulation of symbolic regression (Guimerà et al., 2020) and has shown impressive results for extrapolation and discovery of physical equations.In this work, we tie the ends together towards a hybrid BMS, which enables us to extract distributions over both equations and fixed meaningful parameters in a Bayesian fashion, thus pioneering the path for combining perturbed parameter ensemble (Dagon et al., 2020) with symbolic regression in climate models. Our concrete contributions are as follows: We formulate the BMS for the hybrid modeling setting and showcase its capabilities in a simple but significant problem of ecosystem respiration modeling, where the hybrid Bayesian machine scientist (HBMS): 1) improves extrapolation under different forms of OOD data; 2) recovers the true equations; and 3) shows improved inference on parameters of interest and robustness under equifinality (the lack of identifiability). We find that these results hold for both synthetic and measured data. Our findings showcase that the HBMS could be the preferred tool for offline learning in the context of ESMs. 
hybrid bayesian machine scientist: 
 Prior works have explored integrating symbolic regression and domain knowledge. Famous approaches, such as SINDy (Brunton et al., 2016) or genetic programing (GP) (Koza, 1994) , have been subject to the introduction of various scientific constraints. Asadzadeh et al. ( 2021) demonstrated improved performance by partially fixing the regression tree in hybrid symbolic regression models with GP. Kronberger et al. (2022) introduced constraints on function image and derivatives to incorporate prior knowledge in a GP approach. Cornelio et al. (2023) developed AI-Descartes, which constraints symbolic regression using logical axioms. Additionally, Chen et al. ( 2021) combined sparse regression with physics-informed neural networks (PINNs) for effective use in sparse and noisy data scenarios. To our knowledge, this work is the first to apply prior knowledge constraints to the BMS, focusing on the joint inference of physical parameters in the hybrid setting with a partially fixed parametrization.In the probabilistic formulation of model selection, the posterior p(m|D) over models m given some data D encapsulates all the information about the plausibility of models. This posterior can be written asp(m|D) = exp[-H(m; D)] Z ,where H(m; D) is the description length for jointly encoding the model m and the data D, and Z is a normalizing constant. It can be approximated as H = BIC(m; D)/2 -log p(m), where BIC is the Bayesian information criterion and p(m) is a suitable prior for models (Guimerà, 2020) . The BMS is a symbolic regression algorithm that samples closed-form mathematical models from the posterior p(m|D) using MCMC (Guimerà, 2020) .Here, we formulate the BMS for additional inference on meaningful parameters in the hybrid modeling setting, in which part of the equation is known (with some parameters θ whose values are unknown and need to be learned) and the other part m needs to be inferred through symbolic regression among all possible models M . We obtain the posterior over θ via marginalizationEQUATIONSince, in general, there is no closed-form solution to the posterior p(θ|m, D), we would ideally need to run a nested MCMC, sampling each equation m and then, in turn, sampling from the posterior p(θ|m, D) over θ given the equation. This is computationally intensive, so we instead assume that p(θ|m, D) is peaked around some value θ * and use the approximationp(θ|D) ≈ m∈M δ(θ * m -θ)p(m|D) .We estimate θ * m via maximum likelihood, which amounts to assuming a flat (enough) prior over θ. EQUATIONwhere Q 10 is the parameter describing temperature sensitivity, which is the part of the equation that is well understood, X is a set of meteorological drivers, and R b describes the unknown base respiration, i.e., the baseline emission of CO 2 of an ecosystem at the reference temperature of 15 • C. For a thorough description of the real and synthetic data, see Appendix A.1. A typical hybrid modeling approach amounts to using a neural network (NN) as an estimator for R b , treating Q 10 as a trainable parameter, and fitting everything end-to-end with gradient descent, as in Reichstein et al. (2022) . We will refer to this approach as hybrid Neural Network (HNN) approach in contrast to the purely non-parametric NN approach, and it will serve as a hybrid modeling baseline.Training We report the ensemble mean (EM) for these models. Finally, we run both the HBMS and HNN with T A as an additional driver to create issues of equifinality. We report extrapolation performance and inference on Q 10 for the hybrid models. 
results: 
 HBMS extrapolates better to unseen scenarios Initially, we assess HBMS extrapolation capabilities, focusing on two scenarios detailed in Table 1 . The first involves extrapolation from one year to the next, where BMS-based approaches outperform NN-based methods. Hybrid knowledge (WD cum -α 6 α T A 3 )( SW SM POT α7 + SW SM,diff POT ) -α 7 HBMS (α 3 SW SM,diff POT + α 7 SW SM POT + α 5 )α WDcum 5 Q 0.1(T A -15) 10 Night time synthetic BMS ((WD cum + α 6 + α 2 T A )(α 4 + SW SM POT ) + SW SM,diff POT )α 0 HBMS (α 7 + α 7 SW SM,diff POT + α 4 SW SM POT + α 2 )α WDcum 2 Q 0.1(T A -15) 10 Night time measured BMS (α 5 (T A + α 4 (SW SM,diff POT ) 3 ) + α 7 )SW SM POT HBMS -((WDcum+α SW SM,diff POT 0 )+SW SM POT ) α4 Q 0.1(T A -15) 10enhances BMS, but harms the NN approach. In the second scenario, a shift from training on nighttime data to testing on daytime data challenges pure NN extrapolation, favoring the hybrid variant. HBMS typically shows superior MDL performance, as the optimal representation beats an ensemble. On real data (refer to Table 3 ), PPM excels, particularly when the real mode is harder to identify. For visualizations and a brief discussion on the sampled functions, refer to Appendix A.3. 2 , synthetic data, the HBMS methods all picked up correct relationships of SW SM POT , SW SM,diff POT , including the exponential relationship with WD cum . The BMS instead introduced WD cum linearly and found only on the random data exponential relationship. Potentially, due to the smaller range of values, a linear relationship with temperature was discovered on nighttime data alone. The HBMS found a linear relationship of WD cum on measured data.Examining Q 10 values from the scenario with 1000 randomly sampled data points, HBMS provides a distribution closer to the actual Q 10 value of 1.5 with the ability to capture multiple modes. HBMS remains robust even when introducing an additional temperature predictor, creating an equifinality scenario. In contrast, HNN struggles with this introduction, significantly affecting Q 10 estimation. The resilience of HBMS to the additional predictor aligns with its nature, as the introduction of unnecessary dynamics is penalized by an increase in description length.  
conclusion: 
 For advancing AI-enhanced ESMs, fast, scientifically constrained, interpretable, and probabilistic machine learning models are imperative. The proposed HBMS is a significant step forward, offering a joint distribution over parameters and interpretable equations. Our results demonstrate its superiority over non-hybrid and standard hybrid neural network baselines in extrapolation and robustness under equifinality, thus emphasizing its potential. In future work, we aim to identify more use cases and enhance strategies for approximating the true posterior over structural parameters, potentially using methods like the Laplace approximation. Exploring additional avenues to introduce prior scientific constraints in equation discovery, facilitated by the Bayesian framework, is a crucial aspect to investigate further.visually seems to capture the uncertainty better when the predictive mean deviates from the true curve.  
 
 Appendix:  


Decision: This paper shows that the Bayesian machine scientist, a Bayesian approach to symbolic regression for hybrid modeling tasks. It is expected that authors will be addressing comments by the reviewer in the camera ready version

Reviews: review: title: title: Innovative method for symbolic equation discovery
main_review: main_review: The paper showcases an innovative approach for symbolic equation discovery, a topic that has gained a lot of relevance. The authors show improved results in benchmarking of their method compared to the original BMS method but lack comparisons to other well established methods for symbolic discovery. The paper is well written and has a focus in earth sciences with results shown for the models of the domain.