Physics-Integrated Variational Autoencoders for Robust and Interpretable Generative Modeling
Year: 2022

Abstract:
Integrating physics models within machine learning models holds considerable promise toward learning robust models with improved interpretability and abilities to extrapolate. In this work, we focus on the integration of incomplete physics models into deep generative models. In particular, we introduce an architecture of variational autoencoders (VAEs) in which a part of the latent space is grounded by physics. A key technical challenge is to strike a balance between the incomplete physics and trainable components such as neural networks for ensuring that the physics part is used in a meaningful manner. To this end, we propose a regularized learning method that controls the effect of the trainable components and preserves the semantics of the physics-based latent variables as intended. We not only demonstrate generative performance improvements over a set of synthetic and realworld datasets, but we also show that we learn robust models that can consistently extrapolate beyond the training distribution in a meaningful manner. Moreover, we show that we can control the generative process in an interpretable manner.

heading: 1 Introduction
text: Data-driven modeling is often opposed to theory-driven modeling, yet their integration has also been recognized as an important approach called gray-box or hybrid modeling. In statistical machine learning, incorporation of mathematical models of physics (in a broad sense; including knowledge of biology, chemistry, economics, etc.) has also been attracting attention. Gray-box / hybrid modeling in machine learning holds considerable promise toward learning robust models with improved abilities to extrapolate beyond the distributions that they have been exposed to during training. Moreover, it can bring significant benefits in terms of model interpretability since parts of a model get semantically grounded to concrete domain knowledge.
A technical challenge in deep gray-box modeling is to ensure an appropriate use of physics models. A careless design of models and learning can lead to an erratic behavior of the components meant to represent physics (e.g., with erroneous estimation of physics parameters), and eventually, the overall model just learns to ignore them. This is particularly the case when we bring together simplified or imperfect physics models with highly expressive data-driven machine learning models such as deep neural networks. Such cases call for principled methods for striking an appropriate balance between physics and data-driven models to prevent the detrimental effects during learning.
Integration of physics models into machine learning has been considered in various contexts (see, e.g., [44, 40] and our Section 4), but most existing studies focus on prediction or forecasting tasks and are not directly applicable to other tasks. More importantly, the careful orchestration of physics-based and data-driven components have not necessarily been considered. A notable exception is Yin et al. [47], in which they proposed a method to regularize the action of trainable components of a hybrid
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
model of differential equations. Their method has been developed for dynamics forecasting with additive combinations of physics and trainable models, but application to other situations is not trivial.
In this work, we aim at the integration of incomplete physics models into deep generative models. While we focus on variational autoencoders (VAEs, [15, 29]), our idea is applicable to other models in principle. In our VAE, the decoder comprises physics-based models and trainable neural networks, and some of the latent variables are semantically grounded to the parameters of the physics models. Such a VAE, if appropriately trained, is by construction partly interpretable. Moreover, since it can by construction capture the underlying physics, it will be robust in out-of-distribution regime and exhibit meaningful extrapolation properties. We propose a regularized learning framework for ensuring the meaningful use of the physics models and the preservation of the semantics of the latent variables in the physics-integrated VAEs. We empirically demonstrate that our method can learn a model that exhibits better generalization, and more importantly, can extrapolate robustly in out-of-distribution regime. In addition, we show how the direct access to the physics-grounded latent variables allows us to alter properties of generation meaningfully and explore counterfactual scenarios.

heading: 2 Physics-integrated VAEs
text: We first describe the structure of VAEs we consider, which comprise physics models and machine learning models such as neural nets. We suppose that the physics models can be solved analytically or numerically with a reasonable cost, and the (approximate) solution is differentiable with regard to the quantities on which the solution depends. This assumption holds in most physics models known in practice, which come in different forms such as algebraic and differential equations. If there is no closed-form solution of algebraic equations, we can utilize differentiable optimizers [3] as a layer of the model. For differential equations, differentiable integrators [see, e.g., 7] will constitute a layer. Handling non-differentiable and/or overly-complex simulators remains an important open challenge.

