T OWARD BAYESIAN D EEP G REY- BOX M ODELING

Naoya Takeishi
RCAST, The University of Tokyo
RIKEN Center for Advanced Intelligence Projectntake@g.ecc.u-tokyo.ac.jp

A BSTRACT
Combining scientific models and deep neural networks (deep grey-box or hybrid modeling) is expected
to be a promising strategy for building robust, partly interpretable, and data-adaptive models. This
paper presents a preliminary study to develop a framework for learning deep grey-box models withuncertainty quantification via Bayesian inference of both scientific models and neural nets.
Keywords deep grey-box models · deep hybrid models · Bayesian neural networks

1

Introduction

Despite the success of deep neural networks in a wide range of tasks, their applicability is still limited when extrapolative
prediction or interpretation of predictions is required. In contrast, scientific mathematical models are (believed to be)capable of extrapolation and having implications for the modeled phenomena. Meanwhile, such scientific models
are sometimes incomplete; they are based on the abstraction of complex real phenomena, which hinders accuratequantitative prediction. In contrast, machine learning models including neural nets are meant to be flexible and cannaturally adapt to data. The motivation of deep grey-box modeling (or deep hybrid modeling) [e.g., 11, 6, 7, 10, 2, 8] is
to combine deep neural nets with scientific models to complement the weaknesses of the two.
A technical challenge in deep grey-box modeling is to strike a balance between machine learning and scientific models.
The excess flexibility of deep neural nets may work unfavorably in grey-box modeling because they can overwrite thescientific models. For example, consider the additive combination of two models, a neural net fNN and an incomplete
scientific model fsci , that is, y ≈ fNN (x; θNN ) + fsci (x; θsci ) in a task to predict y from x, where θNN and θsci are theunknown parameters of the neural net and the scientific model, respectively. When fNN can fit to any function, which isthe case with deep neural nets, empirical risk minimization may result in fNN fitting y − fsci (x; θsci ) to similarly good
extents for any value of θsci , meaning choosing a specific value of θsci is impossible. This issue happens as well in thegeneral composition of the two models in the form of fNN ◦ fsci .1 We thus need to somehow limit the flexibility of fNN
by regularization; in fact, it has already been addressed [11, 7], at least partly, and is not the main subject of this paper.
We will work on the uncertainty quantification of deep grey-box models, which is also an important step for gaining
insights of phenomena from inferred models. While most existing studies on deep grey-box modeling, such as [11, 7],deal with point estimation of the model parameters, we will develop a framework for Bayesian inference of theparameters of both parts of a deep grey-box model, that is, a deep neural net and a scientific model. One of therecent studies most relevant to our motivation is by Akhare et al. [1], where they perform the posterior inference ofgrey-box models using the stochastic weight averaging [5] and deep ensembles [3]. However, they do not consider the
regularization of neural nets to limit their flexibility. Since such a regularizer is usually defined on the function value ofa neural network, instead of the network’s parameters, the problem becomes a kind of generalized Bayesian inference,and thus the inference method should be (often slightly) modified accordingly. In this paper, we report a preliminary
study to develop a framework for posterior inference of deep grey-box models that can handle regularization terms to
strike a balance between neural nets and scientific models.
1

The composition in the opposite direction, fsci ◦fNN , is also interesting but out of the scope here because its overall expressivenessis bounded by that of fsci , and thus the balancing the models becomes much less problematic.
International Conference on Scientific Computing and Machine Learning 2024 (SCML2024)

Toward Bayesian Deep Grey-box Modeling2

Deep grey-box models

