Physics-Integrated Variational Autoencoders for Robust and Interpretable Generative Modeling

1 Introduction . 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 . 2 Physics-integrated VAEs . 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 . 2.1 Example . 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 . 2.2 General formulation . 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 . 2.2.1 Latent variables and priors . 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 . 2.2.2 Decoder . 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 ) .

Summary: 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.