heading: 2.1 Example
text: We start with an example to demonstrate the main concepts. Let us suppose that data comprise time-series of the angle of pendulums following an ordinary differential equation (ODE):
d2ϑ(t)/dt2 + ω2 sinϑ(t)︸ ︷︷ ︸ given as prior knowledge, fP + ξdϑ(t)/dt− u(t)︸ ︷︷ ︸ to be learned by NN, fA = 0, (1)
where ϑ is a pendulum’s angle, and ω, ξ, and u are the pendulum’s angular velocity, damping coefficient, and external force, respectively. We suppose that a data point x is a sequence of ϑ(t), i.e., x = [ϑ(0) ϑ(∆t) · · · ϑ((τ − 1)∆t)]T ∈ Rτ for some ∆t ∈ R and τ ∈ N, where ϑ(t) denotes the solution of (1) with a particular configuration of ω, ξ, and u. In this example, we learn a VAE on a dataset comprising such x’s with different configurations of ω, ξ, and u.
Suppose that the first two terms of (1) are given as prior knowledge, i.e., we know that the governing equation should contain fP(ϑ, zP) := ϑ̈ + z2P sinϑ. We will use such prior knowledge, fP, by incorporating it in the decoder of a VAE that we will learn. Since fP misses some effects of the true system (1), we complete it by augmenting the decoder with a neural network fA(ϑ, zA). The VAE’s latent variable will have two parts, zP and zA, respectively linked to fP and fA. On one hand, zA works as an ordinary VAE’s latent variable since fA is a neural net, and we suppose zA ∈ Rd, p(zA) := N (0, I). On the other hand, we semantically ground zP to a physics parameter; in this case, zP ∈ R should work as pendulum’s ω. In summary, the augmented decoder here is E[x] = ODEsolveϑ [ fP(ϑ(t), zP) + fA(ϑ(t), zA) = 0 ] , where ODEsolveϑ denotes some differentiable solver of an ODE with regard to ϑ. The encoder will have corresponding recognition networks for zP and zA. The situation in this example will be numerically examined in Section 5.1.

heading: 2.2 General formulation
text: We now present the concept of our physics-integrated VAEs in a general form. Note that our interest is not limited to the additive model combination nor ODEs. In fact, the general formulation below subsumes non-additive augmentation of various physics models. The notation introduced in this section will be used to explain the proposed regularized learning method later in Section 3.
For ease of discussion, we suppose that a VAE decoder comprises two parts: a physics-based model fP and a trainable auxiliary function fA. More general cases, for example with multiple trainable functions fA,1, fA,2, . . . used in different ways, are handled in Appendix A.

heading: 2.2.1 Latent variables and priors
text: We consider two types of latent variables, zP ∈ ZP and zA ∈ ZA, which respectively will be used in fP and fA. The latent variables can be in any space, but for the sake of discussion, we suppose ZP and ZA are (subsets of) the Euclidean space and set their prior distribution as multivariate normal:
p(zP) := N (zP |mP, v2PI) and p(zA) := N (zA | 0, I), (2) where mP and v2P are defined in accordance with prior knowledge of fP’s parameters. Note that zP will be directly interpretable as they will be semantically grounded to the parameters of the physics model fP; for example in Section 2.1, zP := ω was the angular velocity of a pendulum.

heading: 2.2.2 Decoder
text: The decoder of a physics-integrated VAE comprises two types of functions1, fP : ZP → YP and fA : YP ×ZA → YA. For notational convenience, we consider a functional F that evaluates fP and fA, solves an equation if any, and finally gives observation x ∈ X . X may be the space of sequences, images, and so on. Assuming Gaussian observation noise, we write the observation model as
pθ(x | zP, zA) := N ( x | F [fA, fP; zP, zA],Σx ) , (3)
where zA ∈ ZA and zP ∈ ZP are the arguments of fA and fP, respectively. Note that fA and fP may have other arguments besides zA and zP, respectively, but they are omitted for simplicity. We denote the set of trainable parameters of fA and fP (and Σx) by θ, while fP may have no trainable global parameters other than zP.
Let us see the semantics of the functional2 F first in the light of the example of Section 2.1. Recall that there we considered the additive augmentation of ODE (as in [47] and other studies). It is subsumed by the expression (3) by setting F [fA, fP; zP, zA] := ODEsolve[fP(zP) + fA(zA) = 0]. Let us generalize the idea. Our definition of the decoder in (3) allows not only additive augmentation of ODE but also broader range of architectures. The composition of fP and fA is not limited to be additive because we consider general composition of functions fA and fP. Moreover, the form of the physics model is not limited to ODEs. We list some examples of the configuration:
• If equation fP = 0 has a closed-form solution SfP ∈ YP (assuming that the solution space coincides with YP, just for ease of discussion), then F is simply an evaluation of fA, for example, F [fP, fA; zA] := fA(SfP , zA). • If an algebraic equation fP = 0 or fA ◦ fP = 0 has no closed-form solution, then F will have a differentiable optimizer, e.g., F [fP, fA] := fA(arg min ‖fP‖2) or F := arg min ‖fA ◦ fP‖2. • fP = 0 or fA ◦ fP = 0 can be a stochastic differential equation (and F contains its solver), for which zP and/or zA would become a sequence encoding the realization of the process noise.
The role of fA can also be diverse; it can work not only as a complement of physics models inside equations, but also as correction of numerical errors of solvers or optimizers, downsampling or upsampling, and observables (e.g., from angle sequence to video of a pendulum).