Suppose a task to predict y from x for the sake of argument, though the same discussion holds for other kinds of tasks.
We are interested in models in the following form of the composition of two functions, fNN and fsci :y ≈ f (x; θNN , θsci ) = C fNN , fsci ; x ,(1)where C is a functional that composes the two functions (via, e.g., addition, multiplication, composition, ODE solvers,
optimization, etc., and their further compositions). We denote the unknown parameters of the two functions, fNN andfsci , by θNN and θsci , respectively. We further suppose that fNN is considerably more flexible than fsci , e.g., fNN is a
deep neural network while fsci is a handcrafted (yet scientifically meaningful) mathematical model. Such a compositionof functions makes sense when the scientific model, fsci , is somewhat incomplete, e.g., when it suffices as a qualitative
approximation of a phenomenon but lacks certain aspects of reality due to the abstraction in model-making. See [8] formore detailed definition of such a model.
Care must be taken in learning deep grey-box models because a deep neural net fNN can, it only, fit the relation from x
to y, and the value of fsci may be ignored. Thus we need to strike a balance of the expressivity between the two models,
without which grey-box modeling loses its advantages. One may achieve a good balance by meticulously tuning thecomplexity of fNN by, for example, adjusting the number of layers and units of the network. However, manually tuning
the complexity of neural nets by architectural design is hardly feasible. Instead of designing network architecture, aworking method to inhibit the excess flexibility of fNN is via constraint or regularization [11, 7, 8]. For example, themethod by Yin et al. [11] minimizes the norm of fNN for the additive combination of the two models, and Takeishi and
Kalousis [7] propose a regularizer that minimizes the effect of fNN for more general combination of the two models.
We will denote such a regularizer by R. Also, we will denote the main loss function (e.g., prediction error) by L.
Importantly, R is usually computed with the values of fNN instead of those of the full function, f , because the purposeof R is to minimize the effect of the neural net. For example, it is common to minimize the functional norm of theneural net [11], that is, R = ∥fNN ∥2 . Consequently, R usually does not need the information of y, and thus it is possible
to compute L and R using different minibatches of x, which is particularly helpful when the amount of data is limited.

3

Particle-based variational inference for Bayesian deep grey-box modeling

Despite the variety of methods for Bayesian inference of neural network parameters, in this work we opt to use theparticle-based variational inference (VI) based on the functional space similarity measure [9] because of its controllable
flexibility and principled nature. Note that we here do not intend to claim any definitive superiority of the method; ourinterest does not lie in comparing different methods for Bayesian neural net inference.
3.1

Function space particle-based VI

Suppose a prediction model g(x; θ) : X → R with input x ∈ X and parameters θ ∈ Rp . Particle-based VI approximates
the posterior distribution of θ with a set of particles, {θ (1) , . . . , θ (n) }. The method by Wang et al. [9] performs theparticle-based VI by measuring the similarity between particles in the function space, instead of the parameter space.
Each particle is updated with the following rule:
(i)

(i)
(i)
θℓ+1 ← θℓ −∂gℓ!⊤
 (i) v gℓ ,

(i)
∂θℓ

(i)

(2)

(i)

(i)

(i)

where θℓ denotes the value of the i-th particle at the ℓ-th iteration, and gℓ := [g(x1 ; θℓ ) · · · g(xm ; θℓ )]⊤ ∈ Rm(i)
is the stack of the model’s evaluations (with parameter θℓ ) on a minibatch of inputs, X := {x1 , . . . , xm }.
v : Rm → Rm defines a gradient flow. The gradient flow based on the Stein variational gradient descent [4] isn (i) 
1X
(i)
(j) 
v gℓ =
k gℓ , gℓ
n j=1

(i) 

∂ log p Y | gℓ
(i)

∂gℓ

(i) 

+

∂ log p gℓ
(i)

∂gℓ

!

(i)

+

(j) 

∂k gℓ , gℓ
(i)

!⊤
,

(3)

∂gℓ

(i)
(j) 
(i)
(j)
where k gℓ , gℓ is a kernel function that measures the similarity between two values, gℓ and gℓ , and Y denotes
the set of the labels corresponding to X = {x1 , . . . , xm }.

2

Toward Bayesian Deep Grey-box Modeling

3.2

Applying function space particle-based VI to deep grey-box models

