PDE-regularized Neural Networks for Image Classification

1 INTRODUCTION . It had been discovered that interpreting neural networks as differential equations is possible by several independent research groups ( Weinan , 2017 ; Ruthotto & Haber , 2019 ; Lu et al. , 2018 ; Ciccone et al. , 2018 ; Chen et al. , 2018 ; Gholami et al. , 2019 ) . Among them , the seminal neural ordinary differential equation ( neural ODE ) research work , which considers the general architecture in Figure 1 ( a ) , is to learn a neural network approximating ∂h ( t ) ∂t , where h ( t ) is a hidden vector at layer ( or time ) t ( Chen et al. , 2018 ) . As such , a neural network is described by a system of ODEs , each ODE of which describes a dynamics of a hidden element . While neural ODEs have many good characteristics , they also have limitations , which are listed as follows : Pros . Neural ODEs can interpret t as a continuous variable and we can have hidden vectors at any layer ( or time ) l by h ( l ) = h ( 0 ) + ∫ l 0 o ( h ( t ) , t ; θo ) dt , where o ( h ( t ) , t ; θo ) = ∂h ( t ) ∂t is a neural network parameterized by θo . Pros . Neural ODEs sometimes have smaller numbers of parameters than those of other conven- tional neural network designs , e.g. , ( Pinckaers & Litjens , 2019 ) . Cons . Neural ODEs , which use an adaptive step-size ODE solver , sometimes show numerical instability ( i.e. , the underflow error of the step-size ) or their forward-pass inference can take a long time ( i.e. , too many steps ) in solving integral problems , e.g , a forward-pass time of 37.6 seconds of ODE-Net vs. 9.8 seconds of PR-Net in Table 2 . Several countermeasures have been proposed but it is unavoidable to solve integral problems ( Zhuang et al. , 2020 ; Finlay et al. , 2020 ; Daulbaev et al. , 2020 ) . To tackle the limitation , we propose the concept of partial differential equation-regularized neural network ( PR-Net ) to directly learn a hidden element , denoted h ( d , t ) at layer ( or time ) t ∈ [ 0 , T ] and dimension d ∈ Rm . Under general contexts , a PDE consists of i ) an initial condition at t = 0 , ii ) a boundary condition at a boundary location of the spatial domain Rm , and iii ) a governing equation describing ∂h ( d , t ) ∂t . As such , learning a PDE from data can be reduced to a regression-like problem to predict h ( d , t ) that meets its initial/boundary conditions and governing equation . In training our proposed PR-Net , h ( 0 ) is provided by an earlier feature extraction layer , which is the same as neural ODEs . However , an appropriate governing equation is unknown for downstream machine learning tasks . Therefore , we propose to train a regression model for predicting h ( d , t ) and its governing equation simultaneously ( see Figure 1 ( b ) ) . In other words , neural ODEs directly learn a governing equation ( i.e. , ∂h ( t ) ∂t ) , whereas PR-Net learns a governing equation in conjunction with a regression model that conforms with the learned governing equation . The key advantage in our approach is that we can eliminate the necessity of solving integral problems — in neural ODEs , where we learn a governing equation only , solving integral problems is mandatory . Such forward and inverse problems ( i.e. , solving PDEs for h ( d , t ) and identifying governing equations , respectively ) arise in many important computational science problems and there have been many efforts applying machine learning/deep learning techniques to those problems ( e.g. , in earth science ( Reichstein et al. , 2019 ; Bergen et al. , 2019 ) and climate science ( Rolnick et al. , 2019 ) ) . Recently , physics-informed or physics-aware approaches ( Battaglia et al. , 2016 ; Chang et al. , 2017 ; de Bezenac et al. , 2018 ; Raissi et al. , 2019 ; Sanchez-Gonzalez et al. , 2018 ; Long et al. , 2018 ) have demonstrated that designing neural networks to incorporate prior scientific knowledge ( e.g. , by enforcing physical laws described in governing equations ( Raissi et al. , 2019 ) ) greatly helps avoiding over-fitting and improving generalizability of the neural networks . There also exist several approaches to incorporate various ideas of classical mechanics in designing neural-ODE-type networks ( Greydanus et al. , 2019 ; Chen et al. , 2020 ; Cranmer et al. , 2020 ; Zhong et al. , 2020 ; Lee & Parish , 2020 ) . However , all these works are interested in solving either forward or inverse problems whereas we solve the two different problem types at the same time for downstream tasks . The most similar existing work to our work is in ( Long et al. , 2018 ) . However , this work studied scientific PDEs and do not consider t as a continuous variable but use a set of discretized points of t. Compared to previous approaches , the proposed method has a distinct feature that forward and inverse problems are solved simultaneously with a continuous variable t. Due to this unique feature , the method can be applied to general machine learning downstream tasks , where we do not have a priori knowledge on governing equations , such as image classification . Our proposed PR-Net had the following characteristics : Pros . PR-Net trains a regression model that outputs a scalar element h ( d , t ) ( without solving any integral problems ) , and we can consider both d and t as continuous variables . Therefore , it is possible to construct flexible hidden dimension vectors . Pros . PR-Net does not require solving integral problems . As such , there is no numerical instability and their forward-pass time is much shorter than that of neural ODEs . Pros . By learning a governing equation , we can regularize the overall behavior of PR-Net . Cons . PR-Net sometimes requires a larger number of parameters than that of neural ODEs or conventional neural networks . 2 PARTIAL DIFFERENTIAL EQUATIONS . The key difference between ODEs and PDEs is that PDEs can have derivatives of multiple variables whereas ODEs should have only one such variable ’ s derivative . Therefore , our PDE-based method interprets both the layer of neural network and the dimension of hidden vector as continuous variables , which can not be done in neural ODEs . In our context , h ( d , t ) means a hidden scalar element at layer t ∈ R and dimension d ∈ Rm , e.g. , m = 1 if h ( t ) is a vector , m = 3 if h ( t ) is a convolutional feature map , and so on . h ( d,0 ) Governing Equation Solution h ( d , t ) Neural Network Figure 2 : A neural network predicts solution values at d , t given initial conditions , denoted h ( d , 0 ) for various d , and a governing equation . Table 1 : Two types of PDE problems related to our work Type Data What to infer Forward Problem - Initial condition- Governing equation Solution h ( d , t ) Inverse Problem - Solution h ( d , t ) - Initial condition Governing equation In this section , we first introduce the forward and inverse problems of PDEs in general contexts ( see Table 1 ) . Then , we extend them to design our proposed method in deep-learning contexts . 2.1 FORWARD PROBLEM OF PDES IN GENERAL CONTEXTS . The forward PDE problem in general contexts is to find a solution h ( d , t ) , where d is in a spatial domain Rm and t is in a time domain [ 0 , T ] , given i ) an initial condition h ( d , 0 ) , ii ) a boundary condition h ( dbc , t ) , where dbc is a boundary location of the spatial domain Rm , and iii ) a governing equation g ( Raissi et al. , 2019 ) We note that the boundary condition can be missing in some cases ( Kim , 2018 ) . The governing equation is typically in the following form with particular choices of αi , j ( Raissi , 2018 ; Peng et al. , 2020 ) : g ( d , t ; h ) def = ht − ( α0,0 + α1,0h+ α2,0h 2 + α3,0h 3 + α0,1hd + α1,1hhd + α2,1h 2hd + α3,1h 3hd + α0,2hdd + α1,2hhdd + α2,2h 2hdd + α3,2h 3hdd + α0,3hddd + α1,3hhddd + α2,3h 2hddd + α3,3h 3hddd ) , ( 1 ) where ht = ∂h ( d , t ) ∂t , hd = ∂h ( d , t ) ∂d , hdd = ∂2h ( d , t ) ∂d2 , and hddd = ∂3h ( d , t ) ∂d3 . We also note that g is always zero in all PDEs , i.e. , g ( d , t ; h ) = 0 . In many cases , it is hard to solve the forward problem and hence general purpose PDE solvers do not exist . Nevertheless , one can use the following optimization to train a neural network f ( d , t ; θ ) to approximate the solution function h ( d , t ) as shown in Figure 2 ( Raissi et al. , 2019 ) : argmin θ LI + LB + LG , ( 2 ) LI def = 1 NI ∑ d ( f ( d , 0 ; θ ) − h ( d , 0 ) ) 2 , ( 3 ) LB def = 1 NB ∑ ( dbc , t ) ( f ( dbc , t ; θ ) − h ( dbc , t ) ) 2 , ( 4 ) LG def = 1 NG ∑ ( d , t ) g ( d , t ; f , θ ) 2 , ( 5 ) where NI , NB , NG are the numbers of training samples , LI is to train θ for the initial condition , LB is for the boundary condition , and LG is for the governing equation . Because the governing equation is always zero , we simply minimize its squared term . Note that i ) ft , fd , fdd , fddd can be easily constructed using the automatic differentiation implemented in TensorFlow or PyTorch , and ii ) we only need h ( d , 0 ) , h ( dbc , t ) , which are known a priori , to train the parameters θ . 2.2 INVERSE PROBLEM OF PDES IN GENERAL CONTEXTS . The inverse problem is to find a governing equation given i ) an initial condition h ( d , 0 ) and ii ) a solution function h ( d , t ) ( Raissi , 2018 ) . It learns αi , j in Eq . 1 with the following loss ( if possible , they use reference solutions as well ) : argmin αi , j 1 NG ∑ ( d , t ) g ( d , t ; h ) 2 . Given a solution function h and its partial derivative terms , we train αi , j by minimizing the objective loss . Note that we know h in this case . Therefore , the objective loss is defined with h rather than with f , unlike Eq . 5 . The optimal solution of αi , j is not unique sometimes . However , we note that no trivial solutions , e.g. , αi , j = 0 for all i , j , exist for the inverse problem .

Summary: The paper proposed the method of neural PDE as an improvement of neural ODE. In specific, neural PDE considers both the layer and the hidden dimension as continuous variables of the PDE. The new part of neural PDE compared to neural ODE is essentially solving PDE inverse problems (learning PDE from data) in the computational mathematics and engineering community, and the way of learning PDE (by embedding the PDE and initial condition into the loss function via automatic differentiation) is the physics-informed neural network (PINN) proposed in [Raissi et al., JCP, 2019]. The experiments show that compared to neural ODE, neural PDE achieves comparable accuracy but with less forward-pass inference time; but these experiments are not convincing enough.