heading: 2.2.3 Encoder
text: The encoder of a physics-integrated VAE accordingly comprises two parts: for posterior inference of zP and for that of zA. We consider the following decomposition of the approximated posterior:
qψ(zP, zA | x) := qψ(zA | x)qψ(zP | x, zA), where qψ(zA | x) := N ( zA | gA(x),ΣA ) , qψ(zP | x, zA) := N ( zP | gP(x, zA),ΣP ) . (4)
1The distinction between fP and fA depends on the origin of the functional forms (and not if trainable or not). The form of fP depends on physics’ insight and thus fixed. On the other hand, the form of fA is determined only from utility as a function appoximator, and we can use whatever useful (e.g., feed-forward NNs, RNNs, etc.).
2It is natural to consider that F is a functional (and not a function) because we may need the access to the functions fA and fP themselves, rather than their pointwise values. For example, we need the full access to those functions when the decoder has an ODE solver with arbitrary initial condition.
gA : X → ZA and gP : X ×ZA → ZP are recognition networks. We denote the trainable parameters of gA and gP (and ΣA and ΣP) as ψ. This particular dependency is for our regularization method in Section 3.2, where gP should first remove the information of zA from x and then infer zP.

heading: 2.3 Evidence lower bound
text: The VAE is to be learned as usual by maximizing the lower bound of the marginal log likelihood known as evidence lower bound (ELBO). In our case, it is straightforward to derive:
ELBO(θ, ψ;x) = Eqψ(zP,zA|x) log pθ(x | zP, zA) −DKL [ qψ(zA | x) ‖ p(zA) ] − Eqψ(zA|x)DKL [ qψ(zP | x, zA) ‖ p(zP) ] . (5)

heading: 3 Striking balance between physics and trainable models
text: We propose a regularized learning objective for physics-integrated VAEs. It comprises two types of regularizers. The first is for regularizing unnecessary flexibility of function approximators like neural networks and presented in Section 3.1. The second is for grounding encoder’s output to physics parameters and presented in Section 3.2. The overall objective is summarized in Section 3.3.

heading: 3.1 Regularizing excess flexibility of trainable functions
text: If the trainable component of the physics-integrated VAE (i.e., fA) has rich expression capability, as is often the case with deep neural networks, merely maximizing the ELBO in (5) provides no guarantee that the physics-based component (i.e., fP) will be used in a meaningful manner; e.g., fP may just be ignored. We want to ensure that fA does not unnecessarily dominate the behavior of the entire model and that fP is not ignored. To this end, we borrow an idea from the posterior predictive check (PPC), a procedure to check the validity of a statistical model [see, e.g., 9]. Whereas the standard PPCs examine the discrepancy between distributions of a model and data, we compute the discrepancy between those of the model and its “physics-only” reduced version, for monitoring and balancing the contributions of parts of the model.
For the sake of argument, suppose that a given physics model fP is completely correct for given data. Then, the discrepancy between the original model and its “physics-only” reduced model (where fA is somehow invalidated) should be close to zero because the decoder of both the original model (with fP and fA working) and the reduced model (with only fP working) should coincide in an ideal limit with the true data-generating process. Even if fP captures only a part of the truth, the discrepancy should be kept small, if not zero, to ensure meaningful use of the physics models in the overall model.
The “physics-only” reduced model is created as follows. Recall that the original VAE is defined by Eqs. (3) and (4). We define the decoder of the reduced model by replacing fA : YP ×ZA → YA of (3) with a baseline function hA : YP → YA. That is, the reduced observation model is
prθr(x | zP, zA) := N ( x | F [hA, fP; zP],Σx ) , (3r)
where we omit zA from the argument of F because hA no longer takes it. We denote the set of the trainable parameters of such a model as θr := θ\ param(fA) ∪ param(hA). The corresponding encoder is defined as follows. Recall that in the original model, posterior distributions of both zP and zA are inferred in (4) and then used for reconstructing each input x in (3). On the other hand, in the “physics-only” reduced model, zA is not referred to by (3r), which makes it less meaningful to place a particular posterior of zA for each x. Hence, we define the “physics-only” encoder by marginalizing out zA and using prior3 p(zA) instead. That is, the reduced posterior is
qrψ(zA, zP | x) := p(zA) ∫ qψ(zP, zA | x)dzA. (4r)
Below we give a guideline for the choice of the baseline function, hA:
• If the ranges of fP and fA are the same (i.e., YP = YA), then hA can be an identity function hA = Id. Note that in the additive case fA ◦ fP = fP + fA′ , where fA′ is a trainable function, replacing fA with hA = Id is equivalent to replacing fA′ with hA′ = 0.
3It is just for defining qrψ on the common support with qψ . Any non-informative distributions of zA are fine.
• If YP 6= YA, then hA can be a linear or affine map from YP to YA. For example, if YP = RdP and YA = RdA (dP 6= dA), then we can set hA(fP(zP)) = W fP(zP) where W ∈ RdA×dP .
The idea is to minimize the discrepancy between the full model and the “physics-only” reduced model. In particular, we minimize the discrepancy between the posterior predictive distributions
DKL [ pθ,ψ(x̃ | X) ‖ prθr,ψ(x̃ | X) ] , where
pθ,ψ(x̃ | X) = ∫ pθ(x̃ | zP, zA)qψ(zP, zA | x)pd(x | X)dzPdzAdx,
prθr,ψ(x̃ | X) = ∫ prθr(x̃ | zP, zA)qrψ(zP, zA | x)pd(x | X)dzPdzAdx.
(6)
pd(x | X) is the empirical distribution with the support on data X := {x1, . . . ,xn}. We use x̃, instead of x, just for avoiding notational confusion by clarifying the target of integral ∫ dx.
Unfortunately, analytically computing (6) is usually intractable. Hence, we take the following upper bound of (6) (a proof is in Appendix B, and further remarks are in Appendix C): Proposition 1. Let pθ and prθ be the shorthand of pθ(x̃ | zP, zA) in (3) and prθr(x̃ | zP, zA) in (3r), respectively. Let pP and pA be some distributions of zP and zA, e.g., p(zP) and p(zA) using the priors in (2), respectively. The KL divergence in (6) can be upper bounded as follows:
DKL [ pθ,ψ(x̃ | X) ‖ prθr,ψ(x̃ | X) ] ≤ Epd(x|X) [ Eqψ(zP,zA|x)DKL[pθ ‖ p r θ]
+DKL[qψ(zA | x) ‖ pA] + Eqψ(zA|x)DKL[qψ(zP | zA,x) ‖ pP] ] . (7)
Definition 1. Let us denote the upper bound (7) by Epd(x|X)D̂(θ,param(h), ψ;x). The regularization for inhibiting unnecessary flexibility of trainable functions is defined as minimization of
RPPC(θ,param(h), ψ) := Epd(x|X)D̂(θ,param(h), ψ;x). (8) Remark 1. When multiple trainable functions are differently used in a model (e.g., inside and outside an equation solver), which is often the case in practice, the definition of RPPC should be generalized to consider marginal contribution of every trainable function. See Appendix A.

