Half-Inverse Gradients for Physical Deep Learning

1 INTRODUCTION . The groundbreaking successes of deep learning ( Krizhevsky et al. , 2012 ; Sutskever et al. , 2014 ; Silver et al. , 2017 ) have led to ongoing efforts to study the capabilities of neural networks across all scientific disciplines . In the area of physical simulation , neural networks have been used in various ways , such as creating accurate reduced-order models ( Morton et al. , 2018 ) , inferring improved discretization stencils ( Bar-Sinai et al. , 2019 ) , or suppressing numerical errors ( Um et al. , 2020 ) . The long-term goal of these methods is to exceed classical simulations in terms of accuracy and speed , which has been achieved , e.g. , for rigid bodies ( de Avila Belbute-Peres et al. , 2018 ) , physical inverse problems ( Holl et al. , 2020 ) , and two-dimensional turbulence ( Kochkov et al. , 2021 ) . The successful application of deep learning to physical systems naturally hinges on the training setup . In recent years , the use of physical loss functions has proven beneficial for the training procedure , yielding substantial improvements over purely supervised training approaches ( Tompson et al. , 2017 ; Wu & Tegmark , 2019 ; Greydanus et al. , 2019 ) . These improvements were shown to stem from three aspects ( Battaglia et al. , 2016 ; Holl et al. , 2020 ) : ( i ) Incorporating prior knowledge from physical principles facilitates the learning process , ( ii ) the ambiguities of multimodal cases are resolved naturally , and ( iii ) simulating the physics at training time can provide more realistic data distributions than pre-computed data sets . Approaches for training with physical losses can be divided into two categories . On the one hand , equation-focused approaches that introduce physical residuals ( Tompson et al. , 2017 ; Raissi et al. , 2019 ) , and on the other hand , solver-focused approaches that additionally integrate well-established numerical procedures into the training ( Um et al. , 2020 ; Kochkov et al. , 2021 ) . From a mathematical point of view , training a neural network with a physical loss function bears the difficulties of both network training and physics optimization . In order to obtain satisfying results , it is vital to treat flat regions of the optimization landscapes effectively . In learning , the challenging loss landscapes are addressed using gradient-based optimizers with data-based normalizing schemes such as Adam ( Kingma & Ba , 2015 ) , whereas in physics , the optimizers of choice are higher-order techniques such as Newton ’ s method ( Gill & Murray , 1978 ) that inherently make use of inversion processes . However , Holl et al . ( 2021 ) found that these approaches can not effectively handle the joint optimization of network and physics . Gradient-descent-based optimizers suffer from vanishing or exploding gradients , preventing effective convergence , while higher-order methods do not generally scale to the high-dimensional parameter spaces required by deep learning ( Goodfellow et al. , 2016 ) . We likewise focus on inversion-based approach but propose a new method for joint physics and network optimization which we refer to as half-inverse gradients . At its core lies a partial matrix inversion , which we derive from the interaction between network and physics both formally and geometrically . An important property of our method is that its runtime scales linearly in the number of network parameters . To demonstrate the wide-ranging and practical applicability of our method , we show that it yields significant improvements in terms of convergence speed and final loss values over existing methods . These improvements are measured both in terms of absolute accuracy as well as wall-clock time . We evaluate a diverse set of physical systems , such as the Schrödinger equation , a nonlinear chain system and the Poisson problem . 2 GRADIENTS BASED ON HALF-INVERSE JACOBIANS . Optimization on continuous spaces can be effectively performed with derivative-based methods , the simplest of which is gradient descent . For a target function L ( θ ) to be minimized of several variables θ , using bold symbols for vector-valued quantitites in this section , and learning rate η , gradient descent proceeds by repeatedly applying updates ∆θGD ( η ) = −η · ( ∂L ∂θ ) > . ( 1 ) For quadratic objectives , this algorithm convergences linearly with the rate of convergence depending on the condition number λ of the Hessian matrix ( Lax , 2014 ) . In the ill-conditioned case λ 1 , flat regions in the optimization landscape can significantly slow down the optimization progress . This is a ubiquitous problem in non-convex optimization tasks of the generic form : L ( θ ) = ∑ i l ( yi ( θ ) , ŷi ) = ∑ i l ( f ( xi ; θ ) , ŷi ) ( 2 ) Here ( xi , ŷi ) denotes the ith data points from a chosen set of measurements , f is a function parametrized by θ to be optimized to model the relationship between the data points yi ( θ ) = f ( xi ; θ ) , and l denotes a loss function measuring the optimization progress . In the following , we assume the most common case of l ( yi , ŷi ) = 12 ||yi − ŷi|| 2 2 being the squared l2-loss Physics Optimization . Simulating a physical system consists of two steps : ( i ) mathematically modeling the system by a differential equation , and ( ii ) discretizing the occuring differential operators to receive a solver usable on a computer . If we then wish to manipulate a physical system by an external force to reach a given configuration , we have to solve an inverse problem of the form 2 . In such a control task , the sum reduces to a single data point ( x , ŷ ) with x being the initial state , ŷ the target state and θ the external force we want to find . The physical solver corresponds to the function f representing time evolution y ( θ ) = f ( x ; θ ) . This single data point sum still includes squared summation over vector components of y − ŷ in the l2-loss . Sensitive behavior of the physical system connected to its high-frequency modes is present in the physical solver f , and produces small singular values in its linear approximation . This leads to an ill-conditioned Jacobian and flat regions in the optimization landscape when minimizing 2 . This is addressed by using methods that incorporate more information than only the gradient . Prominent examples are Newton ’ s method or the Gauss-Newton ’ s algorithm ( Gill & Murray , 1978 ) ; the latter one is based on the Jacobian of f and the loss gradient : ∆θGN = − ( ∂y ∂θ ) −1 · ( ∂L ∂y ) > ( 3 ) Here , the inversion of the Jacobian is done with the pseudoinverse . The Gauss-Newton update maps the steepest descent direction in y space to the parameter space θ . Therefore , to first order , the resulting update approximates gradient descent steps in y space , further details of which are given in the appendix . An advantage of this formulation is that the update steps in y space are invariant under arbitrary rescaling of the parameters θ , which cancels inherent scales in f and ensures quick progress in the optimization landscape . Neural Network Training . When f represents a neural network in equation 2 , the optimization matches the typical supervised learning task . In this context , the problem of flat regions in the optimization landscape is also referred to as pathological curvature ( Martens , 2010 ) . Solving this problem with higher-order methods is considered to be too expensive given the large number of parameters θ . For learning tasks , widely used optimizers such as Adam instead use gradient information from earlier update steps , for instance in the form of momentum or adaptive learning rate terms , thereby achieving better optimization progress at little additional computational cost . Additionally , the updates are computed on mini-batches instead of the full data set , which saves computational resources and benefits generalization . Neural Network Training with Physics Objectives . For the remainder of the paper , we consider joint optimization problems , where f denotes a composition of a neural network parameterized by θ and a physics solver . Using classical network optimizers for minimizing equation 2 is inefficient in this case since data normalization in the network output space is not possible and the classical initialization schemes can not normalize the effects of the physics solver . As such , they are unsuited to capture the strong coupling between optimization parameters typically encountered in physics applications . While Gauss-Newton seems promising for these cases , the involved Jacobian inversion tends to result in large overshoots in the updates when the involved physics solver is ill-conditioned . As we will demonstrate , this leads to oversaturation of neurons , hampering the learning capability of the neural network . 2.1 AN ILL-CONDITIONED TOY EXAMPLE . To illustrate the argumentation so far , we consider a data set sampled from ŷ ( x ) = ( sin ( 6x ) , cos ( 9x ) ) for x ∈ [ −1 , 1 ] : We train a neural network to describe this data set by using the loss function l ( y , ŷ ; γ ) = 12 ( y1 − ŷ1 ) 2 + 12 ( λ · y2 − ŷ2 ) 2 , where y1 and y2 denote the first and second component of the y-vectors . For a scale factor of γ = 1 , we receive the well-conditioned mean squared error loss . However , l becomes increasingly ill-conditioned as γ is decreased , which mimics the effects of a physics solver . For real-world physics solvers the situation would typically be even more pronounced , as these scales usually vary strongly in direction and magnitude across different data points and optimization steps . We use a small neural network with a single hidden layer with 7 neurons and a tanh activation . We then compare training with the well-conditioned γ = 1 loss against an ill-conditioned γ = 0.01 loss . In both cases , we train the network using both Adam and Gauss-Newton as representatives of gradient-based and higher-order optimizers , respectively . The results are shown in figure 1 . In the well-conditioned case , Adam and Gauss-Newton behave similarly , decreasing the loss by about three orders of magnitude . However , in the ill-conditioned case , both optimizers fail to minimize the objective beyond a certain point . To explain this observation , we first illustrate the behavior from the physics viewpoint by considering the trajectory of the network output f ( x ) for a single value x during training ( figure 1 , right ) . For γ= 1 , Adam optimizes the network to accurately predict ŷ ( x ) while for γ= 0.01 , the updates neglect the second component preventing Adam to move efficiently along the small-scale coordinate ( blue curve in figure 1b , right ) . To illustrate the situation from the viewpoint of the network , we consider the variance in the outputs of specific neurons over different x ( figure 1 , middle ) . When γ = 1 , all neurons process information by producing different outcomes for different x . However , for γ = 0.01 , Gauss-Newton ’ s inversion of the small-scale component y2 results in large updates , leading to an oversaturation of neurons ( red curve in figure 1b , middle ) . These neurons stop processing information , reducing the effective capacity of the network and preventing the network from accurately fitting ŷ . Facing these problems , a natural questions arises : Is it possible to construct an algorithm that can successfully process the inherently different scales of a physics solver while training a neural network at the same time ?

Summary: The paper proposes a way to interpolate between gradient descent and Gauss-Newton's method for solving nonlinear least squares problems arising from physics informed training. The authors cite past work that physics informed training is often ill conditioned, so gradient descent often performs poorly. They give a classic example of ill-conditioning, and show that gradient descent converges slowly, while Gauss-Newton quickly saturates neuron activations.  This motivates them to introduce Half-Inverse Gradients (HIG), which interpolates between GD and GN using the SVD. The authors then try the method on several test problems in scientific computing: Control of nonlinear oscillators, Poisson, and control of a Quantum dipole. They compare Adam, HIG, and GN.