Physical Gradients for Deep Learning

1 INTRODUCTION . Inverse problems that involve physical systems play a central role in computational science . This class of problems includes parameter estimation ( Tarantola , 2005 ) and optimal control ( Zhou et al. , 1996 ) . Solving inverse problems is integral in determining the age of the universe ( Freedman et al. , 2001 ) , measuring the amount of dark matter ( Padmanabhan & White , 2009 ) , detecting gravitational waves ( George & Huerta , 2018b ) , controlling plasma flows ( Maingi et al. , 2019 ) , searching for neutrinoless double-beta decay ( Agostini et al. , 2013 ; Aalseth et al. , 2018 ) , reconstructing particle trajectories and energies ( Belayneh et al. , 2020 ) , and testing general relativity ( Dyson et al. , 1920 ; Kraniotis & Whitehouse , 2003 ) . Decades of research in optimization have produced a wide range of iterative methods and algorithms for solving inverse problems ( Press et al. , 2007 ) . Higher-order methods such as limited-memory BFGS ( Liu & Nocedal , 1989 ) have been especially successful . These methods have two fundamental drawbacks . First , they require a good initial guess of the solution . When the initial guess is far away from the true solution in parameter space , iterative methods may diverge or converge to a local optimum instead . Second , their computational cost is strongly problem-dependent , as the gradient and possibly even Hessian matrix need to be evaluated for each iteration of the optimization . This may be negligible for simple tasks but when the optimization involves a physical simulation , each iteration can take seconds or longer . However , many experiments continuously or repeatedly collect data , yielding data sets with millions or billions of individual events , each posing an inverse problem . Solving all of these with iterative solvers poses a major challenge which is why , in such settings , data analysis has often turned to machine learning algorithms to quickly infer approximate solutions given the observations ( Carleo et al. , 2019 ; Delaquis et al. , 2018 ; George & Huerta , 2018b ; Agostini et al. , 2013 ) . The inferred solutions can then , for example , be used to filter data sets for interesting events or can serve as the starting point for a more accurate analysis ( CMS Collaboration , 2020 ; George & Huerta , 2018a ; b ) . Interestingly , using machine learning models to solve the inverse problems also reduces the chances of divergence or convergence to undesired local optima . Unlike iterative solvers , learning models are jointly trained on large collections of data . The shared parameters allow updates from one example to avoid local minima from a different example . Fig . 1 shows a generic case illustrating this effect : The task is to localize a wave packetA·sin ( f ·x ) ·exp ( − 12 ( x−x0 ) 2/σ2 ) and determine its amplitude from a noisy recorded time series ( grey curve ) . Iterative solvers fail to reconstruct the wave packet when the initial guess is not already close to the true parameters , even diverging in 0.12 % of our cases . In contrast , a neural network trained on the same data using the same objective quickly learns to fit all wave packets . Details are given in appendix B.1 . In this paper , we derive a hybrid optimization algorithm that integrates higher-order inverse solvers into the traditional deep learning pipeline without imposing these limitations and without discarding the progress made in first-order optimization . Our method can work with arbitrary second-order solvers when the solution space of the inverse problem allows for computation of the Hessian . Otherwise , it requires additional domain knowledge which is used to formulate the inverse solver in a more efficient manner . Our method computes the weight updates using both backpropagation as well as the embedded inverseproblem solver . Specifically , the first-order gradient of the physical process used in backpropagation is replaced by a physical gradient ( PG ) derived from the update of the inverse-problem solver . When used in conjunction with training deep learning models , physical gradients allow each update step to encode nonlinear physics , which yields substantially more accurate optimization directions . A key advantage of our approach is its compatibility with acceleration schemes ( Duchi et al. , 2011 ; Kingma & Ba , 2015 ) and stabilization techniques ( Ioffe & Szegedy , 2015 ; Huang et al. , 2016 ; Ba et al. , 2016 ) developed for training deep learning models . The presented hybrid training approach differs fundamentally from pure first-order training . While PGs take the place of regular gradients in the optimization algorithm , they are based on real physical states instead of the adjoint ones calculated in backpropagation . This feature makes it possible to integrate solution priors into the training without restricting the model architecture ( Kim et al. , 2019 ) , adding soft constraints to the objective ( Raissi et al. , 2018 ) , or relying exclusively on firstorder updates ( Rackauckas et al. , 2020 ) . We test our method on a wide variety of inverse problems including the highly challenging Navier-Stokes equations . In all of our experiments , neural networks trained with PGs significantly outperform state-of-the-art training schemes , yielding accuracy gains of up to three orders of magnitude . 2 METHOD . We consider unconstrained inverse problems that involve a differentiable physical process P : X ⊂ Rdx → Y ⊂ Rdy which can be simulated . Here X denotes the physical parameter space and Y the space of possible observed outcomes . Given an observed or desired output y∗ ∈ Y , the inverse problem consists of finding optimal parameters x∗ = arg min x L ( x ) with L ( x ) = 1 2 ‖P ( x ) − y∗‖22 . ( 1 ) We are interested in approximating this problem using a neural network , x∗ = NN ( y∗ | θ ) , parameterized by θ . Let Y∗ = { y∗i | i = 1 , ... , N } denote a set of N inverse problems , all sharing the same physical process P . Then training the network is equivalent to finding θ∗ = arg min θ N∑ i=1 1 2 ‖P ( NN ( y∗i | θ ) ) − y∗i ‖22 . ( 2 ) This training scheme is unsupervised , i.e . no labels are required over the training data Y . 2.1 ITERATIVE OPTIMIZATION . Inverse problems as in Eq . 1 are classically solved by starting with an initial guess x0 and iteratively applying updates ∆xk . This yields an estimate of the inverse of the physical model P , i.e. , x∗ ≈ P−1n∗ ( y∗ |x0 ) where n∗ denotes the number of iterations required to reach a chosen convergence threshold . A key decision in such an iterative optimization scheme is how to derive the update ∆xk . Figure 2a visually compares gradient descent , Newton ’ s method and PGs which we will introduce in section 2.2 . First-order methods The scale of typical machine learning datasets necessitates the use of minibatches , which has led popular optimization methods in machine learning to rely only on firstorder derivatives of L. These Hessian-free methods are popular in deep learning due to their low computational cost and stable convergence ( Goodfellow et al. , 2016 ) . The simplest such method is gradient descent ( GD ) ( Curry , 1944 ) and its stochastic version ( SGD ) where ∆x = −η· ( ∂L ∂x ) T . Note that ∂L∂x , and therefore of ∆x , scales inversely with x . When x carries physical units , neither ∆x nor ∆xT lie in the same space as x . This scaling behavior is a problem when dealing with sensitive or insensitive functions , i.e . functions whose gradient magnitudes ∣∣∂L ∂x ∣∣ are far from 1 . When small changes in x cause large changes in L , this counter-intuitively leads to large updates ∆x when using GD , resulting in even larger changes to L. This is often referred to as exploding gradients , and the opposite effect , vanishing gradients , occurs with insensitive functions . Machine learning models can be tuned to behave well given these updates ( Ioffe & Szegedy , 2015 ; Loshchilov & Hutter , 2019 ) but in physics-based optimization , where functions may be extremely sensitive in certain parts of the parameter space or on subsets of training examples , GD often prescribes suboptimal optimization directions . While there are a number of first-order optimizers that approximate higher-order information to improve convergence ( Kingma & Ba , 2015 ; Hestenes et al. , 1952 ; Duchi et al. , 2011 ; Martens & Grosse , 2015 ) , they can not take nonlinearities into account , which results in suboptimal optimization directions when optimizing nonlinear functions . Second-order methods Newton ’ s method ( Atkinson , 2008 ) uses the inverse Hessian to determine the optimization update ∆x = −η · ( ∂2L ∂x2 ) −1 ( ∂L ∂x ) T . Direct computation of the Hessian adds a significant computational cost to each iteration . Quasi-Newton methods ( Broyden , 1970 ; Liu & Nocedal , 1989 ; Gill & Murray , 1978 ; Moré , 1978 ; Powell , 1970 ; Berndt et al. , 1974 ; Conn et al. , 1991 ; Avriel , 2003 ) alleviate this issue by approximating the inverse Hessian instead . Unlike GD , Newton-type methods prescribe small updates to sensitive functions and large updates to insensitive ones . The resulting update steps typically progress towards optima much faster than first-order methods ( Ye et al. , 2019 ) . However , the difficulty of approximating the Hessian for batched data sets ( Schraudolph et al. , 2007 ) and the high cost of evaluating the Hessian directly have largely prevented second-order methods from being utilized in machine learning research ( Goodfellow et al. , 2016 ) . Domain-specific inverse solvers Domain knowledge can often be used to formulate more efficient optimization updates or to directly approximate a solution . A classical example would be preconditioners that , when chosen correctly , allow linear solvers to converge much more quickly and reliably . In nonlinear settings , part of the governing equations may be inverted analytically or priors about the solution space may be incorporated into the solver , for example . A time-reversed simulation may be able to more accurately recover prior states than a generic inverse solver . These methods are commonly used in optimization but their use has eluded machine learning applications so far due to the fixed gradient-descent pipeline . When learning with PGs , domain-specific optimizers can be integrated into training like any other optimizers , allowing for greatly-increased training convergence . Appendix appendix A.2 lists necessary conditions that need to be fulfilled . 2.2 PHYSICAL GRADIENTS FOR DEEP LEARNING . We now consider the problem of learning solutions to the inverse problems as in Eq . 2 . As discussed above , GD-based optimizers fail to take the nonlinearity and potentially strongly varying sensitivity of the joint problem into account while higher-order optimizers are hard to apply directly . Instead , we aim to leverage the advantages of higher-order optimizers in a deep learning setting by embedding arbitrary inverse-problem solvers as physical gradients ( PGs ) into the first-order optimization pipeline . The network updates ∆θ can then be computed using any first-order optimizer such as SGD or Adam . Additionally , all state-of-the-art deep learning techniques , such as normalization Ioffe & Szegedy ( 2015 ) ; Ba et al . ( 2016 ) or dropout Srivastava et al . ( 2014 ) can be employed , which often improves training speed and generalization performance Keskar et al . ( 2017 ) . Our method can be applied whenever one of the following is available : ( i ) ∂ 2P ∂x2 can be computed numerically , ( ii ) a higher-order update such as ( ∂2P ∂x2 ) −1 ∂P ∂x can be derived analytically , or ( iii ) an inverse-physics solver can be derived using the available domain knowledge . If ( i ) is fulfilled , we can explicitly compute the Newton direction in in x space . Thus , with any of ( i ) , ( ii ) or ( iii ) we have access to better updates ∆x than ∆xGD ∝ ∂P∂x . These ∆x are the basis for our method which we describe in the following . We derive our method starting from Eq . 2 which describes the combined problem of finding both solutions x∗ to the individual inverse problems y∗ ∈ Y∗ and network weights θ that approximate them . One way to solve this problem is to first consider all inverse problems separately and precompute corresponding solutions X sv = { xsvi : P ( xsvi ) = y∗i } using an iterative optimizer . Then , X sv could be used as labels for supervised training of a neural network : θ∗ = arg min θ N∑ i=1 1 2 ‖NN ( y∗i | θ ) − xsvi ‖22 . ( 3 ) This precomputation makes it possible to use an efficient optimization method forX sv while training the neural network with a first-order optimizer , retaining the ability to employ state-of-the-art deep learning techniques , such as normalization ( Ioffe & Szegedy , 2015 ; Ba et al. , 2016 ) or dropout ( Srivastava et al. , 2014 ) . However , this approach has severe drawbacks . For one , the individual optimizations can more easily get stuck in local optima than the combined training ( Holl et al. , 2020 ) . Also , many inverse problems are inherently multi-modal , i.e . multiple solutions x∗ exist for one y∗ . In such settings xsv heavily depends on the initial guess x0 used for precomputation , which may cause the network to interpolate between possible solutions , leading to subpar convergence and generalization performance . To avoid these problems , we alter the training procedure from Eq . 3 in two ways . First , we treat the inverse problem solver P−1 as part of a coupled optimization and consider it as part of the training loop , yielding θ∗ = arg min θ N∑ i=1 1 2 ‖NN ( y∗i | θ ) − P−1n∗i ( y ∗ i ) ‖22 . ( 4 ) Next we can condition the embedded inverse problem solver on the neural network prediction , P−1n∗i ( y ∗ i ) → P −1 n∗i ( y∗i |NN ( y∗i | θ ) ) , by using it as an initial guess . The embedded optimizer can now find a minimum close to the prediction , which allows the neural network to choose any solution of a multi-modal problem . Also , since all inverse problems from Y∗ are optimized jointly , this reduces the likelihood that an individual solution gets stuck in a local minimum . It is furthermore not necessary to run the iterative inverse solver to convergence . Any choice of n ∈ N steps likely leads to convergence ( proof given in appendix A.2 ) , and choosing small n can increase the training speed . We also observe this in practice : Optimizations with n = 1 converged reliably in all our experiments . With these modifications , we arrive at θ∗ = arg min θ N∑ i=1 1 2 ‖NN ( y∗i | θ ) − P−1n ( y∗i |NN ( y∗i | θ ) ) ‖22 . ( 5 ) The central quantity here is the difference between prediction and correction obtained from P−1n , which we refer to as the physical gradient . It encodes a valid physical state and takes the place of ( ∂P ∂x ) T , which would otherwise be computed by backpropagation . Since the PG stems from a generic higher-order or a domain-specific solver , the resulting updates ∆θ will also be non-linear without computing the Hessian w.r.t . θ . Assuming the physics optimizer can account for varying function sensitivity in its updates ∆x , the resulting PG updates ∆θ will inherit that property since the network is generally tuned to retain gradient magnitude . The properties of PG as well as GD and Newton ’ s method are summarized in Fig . 2b . While the L2 measure in Eq . 5 looks like a supervised objective in x space , the actual network optimization does not behave as such . To see this more clearly , consider P−1 performing a single GD step with unit learning rate , P−1 ( y∗ |x ) = x− ( ∂P ∂x ) T ( P ( x ) − y∗ ) . Then the gradient for Eq . 5 is ( ∂P ∂x ∂x ∂θ ) T ( P ( x ) − y∗ ) , which is identical to the unsupervised training in Eq . 2 ( the proof is given in appendix A.1 ) . Algorithm 1 : Neural network training with physical gradients according to Eq . 5 for each training sample do x0 ← NN ( y∗ | θ ) Solution inference with NN for k = 1 , ... , n do xk ← P−1 ( y∗ |xk−1 ) Accumulate iterative updates for PG end ∆θ ← ( ∂x ∂θ ) T · ( x0 − xn ) T Backpropagation through NN θ ← Step ( θ , ∆θ , η ) Parameter update end Algorithm 1 details the corresponding neural network training procedure . There , the effective gradient ∆θ can also be written as ∂∂θ 1 2‖x0 − xn‖ 2 2 with xn taken as constant as in Eq . 5 , which lends itself to straightforward integration into machine learning frameworks where a simple L2 loss can be used to optimize the network . This L2 loss acts as a proxy to embed the physical gradient into the network training pipeline and is not to be confused with a traditional supervised loss in x space ( see appendix A.1 for a more detailed discussion ) .

Summary: This paper proposes an approach for learning neural network-based approximators for inverse problems in physical domains. The crux of this approach involves rewriting the training loss function to capture the difference between (a) the NN's estimate of the problem input and (b) the solution obtained when using this estimate as an initial point for an approximate inverse problem solver. This loss function is then optimized via gradient descent. This approach is demonstrated on several physical systems, and shown to avoid issues due to multi-modality (which supervised training schemes fall prey to), instability in the gradients through the physical process (faced in some settings by first-order training methods applied to the "standard" inverse problem), and issues such as expense of computation or ill-conditioning (faced in some settings by training method using e.g. inverse Hessian information).