heading: 3.2 Grounding physics encoder by physics-based data augmentation
text: Toward properly learning physics-integrated VAEs, minimizing RPPC solely may not be enough because inferred zP may be still meaningless but makes RPPC not that large (e.g., with solution of fP fluctuating around the mean pattern of data), and then optimization may not be able to escape such local minima. Though it is difficult to avoid such a local solution perfectly, we can alleviate the situation by considering additional objectives to encourage a proper use of the physics.
The idea is to use the physics model as a source of information for data augmentation, which helps us to ground the output of the recognition network, gP in (4), to the parameters of fP. We want to draw some zP, feed it to the physics model fP (and a solver if any), and use the generated signal as additional data during training. A technical challenge to this end is that because the physics model may be incomplete, the artificial signals from it and the real signals may have different natures. To compensate such difference, we arrange a particular functionality of the physics encoder, gP.
Let z?P be a sample drawn from some distribution of zP (e.g., prior p(zP)). We artificially generate signals xr(z?P) by feeding z ? P to the “physics-only” decoding process in (3r), that is,
xr(z?P) := F [hA, fP; zP = z?P]. (9)
We want the physics-part recognition network, gP, to successfully estimate z?P given the corresponding x
r(z?P), which is necessary to say that the result of the inference by gP is grounded to the parameters of fP. However, in general, real data x and the augmented data xr(z?P) have different natures because fP may miss some aspects of the true data-generating process.
We handle this issue by considering a specific design of the physics-part recognition network, gP. We decompose gP into two stages as gP(x, zA) = gP,2(gP,1(x, zA)) without loss of generality. On one hand, gP,1 should transform real data x to signals that resemble the physics-based augmented signal, xr. In other words, gP,1 should “cleanse” real data into a virtual “physics-only” counterpart. We enforce such a functionality of gP,1 by making its output close to the following quantity:
xr(gP(x, zA)) = F [hA, fP; zP = gP(x, zA)]. (10) On the other hand, gP,2 should receive such “cleansed” input and return the (sufficient statistics of) posterior of zP. If the aforementioned functionality of gP,1 is successfully realized, we can directly self-supervise gP,2 with xr(z?P) because x r(gP(x, zA)) and xr(z?P) should have similar nature.
In summary, we define a couple of regularizers for setting such functionality of gP,1 and gP,2 as follows (with the corresponding diagrams of computation shown in Figure 1): Definition 2. Let sg[·] be the stop-gradient operator. The regularization for the physics-based data augmentation is defined as minimization of
RDA,1(ψ) := Epd(x|X)q(zA|x) ∥∥gP,1(x, zA)− sg [xr(gP(x, zA))]∥∥22 and (11)
RDA,2(ψ) := Ez?P ∥∥gP,2( sg [xr(z?P)])− z?P∥∥22. (12)

