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: This manuscript introduces a new optimization fashion bridging the Gauss-Newton method and the vanilla gradient descent method in the middle for physical optimization with neural networks. Traditionally, when a neural network is used for a physical problem, it will inevitably be affected by the unbalanced magnitudes dramatically. This manuscript proposes HIG, which is a middle point of two popular and individually advantageous optimization methods. It also provides a nice guideline on how to choose the hyperparameters and examine the efficacy of HIG on a concept-illustrating synthetic problem and three realistic physical problems.