We apply the function space particle-based VI [9] to the deep grey-box model f in eq. (1). We replace the likelihood
and the prior distribution in eq. (3) by the loss function L and the regularizer R defined in section 2, respectively. Thatis, the gradient flow becomesn (i) 
1X
(i)
(j) 
v fℓ =
k fℓ , fℓ
n j=1

∂L
(i)

∂fℓ

∂R

+

!

(i)

+

(i)

(j) 

∂k fℓ , fℓ
(i)

∂fℓ

!⊤
.

(4)

∂fℓ

Recall that R is usually computed on the values of fNN and not on those of f . It is thus possible to use different inputs
to evaluate R (via fNN ’s values) and L (via f ’s values); it means that there may be two independent computation graphs,
X → f → L and X ′ → fNN → R, where X ̸= X ′ . In this case it is not straightforward to compute ∂R/∂f . We
approximate this derivative as follows. By the chain rule, we have(i)∂R ∂fℓ∂R=,(i) ∂θ
∂θ
∂f

(5)

ℓ

(i)

⊤
⊤ ⊤
where θ = [θNNθsci
] . We compute the minimum norm estimation of ∂R/∂fℓ , that is,

∂R
≈
(i)
∂θ
∂f
∂R
ℓ

(i)

∂fℓ
∂θ

!†
,

(6)

†

where · denotes the pseudoinverse of a matrix. Computing eq. (6) using automatic differentiation is straightforward(i)because both R and fℓ comes after θ in the computation graph.
3.3

Numerical example

Dataset We used simulated data of a frictionless compound pendulum controlled by some regulator2 . The original
dataset comprises sequences of the pendulum’s state, (s1 , . . . , stmax ), where st is the angle and angular velocity of thependulum at time t. We collected pairs (x = st , y = (st+∆t , . . . , st+10∆t ))) with ∆t = 0.05, so the task is to predictthe pendulum’s state up to 10-steps ahead given the current state. We added Gaussian noise with a standard deviation
0.05 as observation noise. We used 3600, 900, and 4500 such pairs as training, validation, and test sets, respectively.
Model We know the equation of motion of an uncontrolled frictionless compound pendulum but do not know theexact formulation of the regulator controlling the pendulum in the dataset. Also, we suppose that the parameters of theequation of motion (the gravity constant g and the length of the pendulum l) are unknown. The grey-box model is
f (x) = ODESolve∆t,...,10∆t ṡ = fNN (s; θNN ) + fsci (s; θsci ) | s0 = x ,(7)where ODESolve∆t,...,10∆t denotes the operation to solve the initial value problem in the argument and return thevalues of the solution evaluated at time steps t = ∆t, . . . , 10∆t. fNN is a neural network with fully connected layers
having one hidden layer of size 32. fsci is from the equation of motion of an uncontrolled frictionless compound
pendulum with unknown parameters θsci = [g l]⊤ . We used 50 particles for posterior approximation.
Loss and regularizer The main loss function L isP
the mean squared error of the prediction, that is, L = ∥f (x) − y∥22 .
The regularizer is the norm of fNN , that is, R = s fNN (s; θNN ). The summation about s was taken using all thevalues of s that appeared in the computation regarding each minibatch of data.
Results Figure 1 depicts the inferred posterior of θsci = [g l]⊤ . For both elements, the distributions are roughly
centered around the data-generating values, g = 10 and l = 1. The distribution of g happens to have multiple modes,
but we do not have enough information so far to discuss if it is something essential or not. Figure 2 is an example
of prediction by the learned models. It works as a proof of concept as it seems to capture the uncertainty due to theobservation noise, at least to some extent.
2

Available online at stable_baselines/gail/dataset/expert_pendulum.npz of repository https://github.com/
Stable-Baselines-Team/stable-baselines; retrieved on 19 January 2024.

3

Toward Bayesian Deep Grey-box Modeling

Figure 1: Posterior particles (red dots) of the elements of θsci , (left) g and (right) l. The data-generating values wereg = 10 and l = 1. The densities estimated by the kernel density estimation (blue lines) are present for visualization.