heading: 3.3 Overall regularized learning objective
text: The overall regularized learning problem of the proposed physics-integrated VAEs is as follows:
minimize θ,param(h),ψ − Epd(x|X)ELBO(θ, ψ;x) + αRPPC(θ,param(h), ψ) + βRDA,1(ψ) + γRDA,2(ψ),
where each term appears in (5), (8), (11), and (12), respectively. Recall that θ and ψ are the sets of the parameters of the full model’s decoder (3) and encoder (4), respectively, and that param(h) denotes the set of the parameters of h, which may be empty. If we cannot specify a reasonable sampling distribution of z?P needed in (12), we do not use RDA,1 and RDA,2; it may happen when the semantics of zP are not inherently grounded, e.g., when fP is a neural Hamilton’s equation [39].

heading: 4 Related work
text: The integration of theory-driven and data-driven methodologies has been sought in various ways. We overview some perspectives in this section and more in Appendix D.
Physics+ML in model design Integration in model design, often called gray-box or hybrid modeling, has been studied for decades [e.g., 24, 30, 38] and is still active, with deep neural networks utilized in various areas [e.g., 48, 27, 21, 41, 23, 1, 2, 8, 49, 42, 33, 16, 22, 5, 34, 26, 19, 25, 35]. Most recent studies focus on prediction, and the generative modeling has been less investigated. Moreover, mechanisms to regularize the flexibility of trainable components have hardly been addressed.
The work of Yin et al. [47] is notable here because they consider a mechanism to regularize the flexibility a trainable component to preserve the utility of physics in the model, even though it is only focused on dynamics learning for forecasting. They learn an additive hybrid ODE model ẋ = fP(x) + fA(x), where fP is a prescribed physics model, and fA is a neural network. Such a model is subsumed in our architecture as exemplified in Section 2. Moreover, Yin et al. [47] propose to regularize fA by minimizing ‖fA‖2. Such a term also appears in one of our regularizers, RPPC; when the observation noise is Gaussian, the first term of the right-hand side of (7) becomes E‖(fA ◦ fP)− fP‖22 = E‖fP + fA′ − fP‖22 = E‖fA′‖22. Therefore, we get a “VAE variant” of Yin et al. [47] by switching off a part of RPPC and the other regularizers, RDA,1 and RDA,2. We examine cases similar to it in our experiment for comparison.
Yıldız et al. [46] and Linial et al. [20] developed VAEs whose latent variable follows ODEs. Linial et al. [20] also suggest grounding the semantics of the latent variable by providing sparse supervision on it. It is feasible only when we have a chance to observe the latent variable (e.g., with an increased cost) and may often be inherently infeasible in some problem settings including ours. In our method, we never assume availability of observation of latent variables and instead use the physics models in a self-supervised manner. While direct comparison is not meaningful due to the difference of settings, we examine a baseline close to the base model of Linial et al. [20] in our experiment for comparison.
0 2 4 6 8
−2 −1
0
1 extrapolation→← reconstruction
time t
ϑ (t )
Truth Phys-only NN+solver NN+phys+reg
Figure 2: Reconstruction and extrapolation of a test sample of the pendulum data. Range 0 ≤ t < 2.5 is reconstruction, whereas t ≥ 2.5 is extrapolation.
0 1 2 −2 −1
0
1
2 ω=0.86
ϑ (t )
Truth
0 1 2
ω=1.29
NN+phys
0 1 2
ω=1.72
NN+phys+reg
0 1 2
ω=E[zP] =2.15
original 0 1 2
ω=3.22
0 1 2
ω=4.29
0 1 2
ω=5.36
Figure 3: Counterfactual generation for the pendulum data. Horizontal axis is time t. The center panel shows the original data, and the rest is the generation with zP (i.e., ω) altered while zA fixed.
Toth et al. [39] propose a model where the latent variable sequence is governed by the Hamiltonian mechanics with a neural Hamiltonian. While it does not suppose very specific physics models but considers general mechanics, they can also be included in our framework; that is, fP can be a Hamilton’s equation with a neural Hamiltonian. We try such a model in one of our experiments.
Physics+ML in objective design Another prevailing strategy is to define objective functions based on physics knowledge [e.g., 36, 14, 28, 12, 45, 13, 50, 31, 6, 43]. In generative modeling, for example, Stinis et al. [37] use residuals from physics models as a feature of GAN’s discriminator. Golany et al. [10] regularize the generation from GANs by forcing it close to a prescribed physics relation. These approaches are often easy to deploy, but an inherent limitation is that given physics knowledge should be complete to some extent, otherwise a physics-based loss is not well-defined.

