Fully differentiable model discovery

Model discovery aims at autonomously discovering differential equations underlying a dataset . Approaches based on Physics Informed Neural Networks ( PINNs ) have shown great promise , but a fully-differentiable model which explicitly learns the equation has remained elusive . In this paper we propose such an approach by integrating neural network-based surrogates with Sparse Bayesian Learning ( SBL ) . This combination yields a robust model discovery algorithm , which we showcase on various datasets . We then identify a connection with multitask learning , and build on it to construct a Physics Informed Normalizing Flow ( PINF ) . We present a proof-of-concept using a PINF to directly learn a density model from single particle data . Our work expands PINNs to various types of neural network architectures , and connects neural network-based surrogates to the rich field of Bayesian parameter inference . 1 INTRODUCTION . Many physical , chemical and biological systems can be modelled with ( partial ) differential equations . They capture a system ’ s interactions , scales and conserved quantities in an interpretive manner . Unfortunately , manually deriving these equations from first principles is a time-consuming process and requires expert knowledge of the underlying dynamics . The last few years have seen a rising interest in automating this process , also known as model discovery . As the model space is exponentially large , a popular approach is to select a large set of candidate terms ( features ) and perform sparse regression on these features , effectively turning model discovery into variable selection ( Brunton et al. , 2016 ; Rudy et al. , 2017 ) . A uniquely challenging aspect of discovering PDEs is that many of the candidate features contain higher-order derivatives , which are challenging to calculate accurately using numerical differentiation . This essentially limited model discovery to densely sampled datasets with low noise levels . Various works have explored the use of neural networks to generate a surrogate of the data ( Rasheed et al. , 2020 ) , and perform model discovery on this surrogate instead ( Berg & Nyström , 2019 ; Both et al. , 2019 ) . By using a neural network g to learn the data u such that u ≈ g ( x , t ) , the network denoises the data , while automatic differentiation can be used to accurately calculate the ( higher- ) order derivatives of u used in the candidate features . These approaches show significant improvements when the neural network is constrained to solutions allowed by the candidate features . This essentially yields a Physics Informed Neural Network ( PINN ) ( Raissi et al. , 2017 ) , with the important distinction that the form of the constraint , i.e . which terms of the candidate features are active and make up the underlying equation , is also learned . As it is composed of all candidate features , the constraint itself is prone to overfitting , and the discovered equation will contain more terms than required . Applying ` 1 regularisation can alleviate this problem , but raises the question of how strongly to apply it . Alternatively , inactive terms can be pruned from the constraint by applying a mask , but this is a non-differentiable operation , and training the network does not take sparsity into account . The open question then is how to optimally apply a constraint , which itself is learned and sensitive to overfitting , all the while remaining fully differentiable . In this work we introduce a fully differentiable model discovery algorithm consisting of a neuralnetwork based surrogate with a constraint based on Sparse Bayesian Learning ( SBL ) . We summarise our contributions as follows : • We show how Bayesian parameter inference methods can be used as a constraint in a PINN . Specifically , we use Sparse Bayesian Learning ( SBL ) to create a fully-differentiable , robust model discovery algorithm and showcase this on various datasets . • We identify a connection with multitask learning using uncertainty and exploit this to generalise PINNs to probabilistic surrogates . We introduce a conditional normalizing flow constrained by SBL , a so called Physics Informed Normalizing Flow ( PINF ) . • We present a proof-of-concept where PINF learns a time-dependent density from unlabelled single-particle data , allowing the constraint to discover a density model directly from single particle data . 2 BACKGROUND . Model discovery with sparse regression Model discovery aims to discover the PDE from a large set of M candidate features { u , uxx , uux , . . . } . Assuming the underlying equation can be written as a linear combination of the candidate features , model discovery can be approached as a regression problem ( Brunton et al. , 2016 ) by solving ξ̂ = min ξ ‖∂tu−Θξ‖2 + λR ( ξ ) , ( 1 ) where Θ ∈ RN×M contains all candidate features , ξ ∈ RM the unknown coefficient vector and R ( ξ ) some sparsity-promoting penalty ; the number of candidate features is typically much larger than the number of terms in the underlying equation . The main challenge of discovering the underlying equation using this approach is dealing with large , possible correlated errors in the features containing derivatives ; using numerical differentiation to calculate these higher-order derivatives accurately from noisy and sparse data is extremely challenging , even after denoising . One line of work has focused on constructing more robust and sparser approaches to solving eq . 1 , for example SR3 ( Zheng et al. , 2019 ) or using stability criteria ( Maddu et al. , 2020 ) . Alternatively , several works ( Both et al. , 2019 ; Berg & Nyström , 2019 ) have explored the use of neural networks to create a surrogate of the data and perform model discovery on this surrogate instead . Automatic differentiation can then be used to calculate the derivatives , yielding much more accurate features . PINNs Physics Informed Neural Networks ( PINNs ) ( Raissi et al. , 2017 ) have become a very popular method to either solve a differential equation or perform parameter inference with neural networks . Here we focus on parameter inference and consider a ( noisy ) dataset { ui , xi , ti } Ni=1 , governed by a differential equation of form ∂tu = Xw with X the terms of the equation and w the unknown parameters . PINNs inferw by using a neural network to approximate u , and constrain the network to the given differential equation by minimising LPINN ( θ , w ) = 1 N N∑ i=1 ‖ûi − ui‖2 + λ N N∑ i=1 ‖∂tûi −Xiw‖2 . ( 2 ) Here ûi = gθ ( xi , ti ) is the prediction of the neural network and λ sets the strength of the regularisation . The constraint ensures the network approximates the data u consistently with the given differential equation , and terms containing derivatives can be calculated using automatic differentiation . These two features make PINNs especially useful in noisy and sparse datasets . Model discovery with PINNs PINNs can easily be adapted to perform model discovery by replacing the given differential equation X with a larger set of candidate features Θ. Additionally , a maskm is applied to the coefficients , yielding as loss function , LMD ( θ , ξ ) = 1 N N∑ i=1 ‖ûi − ui‖2 + λ N N∑ i=1 ‖∂tûi −Θi ( m ξ ) ‖2 = Lfit ( θ ) + λLreg ( θ , ξ ) . ( 3 ) The maskm describes which terms feature in the equation , and hence the form of the constraint ; this approach can be interpreted as a PINN in which the constraint is also learned . The mask is updated periodically by some sparse regression technique , and as terms are pruned , the constraint becomes stricter , preventing overfitting of the constraint itself and improving the approximation of the network , boosting performance significantly ( Both et al. , 2021 ) . However , the non-differentiability of the mask can lead to issues during training , for example when it is updated at the wrong time , or when the wrong terms are accidentally pruned . Our goal here is thus to construct an approach where the mask M is learned together with the networks parameters , while still maintaining the benefits of iteratively refining the approximation : a fully-differentiable model discovery algorithm . Removing free variables Training with eq . 3 optimises two sets of parameters : the network parameters θ and the coefficients ξ . Typically both are minimised together using gradient descent , but the optimisation of ξ can be performed analytically ( Both et al. , 2021 ) . Given a configuration of the network parameters θ∗ , the minimisation over ξ is a regression problem as given by eq . 1 and can be solved exactly . Referring to this solution as the maximum likelihood estimate ξMLE , we define a loss function L̃MD ( θ ) ≡ LMD ( θ , ξMLE ) , which optimises only the network parameters θ using gradient descent . This significantly speeds up convergence and reduces the variance of the discovered coefficients across initialisations . We shall adopt this approach in the rest of the paper , and define the convention of L̃ denoting the loss functionLwith the independent variables calculated analytically . 3 FULLY DIFFERENTIABLE MODEL DISCOVERY . Our goal is to create a fully-differentiable model discovery algorithm , which , considering eq . 3 , requires making the maskm differentiable . Differentiable masking is challenging due to the binary nature of the problem , and instead we relax the application of the mask to a regularisation problem . Specifically , we propose to use Sparse Bayesian Learning ( Tipping , 2001 ) to select the active features and act as constraint . We start this section by reviewing SBL and how it can be used for differentiable variable selection , next show to integrate it in PINNs and finally introduce Physics Informed Normalizing Flows . 3.1 DIFFERENTIABLE MASKING WITH SBL . Sparse Bayesian Learning Sparse Bayesian Learning ( SBL ) ( Tipping , 2001 ) is a Bayesian approach to regression yielding sparse results . SBL defines a hierarchical model , starting with a Gaussian likelihood with noise precision β ≡ σ−2 , and a zero-mean Gaussian with precision αj on each component ξj as prior , p ( ∂tû ; Θ , ξ , β ) = N∏ i=1 N ( ∂tûi ; Θiξ , β−1 ) , ( 4 ) p ( ξ ; A ) = M∏ j=1 N ( ξj ; 0 , α−1j ) , ( 5 ) with ∂tû ∈ RN , Θ ∈ RN×M , ξ ∈ RM , and we have defined A ≡ diag ( α ) . The posterior distribution of ξ is a Gaussian with mean µ and covariance Σ , Σ = ( βΘTΘ +A ) −1 µ = βΣΘT∂tû . ( 6 ) Many of the terms in A will go to infinity when optimised , and correspondingly the prior for term j becomes a delta peak . We are thus certain that that specific term is inactive and can be pruned from the model . This makes SBL a very suitable choice for model discovery , as it gives a rigorous criterion for deciding whether a term is active or not . Additionally it defines hyper-priors over α and β , p ( α ) = M∏ j=1 Γ ( αj ; a , b ) p ( β ) = Γ ( β ; c , d ) ( 7 ) The inference of A and β can not be performed exactly , and SBL uses type-II maximum likelihood to find the most likely values of Â and β̂ by minimising the negative log marginal likelihood 1 , LSBL ( A , β ) = 1 2 [ β ‖ut −Θµ‖2 + µTAµ− log|Σ| − log|A| −N log β ] − M∑ j=1 ( a logαj − bαj ) − c log β + dβ , ( 9 ) using an iterative method ( see Tipping ( 2001 ) ) . Continuous relaxation The marginal likelihood also offers insight how the SBL provides differentiable masking . Considering only the first two terms of eq . 9 , β ‖ut −Θµ‖2 + µTAµ ( 10 ) we note that the SBL essentially applies a coefficient-specific ` 2 penalty to the posterior mean µ . If Aj → ∞ , the corresponding coefficient µj → 0 , pruning the variable from the model . Effectively , the SBL replaces the discrete mask mj ∈ { 0 , 1 } by a continuous regularisation Aj ∈ ( 0 , ∞ ] , and we thus refer to our approach as continuous relaxation .

Summary: This paper studies the problem of identifying sets of equations that govern the dynamics of observed natural systems. Building off of work that uses sparse regression to identify a small set of active terms from a large dictionary of possible terms, this paper proposes a soft relaxation of a hard constraint used for selecting out possible terms in the optimization problem. The paper applies this technique to recover the underlying equations from data generated by a few canonical PDEs.