Figure 2: Example of predictions on a test sample. The blue dots are the ground truth; the gray lines show the predictions
by each particle; the red lines are the ensemble means; and the green lines indicate the two-sigma areas.

4

Conclusion

We presented a method for Bayesian inference of deep grey-box models. We particularly used the function spaceparticle-based VI [9], for which the gradient of the regularizer should be approximated to avoid the practical difficultyof automatic differentiation. This study is still a work in progress, and there remain several issues toward effective
inference for Bayesian deep grey-box modeling. For example, we may want to use different numbers of particles (ordifferent representations) for the posteriors of θsci and θNN ; as θsci is usually significantly lower dimensional than θNN ,it may be possible to assign more particles (or more flexible representation) to the former without much sacrificing thecomputational efficiency. It would be helpful in practice because we are often interested in the detailed shape of theposterior of θsci .

Acknowledgments
This work was supported by JSPS KAKENHI Grant Numbers JP20K19869 and JP21H05300 and JSPS International
Joint Research Program Grant Number JPJSJRP-20221501. The author would like to thank Alexandros Kalousis forvaluable discussion in the early phase of this work while the author was working at the University of Applied Sciences
and Arts Western Switzerland.

References[1] Deepak Akhare, Tengfei Luo, and Jian-Xun Wang. DiffHybrid-UQ: Uncertainty quantification for differentiablehybrid neural modeling, 2023. arXiv:2401.00161.
[2] Victoriya Kashtanova, Ibrahim Ayed, Andony Arrieula, Mark Potse, Patrick Gallinari, and Maxime Sermesant.
Deep learning for model correction in cardiac electrophysiological imaging. In Proceedings of the 5th International
Conference on Medical Imaging with Deep Learning, pages 665–675, 2022.
[3] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty
estimation using deep ensembles. In Advances in Neural Information Processing Systems 30, pages 6405–6416,
2017.
[4] Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose Bayesian inference algorithm.
In Advances in Neural Information Processing Systems 29, pages 2378–2386, 2016.
4

Toward Bayesian Deep Grey-box Modeling[5] Wesley J Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple
baseline for Bayesian uncertainty in deep learning. In Advances in Neural Information Processing Systems 32,
pages 13153–13164, 2019.
[6] Zhaozhi Qian, William R. Zame, Lucas M. Fleuren, Paul Elbers, and Mihaela van der Schaar. Integrating expert
ODEs into neural ODEs: Pharmacology and disease progression. In Advances in Neural Information Processing
Systems 34, pages 11364–11383, 2021.
[7] Naoya Takeishi and Alexandros Kalousis. Physics-integrated variational autoencoders for robust and interpretable
generative modeling. In Advances in Neural Information Processing Systems 34, pages 14809–14821, 2021.
[8] Naoya Takeishi and Alexandros Kalousis. Deep grey-box modeling with adaptive data-driven models toward
trustworthy estimation of theory-driven models. In Proceedings of the 26th International Conference on Artificial
Intelligence and Statistics, pages 4089–4100, 2023.
[9] Ziyu Wang, Tongzheng Ren, Jun Zhu, and Bo Zhang. Function space particle optimization for Bayesian neuralnetworks. In Proceedings of the 9th International Conference on Learning Representations, 2019.
[10] Antoine Wehenkel, Jens Behrmann, Hsiang Hsu, Guillermo Sapiro, Gilles Louppe, and Jörn-Henrik Jacobsen.
Robust hybrid learning with expert augmentation. Transactions of Machine Learning Research, 2023.
[11] Yuan Yin, Vincent Le Guen, Jérémie Dona, Emmanuel de Bézenac, Ibrahim Ayed, Nicolas Thome, and Patrick
Gallinari. Augmenting physical models with deep networks for complex dynamics forecasting. In Proceedings of
the 9th International Conference on Learning Representations, 2021.

5