heading: 5 Experiments
text: We performed experiments on two synthetic datasets and two real-world datasets, for which we prepared instances of physics-integrated VAEs. We show each particular architecture of physicsintegrated VAEs and the corresponding results; some details are deferred to Appendix E. While direct comparison is impossible due to the differences of the problem settings, the baseline methods we examined (listed below) are similar to some existing methods [4, 46, 39, 20, 47].
NN-only Ordinary VAE [15, 29]; the decoder is Ex = fA(zA), where fA is a neural net. Phys-only Physics VAE; the decoder is Ex = F [fP; zP] with no neural nets. The encoder
is with neural nets as ordinary VAEs. This is almost equivalent to the method of Aragon-Calvo and Carvajal [4] when the problem is as in Section 5.3.
NN+solver VAE with physics solvers; the decoder is Ex = F [fA; zA], where fA is a neural net, and F includes some equation-solving process (e.g., ODE/PDE solver), but no more physics-based knowledge is given (i.e., there is no fP). This is similar to the methods of, for example, Yıldız et al. [46] and Toth et al. [39].
NN+phys Physics-integrated VAE learned without the regularizers (i.e., α = β = γ = 0); this is similar to the base models of Linial et al. [20] and Qian et al. [25]. Finer ablations are also studied, among which the cases with β = 0 or γ = 0 are similar to the model of Yin et al. [47].
NN+phys+reg Our proposal; physics-integrated VAE learned with the proposed regularizers.
We aligned the total dimensionality of the latent variables of each method (except phys-only); when dim zA = dA and dim zP = dP in NN+phys(+reg), we set dim zA = dA + dP in NN-only and NN+solver. The hyperparameters, α, β, and γ, were chosen with validation set performance. We investigated the performance sensitivity to them; no large degradation of performance was observed even if we changed the values by ×10 or × 110 from the chosen values; details are in Appendix F.

heading: 5.1 Forced damped pendulum
text: Dataset We generated data from (1) with u(t) = Aω2 cos(2πφt). Each data-point x is a sequence x := [ϑ1 · · ·ϑτ ] ∈ Rτ , where ϑj is the value of a solution ϑ(tj) at tj := (j − 1)∆t. We randomly
drew a sample of the initial condition ϑ1 (with ϑ̇1 = 0 fixed) and the values of ω, ζ, A, and φ for each sequence. We generated 2,500 sequences of length τ = 50 with ∆t = 0.05 and separated them into a training, validation, and test sets with 1,000, 500, and 1,000 sequences, respectively.
Setting We set fP as in Section 2.1, i.e., fP(ϑ, zP) := ϑ̈ + z2P sin(ϑ), where zP ∈ R should work as angular velocity ω. We augmented it by fA,1(ϑ, zA,1) additively, where fA,1 was a multilayer perceptron (MLP) and zA,1 ∈ R. The ODE fP + fA,1 = 0 was solved with the Euler update scheme in the model. The model had another MLP4 fA,2 with another latent variable zA,2 ∈ R2 for further modifying the solution of the ODE. In summary, the decoding process is F := fA,2(solveϑ[fP(ϑ, zP) + fA,1(ϑ, zA,1) = 0], zA,2). The construction of the proposed regularizer for such multiple fA’s is elaborated in Appendix A. We used hA,1 = 0 and hA,2 = Id as the baseline functions. The recognition networks, gA,1, gA,2, and gP, were modeled with MLPs. We used the initial element of each x as an estimation of the initial condition ϑ1.
Results Figure 2 demonstrates a unique benefit of the hybrid modeling. We show an example of reconstruction with extrapolation. Recall that the training data comprise sequences of range 0 ≤ t < 2.5 only; so the results in t ≥ 2.5 are extrapolation (in time) rather than mere reconstruction. We can observe that while NN+solver cannot extrapolate even if it is equipped with an neural ODE, NN+phys+reg can reconstruct and extrapolate correctly.
Figure 3 illustrates well the advantage of the proposed regularizers. We show an example of generation from learned models with zP manipulated. Recall that zP is expected to work as pendulum’s angular velocity ω. We took a test sample with ω ≈ E[zP] ≈ 2.15 and generated signals with the original and different values of zP, keeping the values of zA to be the original posterior mean. We can see that the generation from NN+phys+reg matches better with the signals from the true process.
Table 1 (left half) summarizes the performance in terms of the reconstruction error and the inference error of physics parameter ω on the test set. The errors are reported in mean absolute errors (MAEs). The inference error of ω is evaluated by |E[zP]− ωtrue|. NN+phys+reg achieves small values in both reconstruction error and inference error. Meanwhile, the MAE of reconstruction by phys-only is significantly worse than those of the other methods, and the MAE of ω inferred by NN+phys is significantly worse than the others. These facts imply the effectiveness of the hybrid modeling and the proposed regularizers.

heading: 5.2 Advection-diffusion system
text: Dataset We generated data from advection-diffusion PDE ∂T/∂t− a · ∂2T/∂s2 + b · ∂T/∂s = 0, where s is the 1-D spatial dimension. We approximated the solution T (s, t) on the 12-point even grid from s = 0 to s = smax, so each data-point x is a sequence of 12-dim vectors, i.e., x := [T1 · · · Tτ ] ∈ R12×τ , where Tj := [T (0, tj) · · · T (smax, tj)]T at tj := (j − 1)∆t. We set the boundary condition as T (0, t) = T (smax, t) = 0 and the initial condition as T (s, 0) = c sin(πs/smax). We randomly drew a, b, and c for each x. We generated 2,500 sequences with τ = 50 and ∆t = 0.02 and separated them into a training, validation, and test sets with 1,000, 500, and 1,000 sequences, respectively.
4We used MLP as the data are fixed length. The same holds hereafter. Extension to other networks is easy.
Setting We set fP as the diffusion PDE, i.e., fP(T, zP) := ∂T/∂t− zP∂2T/∂s2, where zP ∈ R should work as diffusion coefficient a. We augmented it by fA(T, zA) additively, where fA was an MLP and zA ∈ R4. Hence, the decoding process is F := solveT [fP(T, zP) + fA(T, zA) = 0]. We used hA = 0 as the baseline function. The recognition networks, gA and gP, were modeled with MLPs. We used the initial snapshot of each sequence x as an estimation of the initial condition T1.
Results Figure 4 shows an example of reconstruction with extrapolation. As the training data only comprise sequences of range 0 ≤ t < 1, the remaining range t ≥ 1 is extrapolation. Only NN+phys+reg (the bottom panel) achieves adequate extrapolation; phys-only lacks advection, NN+solver has unnatural artifacts, and NN+phys infers zP (i.e., diffusion coefficient a) wrongly.
Table 1 (right half) summarizes the reconstruction and inference errors, which are basically consistent with the results in the pendulum example, in the sense that NN+phys+reg achieves reasonable performance both in reconstruction and inference, while phys-only fails reconstruction, and NN+phys fails inference. Note that the reconstruction performance of NN+phys+reg is slightly worse than some baselines, which is probably due to suboptimal hyperparameters. In fact, with finer tuning of the hyperparameters, NN+phys+reg can achieve the reconstruction error closer to other methods while almost keeping the inference error5. We also show the performance of ablations of NN+phys+reg, where either of the regularizers was turned off (i.e., α = 0, β = 0, or γ = 0). Not surprisingly their performance is worse than the full regularization, especially in terms of the inference error.

heading: 5.3 Galaxy images
text: Dataset We used images of galaxy of the Galaxy10 dataset [18]. We selected the 589 images of the “Disk, Edge-on, No Bulge” class and separated them into training, validation, and test sets with 400, 100, and 89 images, respectively. Each image is of size 69× 69 with three channels. We performed data augmentation with random rotation and increased the size of the training set by 20 times.
Setting We set fP : R4>0 → R69×69 as an exponential profile of the light distribution of galaxies [see 4, and references therein] whose input is zP := [I0 A B ϑ]T ∈ R4>0. Let [fP(zP)]i,j denote the (i, j)-element of the output of fP. Then, for 1 ≤ i, j ≤ 69, [fP(zP)]i,j := I0 exp(−ri,j), where r2i,j := (Xj cosϑ − Yi sinϑ)2/A2 + (Xj sinϑ + Yi cosϑ)2/B2, and (Xj ,Yi) is the coordinate on the 69× 69 even grid on [−1, 1]× [−1, 1]. We modify the output of fP using a U-Net-like neural network fA : R69×69 × Rdim zA → R69×69×3. Thus, the decoding process is F := fA(fP(zP), zA). We set dim zA = 2 for NN+phys+reg. We set hA : R69×69 → R69×69×3 to be the repeat operator along the channel axis. The encoding process is as follows: first, features are extracted from an image x by a convolutional net like [4]. The extracted features are flattened and fed to MLPs gP and gA.
5In the experiment with the advection-diffusion dataset reported in Table 1, the selected values of the hyperparameters were α = 0.1, β = 0.01, and γ = 106, which were chosen from only eight candidates (see Appendix E for detail). When we instead set α = 0.032, β = 0.01, and γ = 106 in the sensitivity experiment (shown in Appendix F), the reconstruction error of NN+phys+reg was 0.0390 (4.5× 10−4), which is comparable to the baselines’ performance in Table 1. In this setting, the inference error of NN+phys+reg was 0.0103 (1.5 × 10−3). We only reported the suboptimal values in Table 1 to align the granularity of the hyperparameter tuning grid with that in the experiment with the pendulum dataset.
Results Figure 5 shows an example of original data and random generation from the learned models. NN-only tends to generate non-realistic images, and NN+phys generates slightly better but still spuriously, whereas NN+phys+reg consistently generates galaxy-like images. More results (reconstruction, counterfactual generation, and inspection of latent variable) are deferred to Appendix F.

heading: 5.4 Human gait
text: Dataset We used a part of the dataset provided by [17], which contains measurements of locomotion at different speeds of 50 subjects. We extracted the angles of hip, knee, and ankle in the sagittal plane. Data originally comprise sequences of each stride normalized to be 100 steps, so each data-point x is a sequence x := [ϑ1 · · ·ϑ100] ∈ R3×100, where ϑj := [ϑhip,j ϑknee,j ϑankle,j ]T. We used different 400, 100, and 344 sequences as training, validation, and test sets, respectively.
Setting Biomechanical modeling of gait is a long-standing problem [see, e.g., 32]. We did not choose a specific model but let fP be a trainable Hamilton’s equation as in [39, 11]. zP ∈ R2dH worked as the initial conditions of it, where dH was the dimensionality of the generalized position. We let dH = 3 and modeled the neural Hamiltonian with an MLP. The solution of fP = 0 was transformed by fA that also took zA ∈ R15 as an argument. In summary, the decoding process is F = fA(solve[fP = 0], zA). We set hA to be an affine transform at each timestep, which had a weight matrix and a bias as param(h). The recognition networks were modeled with MLPs.
Results Figure 6 is for visually comparing the difference of the learned models’ behavior due to the proposed regularizers. We compare the reconstructions by NN+phys and NN+phys+reg. The dashed lines show an intermediate of the decoding process, i.e., solve[fP = 0], and the red solid lines show the final reconstruction, i.e., fA(solve[fP = 0]). Without the regularization (upper row), solve[fP = 0] returns almost meaningless signals, and fA bears the most effort of reconstruction. On the other hand, with the regularization (lower row), solve[fP = 0] already matches well the data, and fA modifies it only slightly. Superiority of the regularized model was also confirmed quantitatively; the average test reconstruction errors were 0.273 with NN+phys and 0.259 with NN+phys+reg.

heading: 6 Conclusion
text: Physics-integrated VAEs by construction attain partial interpretability as some of the latent variables are semantically grounded to the physics models, and thus we can generate signals in a controlled manner. Moreover, they have extrapolation capability due to the physics models. In this work, we proposed a regularized learning objective for ensuring a proper functionality of the integrated physics models. We empirically validated the aforementioned unique capability of physics-integrated VAEs and the importance of the proposed regularization method. In future studies, it would be interesting to investigate possibility and extension to learn a hybrid generative model with a highly complex observation process.
Acknowledgments and Disclosure of Funding
This work was supported by the Innosuisse project Industrial artificial intelligence for intelligent machines and manufacturing digitalization (39453.1 IP-ICT) and the Swiss National Science Foundation Sinergia project Modeling pathological gait resulting from motor impairments (CRSII5_177179).

Review summaries:
This paper develops a general approach to integrating physical domain knowledge into the latent space, encoder, and decoder of VAEs. The latent space is separate into a physics component and learned component, where the former is grounded to some unobserved physical quantity or system variable. The decoder is a composition of a physics-based component with known functional form (e.g. an ODE of the latent physical variable) and a neural component which also takes the learned part of the latent space as input. The encoder is a composition of two stages, the first of which extracts the learned latents, and the second of which recovers the physically grounded latent variables. In order to train this model to encourage interpretability and usage of the physics components, the authors develop several regularization terms which are added to the ELBO loss, thereby constrain the power of highly flexible neural encoders and decoders, ensuring that the physical components play a meaningful role in explaining the data. The importance of such regularization (and the value of physics components) is demonstrated on a series of physical system identification tasks.

The authors propose to incorporate physics models into variational autoencoder (VAE) architectures in a principled manner. They design regularization terms in the ELBO loss to achieve a VAE with a latent space that is partially grounded in interpretable physics elements. The experiments demonstrate that the resulting VAE performs well in reconstruction error, is able to infer physically interpretable parameters in its latent space, and generalizes beyond the provided data.

This paper describes a VAE architecture in which some latent variables are semantically related to our prior knowledge of the physical model. For instance, this represents the (unknown) angular velocity or diffusion coefficient. The semantic is defined the functional form of  the decoder part. Then the objective function is augmented with regularization terms to better control how the model will take into account the physical model and the prior knowledge. The experimental part relies on two synthetic datasets and two reals ones.