Year: 2021

Abstract:
Recent works in deep learning have shown that integrating differentiable physics simulators into the training process can greatly improve the quality of results. Although this combination represents a more complex optimization task than supervised neural network training, the same gradient-based optimizers are typically employed to minimize the loss function. However, the integrated physics solvers have a profound effect on the gradient flow as manipulating scales in magnitude and direction is an inherent property of many physical processes. Consequently, the gradient flow is often highly unbalanced and creates an environment in which existing gradient-based optimizers perform poorly. In this work, we analyze the characteristics of both physical and neural network optimizations to derive a new method that does not suffer from this phenomenon. Our method is based on a halfinversion of the Jacobian and combines principles of both classical network and physics optimizers to solve the combined optimization task. Compared to state-ofthe-art neural network optimizers, our method converges more quickly and yields better solutions, which we demonstrate on three complex learning problems involving nonlinear oscillators, the Schrödinger equation and the Poisson problem.

heading: 1 INTRODUCTION
text: 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.

heading: 2 GRADIENTS BASED ON HALF-INVERSE JACOBIANS
text: 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 cannot 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.

heading: 2.1 AN ILL-CONDITIONED TOY EXAMPLE
text: 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?

heading: 2.2 UPDATES BASED ON HALF-INVERSE JACOBIANS
text: We propose a novel method for optimizing neural networks with physics objectives. Since pure physics or neural network optimization can be thought of as special cases of the joint optimization, we analogously look for a potential method in the continuum of optimization methods between gradient descent and Gauss-Newton. We consider both of them to be the most elementary algorithms representing network and physics optimizers, respectively. The following equation describes an update that lies in between the two.
∆θ(η, β, κ) = −η · ( ∂y
∂θ
)<β,κ> · ( ∂L
∂y
)> (4)
Here, the exponent < β, κ > of the Jacobian denotes the following procedure defined with the aid of the singular value decomposition J = UΛV > as:
J<β,κ> := max{diag(Λ)}β · V ΛκU>, (5)
When β = 0 and κ = 1, equation 4 reduces to the well-known form of gradient descent. Likewise, the case β = 0 and κ = −1 yields Gauss-Newton since the result of the Jacobian exponentiation then gives the pseudoinverse of the Jacobian. Unlike other possible interpolations between gradient descent and Gauss-Newton, exponentiation by κ as in equation 4 significantly affects the scales inherent in the Jacobian. This is highly important to appropriately influence physics and neural network scales.
To determine κ, we recall our goal to perform update steps which are normalized in both θ- and y-space. However, since any update ∆θ and its corresponding effect on the network output ∆y are connected by the inherent scales encoded in the Jacobian, no single κ exists that normalizes both at the same time. Instead, we distribute the burden equally between network and physics by choosing κ = −1/2. From a geometric viewpoint, the resulting update can be regarded as a steepest descent step measured by the norm given in table 1. In addition to the exponentiation, the scalar prefactor in equation 5 enables us to obtain a weaker form of scale invariance: Setting β = −1 − κ yields algorithms that rescale the largest singular value to 1, which ensures that the resulting updates cannot produce arbitrarily large updates in y-space.
For a batch of size b, we define an update with a learning rate η by stacking network-solver Jacobians ∂yi ∂θ ∣∣ xi and loss gradients ∂L∂yi ∣∣ xi,ŷi
of different data points (xi, ŷi) to form the following update step for our method:
∆θHIG = −η ·  ∂y1 ∂θ ∣∣ x1 ∂y2 ∂θ ∣∣ x2
... ∂yb ∂θ ∣∣ xb
 <−1/2,−1/2> ·  ∂L ∂y1 ∣∣> x1,ŷ1 ∂L ∂y2 ∣∣> x2,ŷ2 ...
∂L ∂yb ∣∣> xb,ŷb  (6) Besides batch size b and learning rate η, we specify a truncation parameter τ as an additional hyperparameter enabling us to suppress numerical noise during the half-inversion process in equation 5. As with the computation of the pseudoinverse via SVD, we set the result of the − 12 -exponentiation of every singular value less than τ to 0.
The use of a half inversion – instead of a full inversion – helps to prevent exploding updates of network parameters while still guaranteeing substantial progress in directions of low curvature. With the procedure outlined above, we arrived at a balanced method that combines the advantages of optimization methods from deep learning and physics. As our method uses gradients based on halfinverse Jacobians, we refer to them in short as half-inverse gradients (HIGs) in the following.
Half-inverse Gradients in the Toy Example. With the definition of HIGs from equation 6, we optimize the toy example introduced in section 2.1. The results in figure 1 show that for γ = 1, HIGs minimize the objective as well as Adam and Gauss-Newton’s method. More interestingly, HIGs achieve a better result than the other two methods for γ = 0.01. On the one hand, the physics trajectory (figure 1b, right) highlights that HIGs can process information along the smallscale component of y2 well and successfully progress along this direction. On the other hand, by checking neuron saturation (figure 1b, middle), we see that HIGs – in contrast to Gauss Newton – avoid oversaturating neurons.

heading: 2.3 PRACTICAL CONSIDERATIONS
text: Computational Cost. A HIG update step consists of constructing the stacked Jacobian and computing the half-inversion. The first step is handled via a vectorized Jacobian computation, and therefore induces a runtime cost comparable to regular backpropagation at the expense of higher memory requirements. The latter is not too problematic since our method achieves its optimal performance at lower batch sizes than the established gradient methods. In situations where the length of the HIG step is dominated by the half-inversion, memory requirements can be further reduced by parallelizing the Jacobian computation only partially. Next, the half-inversion is performed, at the heart of which lies a divide and conquer algorithm for the singular value decomposition (Trefethen & Bau, 1997). Hence, the length of a HIG step scales as O(|θ| · (b · |y|)2), i.e. it linear in the number of network parameters |θ|, and quadratic in the batch size b and the dimension of the physical state |y|. Concrete numbers for memory requirements and duration of a HIG step are listed in the appendix.
Hyperparameters. Our method requires choosing several hyperparameters. First, we need a suitable choice of the learning rate. The normalizing effects of HIGs allow for larger learning rates than commonly used gradient descent variants. We are able to use η = 1 for many of our experiments. Second, the batch size b affects the number of data points included in the half-inversion process. It should be noted that the way the feedback of individual data points is processed is fundamentally different from the standard gradient optimizers: Instead of the averaging procedure of individual gradients of a mini batch, our approach constructs an update that is optimized for the complete batch. Consequently, the quality of updates increases with higher batch size. However, overly large batch sizes can cause the Jacobian to become increasingly ill-conditioned and destabilize the learning progress (Trefethen & Bau, 1997). In the appendix C, we discuss the remaining parameters τ , β, and κ with several ablation experiments to illustrate their effects in detail.

heading: 3 EXPERIMENTS
text: We evaluate our method on three physical systems and tasks: controlling a nonlinear oscillator system, the Poisson problem, and a dipole problem from quantum mechanics. Details of the numerical setups are given in the appendix, together with results for a broad range of hyperparameter choices. For a fair comparison, we show results with the best set of hyperparameters for each of the methods below and plot the loss against wall-clock time. All learning curves are recorded on a previously unseen data set.

heading: 3.1 CONTROL OF NONLINEAR OSCILLATORS
text: First, we consider a control task for a system of coupled oscillators with a nonlinear interaction term. This system is of practical importance in many areas of physics, such as solid state physics (Ibach & Lüth, 2003). Its equations of motions are governed by the Hamiltonian
H(xi, pi, t) = ∑ i ( x2i 2 + p2i 2 + α · (xi − xi+1)4 + u(t) · xi · ci ) , (7)
where xi and pi denote the Hamiltonian conjugate variables of oscillator i, α the interaction strength. The constant vector c specifies how the control is applied. u(t) denotes the scalar-valued control signal to be learned by the neural network. For the integration of the 96 time steps, we use a 4th order Runge-Kutta scheme. The control task is to learn the signal u(t) that transforms a given initial state into a given target state. We use a dense neural network with three hidden layers totalling 2956 trainable parameters and ReLU activations. The Mean-Squared-Error loss is used to quantify differences between predicted and target state. A visualization of this control task is shown in figure 2a.
Optimizer comparison. The goal of our first experiments is to give a broad comparison of the proposed HIGs with commonly used optimizers. This includes stochastic gradient descent (SGD), Adagrad (Duchi et al., 2011), Adadelta (Zeiler, 2012), RMSprop (Hinton et al., 2012), Adam (Kingma & Ba, 2015), and Gauss-Newton (GN) applied to mini batches. The results are shown in figure 2b where all curves show the best runs for each optimizer with suitable hyperparameters independently selected, as explained in the appendix. We find that the state-of-the-art optimizers stagnate early, with Adam achieving the best result oscillating around a final loss value of 10−4. In comparison, our method and GN converge faster, exceeding Adam’s accuracy after about three minutes. While GN exhibits stability problems, the best stable run from our hyperparameter search reaches a loss value of 10−6. HIGs, on the other hand, yield the best result with a loss value of 10−7. These results clearly show the potential of our method to process different scales of the physics solver more accurately and robustly. They also make clear that the poor result of the widely-used network optimizers cannot be attributed to simple numerical issues as HIG converges to better levels of accuracy with an otherwise identical setup.
Role of the batch size. We conduct multiple experiments using different values for the batch size b as a central parameter of our method. The results are shown in figure 2c. We observe that for
Adam, all runs converge about equally quickly while HIGs, and GN, show notable improvements from larger batch sizes. This illustrates an important difference between Adam and HIG: Adam uses an average of point-wise gradients in the mini batch. As soon as the batch size is large enough, this average becomes representative of the expectation. Further increasing the batch size has little influence on the updates. Our method, on the other hand, includes the individual data point gradients without averaging. As shown in equation 6, we construct updates that are optimized for the whole batch by solving a linear system. This gives our method the ability to hit target states very accurately with increasing batch size. To provide further insights into the workings of HIGs, we focus on detailed comparisons with Adam as the most widely used gradient descent variant.

heading: 3.2 POISSON PROBLEM
text: Next we consider Poisson’s equation to illustrate advantages and current limitations of HIGs. Poisson problems play an important role in electrostatics, Newtonian gravity, and fluid dynamics (Ames, 2014). For a source distribution ρ(x), the goal is to find the corresponding potential field φ(x) fulfilling the following differential equation:
∆φ = ρ (8)
Classically, Poisson problems are solved by solving the corresponding system of linear equations on the chosen grid resolution. Instead, we train a dense neural network with three hidden layers and 41408 trainable parameters to solve the Poisson problem for a given right hand side ρ. We consider a two-dimensional system with a spatial discretization of 8×8 degrees of freedom. An example distribution and solution for the potential field are shown in figure 3a.
Convergence and Runtime. Figure 3b shows the learning curves when training the network with Adam and HIGs for different learning rates. As we consider a two-dimensional system, this optimization task is challenging for both methods and requires longer training runs. We find that both Adam and HIGs are able to minimize the loss by up to three orders of magnitude. The performance of Adam varies, and its two runs with larger η quickly slow down. In terms of absolute convergence per time, the Adam curve with the smallest η (in light blue) shows advantages in this scenario. However, choosing a log-scale for the time axis reveals that both methods have not fully converged. In particular, while the Adam curve begins to flatten at the end, the slope of the HIG curve remains constant and decreases with a steeper slope than Adam. The performance of Adam in this example can be explained by two reasons. First, the time to compute a single Adam update is much smaller than for HIGs, which requires the SVD solve from equation 5. While these could potentially be sped up with appropriate methods (Foster et al., 2011; Allen-Zhu & Li, 2016), the absolute convergence per iteration, shown in the appendix in figure 7, shows how much each HIG update improves over Adam. Second, compared to the other examples, the Poisson problem is relatively simple, requiring only a single matrix inversion. This represents a level of difficulty which Adam is still able to handle relatively well.
HIGs with Adam Pretraining. To further investigate the potential of HIGs, we repeat the training, this time using a state of the Adam model from figure 3b for the network initialization. While Adam progresses slowly, the HIG training is able to quickly improve the state of the neural network, resulting in a significant drop of the loss values, followed by a faster descent than Adam. Interestingly, this experiment indicates that the HIG updates are able to improve aspects of the solution which Adam is agnostic to. Despite outlining the potential gains from faster SVD calculations, this example also highlights the quality of the HIG updates for simpler, field-based PDEs.

heading: 3.3 QUANTUM DIPOLE
text: As a final example, we target the quantum dipole problem, which is a standard control task formulated on the Schrödinger equation and is highly relevant in quantum physics (Von Neumann, 2018). Given an initial and a target state, we train a neural network to compute the temporal transition function u(t) in an infinite-well potential V according the evolution equation of the physical state Ψ:
i∂tΨ = ( −∆ + V + u(t) · x̂ ) Ψ (9)
We employ a modified Crank-Nicolson scheme (Winckel et al., 2009) for the discretization of spatial and temporal derivatives. Thus, each training iteration consists of multiple implicit time integration steps – 384 in our setup – for the forward as well as the backward pass of each mini-batch. The control task consists of inferring a signal that converts the ground state to a given randomized linear combination of the first and the second excited state. We use a dense neural network with three hidden layers, 9484 trainable parameters and tanh activations. As loss function, we use the inner product of quantum mechanics L(Ψa,Ψb) = 1− |〈Ψa,Ψb〉|2. A visualization of this control task is shown in figure 4a.
Speed and Accuracy. We observe that HIGs minimize the loss faster and reach a better final level of accuracy than Adam (figure 4b). While the Adam run with the largest learning rate drops faster initially, its final performance is worse than all other runs. In this example, the difference between the final loss values is not as large as for the previous experiments. This is due to the numerical accuracy achievable by a pure physics optimization, which for our choice of parameters is around 10−6. Hence, we can not expect to improve upon this lower bound for derived learning problems. Our results indicate that the partial inversion of the Jacobian successfully leads to the observed improvements in convergence speed and accuracy.
Low and High Energy Components. The quantum control problem also serves to highlight the weakness of gradient-based optimizers in appropriately processing different scales of the solutions. In the initial training stage, the Adam curves stagnate at a loss value of 0.5. This is most pronounced for η = 10−4 in dark blue. To explain this effect, we recall that our learning objective targets transitions to combinations of the 1st and 2nd excited quantum states, and both states appear on average with equal weight in the training data. The transition to the energetically higher state is more difficult and connected to smaller scales in the physics solver causing Adam to fit the lower-energetic
component first. In contrast, our method is constructed to consider all scales via the partial inversion of the Jacobian and can therefore process the small scales more efficiently. As a consequence, it does not stagnate at 0.5. We support this explanation by explicitly plotting separate loss curves in figure 4c quantifying how well the low and high energy component of the target state was learned. Not only does Adam prefer to minimize the low-energy loss, it also increases the same loss again before it is able to minimize the high-energy loss. In contrast, we observe that the HIG minimizes both losses uniformly. This is another indication that the theory outlined above carries over to an even processing of different gradient magnitudes in joint physics and neural network objectives.

heading: 4 RELATED WORK
text: Optimization algorithms. Optimization on continuous spaces is a huge field that offers a vast range of techniques (Ye et al., 2019). Famous examples are gradient descent (Curry, 1944), GaussNewton’s method (Gill & Murray, 1978), Conjugate Gradient (Hestenes et al., 1952), or the limitedmemory BFGS algorithm (Liu & Nocedal, 1989). In deep learning, the preferred methods instead rely on first order information in the form of the gradient, such as SGD (Bottou, 2010) and RMSProp (Hinton et al., 2012). Several methods approximate the diagonal of the Hessian to improve scaling behavior, such as Adagrad (Duchi et al., 2011), Adadelta (Zeiler, 2012), and most prominently, Adam (Kingma & Ba, 2015). However, due to neglecting all inter-dependencies of the parameters, these methods are limited in their capabilities to handle uneven learning objectives. Despite the computational cost, higher-order methods have also been studied in deep learning (Pascanu & Bengio, 2013) . Practical methods have been suggested by using a Kroenecker-factorization of the Fisher matrix (Martens & Grosse, 2015) via iterative linear solvers in (Martens, 2010), or by recursive approximations of the Hessian (Botev et al., 2017). To the best of our knowledge, the only other technique specifically targeting optimization of neural networks with physics objectives are physical gradients (PGs) (Holl et al., 2021). Despite targeting inversions, they fundamentally differ from our approach as PGs are based on inverse physics solvers, while we address the problem by treating network and solver as an entity and half-inverting its Jacobian. Thus, we work on the level of linear approximations while PGs are able to use higher-order information provided that an higher-order inverse solver exists. Additionally, PGs compute their update by averaging gradients over different data points, in line with typical gradient-based neural network optimizers. HIGs instead process the feedback of different data points via collective inversion.
Incorporating physics. Many works involve differentiable formulations of physical models, e.g., for robotics (Toussaint et al., 2018), to enable deep architectures (Chen et al., 2018), as a means for scene understanding (Battaglia et al., 2013; Santoro et al., 2017), or the control of rigid body environments de Avila Belbute-Peres et al. (2018). Additional works have shown the advantages of physical loss formulations Greydanus et al. (2019); Cranmer et al. (2020). Differentiable simulation methods were proposed for a variety of phenomena, e.g. for fluids (Schenck & Fox, 2018), PDE discretizations (Bar-Sinai et al., 2019), molecular dynamics (Wang et al., 2020), reducing numerical errors (Um et al., 2020), and cloth (Liang et al., 2019; Rasheed et al., 2020). It is worth noting that none of these works question the use standard deep learning optimizers, such as Adam. In addition, by now a variety of specialized software frameworks are available to realize efficient implementations (Hu et al., 2020; Schoenholz & Cubuk, 2019; Holl et al., 2020).

heading: 5 DISCUSSION AND OUTLOOK
text: We investigated the training of neural networks for physical simulations. As including a physics solver greatly affects the gradient flow, we questioned the current practice of using the standard gradient-based network optimizers for training. By considering ways to smoothly transition from gradient descent to the Gauss-Newton algorithm, we identified a procedure to better address the joint optimization problem. Compared to existing optimizers, our method can achieve a faster and more accurate minimization of the learning objectives.
We believe that our work can provide a starting point for further research into improved learning methods for physical problems. In the context of HIGs, highly interesting avenues for future work are efficient methods for the inversion of the Jacobian matrix, or including unstructured solvers. The combination of finite elements (Zienkiewicz et al., 1977) with graph-neural networks (Battaglia et al., 2013; Ummenhofer et al., 2020; Pfaff et al., 2020) provides a sound foundation that is especially interesting for future applications of HIGs.
REPRODUCIBILITY STATEMENT
To ensure reproducibility of our work we will publish all source code that is necessary to reproduce the results of our paper upon acceptance. This includes data generation, neural network training and evaluation. Hyperparameters are also provided in the appendix for completeness.

heading: A GAUSS-NEWTON’S METHOD IN THE CONTEXT OF PHYSICAL DEEP LEARNING
text: Our work considers optimization algorithms for functions of the form f(x;θ) = y with θ,∆θ ∈ Rt, denoting weight vector and weight update vector, respectively, while x ∈ Rn and y ∈ Rm denote input and output. The learning process solves the minimization problem argminθL(f(x;θ), ŷ) via a sequence θk+1 = θk + η∆θ. Here, ŷ are the reference solutions, and we target losses of the form L(x, ŷ;θ) = ∑ i l ( f(xi;θ), ŷi ) with i being an index for multiple data points (i.e., observa-
tions). l denotes the l2-loss ∑ j ||xj − ŷj ||22 with j referencing the entries of a mini batch of size b.
Update Step. With this notation, the update step of the Gauss-Newton algorithm (Adby, 2013) for η = 1 is given by:
∆θGN = −
(( ∂y
∂θ
)T · ( ∂y
∂θ
))−1 · ( ∂y
∂θ
)T · ( ∂L
∂y
)> (10)
The size of the Jacobian matrix is given by y-space and parameter dimensions. For a full-rank Jacobian corresponding to non-constrained optimization, the Gauss-Newton update is equivalent to:
∆θGN = − ( ∂y
∂θ
)−1 · ( ∂L
∂y
)> (11)
This shortened form of the update step is given in equation 3, and is the basis for our discussion in the main text.
Geometric Interpretation. These updates can be regarded as a gradient descent in the y-variables up to first order in the update step as can be seen with the following equation:
∆y = − ( ∂y
∂θ
) ·∆θGN + o(∆θGN) = − ( ∂L
∂y
)> + o(∆θGN) (12)
In a constrained setting, we can reparametrize the coordinates to obtain an unconstrained optimization problem on the accessible manifold. We can then repeat the argument above. Hence, even in the constrained setting the Gauss-Newton algorithm can be seen as gradient descent on the y-manifold.
In figure 1 of the main paper this behavior is visible in both the well- and ill-conditioned case in terms of the straight line that the GN curve takes towards the target even in the presence of a neural network. The first few, seemingly erratic, movements in y-space which are observable for both GN and HIG stem from the inversion operations.
The behavior of GN to perform steepest descent in y-space stands in contrast to gradient descent methods, which instead perform steepest descent on the θ-manifold. In combination, this geometric viewpoint, as summarized in table 1 above motivates our construction of the HIGs
Stability of Inversions in the Context of Physical Deep Learning. In the following, we illustrate how the full inversion of GN can lead to instabilities at training time. Interestingly, physical solvers are not the only cause of small singular values in the Jacobian. They can also occur when applying equation 11 to a mini batch to train a neural network and are not caused by numerical issues. Consider the simple case of two data points (x1, ŷ1) and (x2, ŷ2) and a one-dimensional output. Let f be the neural network and J the Jacobian, which is in this case the gradient of the network output. Then equation 11 yields:
( Jf (x1) Jf (x2) ) ·∆θGN = ( f(x1)− ŷ1 f(x2)− ŷ2 ) (13)
Next, we linearly approximate the second row by using the HessianH by assuming the function to be learned is f̂ , i.e. f̂(x1) = y1 and f̂(x2) = y2. Neglecting terms beyond the linear approximation, we receive: ( Jf (x1)
Jf (x1) +Hf (x1) · (x2 − x1)
) ·∆θGN = ( f(x1)− y1
f(x1)− y1 + (Jf (x1)− Jf̂ (x1)) · (x2 − x1) ) (14)
Considering the case of two nearby data points, i.e. x2 −x1 being small, the two row vectors in the stacked Jacobian on the left-hand side are similar, i.e. the angle between them is small. This leads to a small singular value of the stacked Jacobian. In the limit of x2 = x1 both row vectors are linearly dependant and hence, one singular value becomes zero.
Moreover, even if x2 is not close to x1, small singular values can occur if the batch size increases: for a growing number of row vectors it becomes more and more likely that the Jacobian contains similar or linearly dependent vectors.
After inversion, a small singular value becomes large. This leads to a large update ∆θGN when the right-hand side of equation 14 overlaps with the corresponding singular vector. This can easily happen if the linear approximation of the right-hand side is poor, for instance when f̂ is a solution to an inverse physics problem. Then f̂ can have multiple modes and can, even within a mode, exhibit highly sensitive or even singular behavior. In turn, applying large updates to the network weights naturally can lead to the oversaturation of neurons, as illustrated above, and diverging training runs in general.
As illustrated in the main paper, these inherent problems of GN are alleviated by the partial inversion of the HIG. It yields a fundamentally different order of scaling via its square-root inversion, which likewise does not guarantee that small singular values lead to overshoots (hence the truncation), but in general strongly stabilizes the training process.

heading: B EXPERIMENTAL DETAILS
text: In the following, we provide details of the physical simulations used for our experiments in the main paper. For the different methods, we use the following abbreviations: half-inverse gradients (HIG), Gauss-Newton’s method (GN), and stochastic gradient descent (GD). Learning rates are denoted by η, batch sizes by b, and truncation parameters for HIG and GN by τ . All loss results are given for the average loss over a test set with samples distinct from the training data set.
For each method, we run a hyperparameter search for every experiment, varying the learning rate by several orders of magnitude, and the batch size in factors of two. Unless noted otherwise, the best runs in terms of final test loss were selected and shown in the main text. The following sections contain several examples from the hyperparameter search to illustrate how the different methods react to the changed settings.
Runtime Measurements Runtimes for the non-linear chain and quantum dipole were measured on a machine with Intel Xeon 6240 CPUs and NVIDIA GeForce RTX 2080 Ti GPUs. The Poisson experiments used an Intel Xeon W-2235 CPU with NVIDIA Quadro RTX 8000 GPU.We experimentally verified that these platforms yield an on-par performance for our implementation. As deep learning API we used TensorFlow version 2.5. If not stated otherwise, each experiment retained the default settings.
All runtime graphs in the main paper and appendix contain wall-clock measurements that include all steps of a learning run, such as initialization, in addition to the evaluation time of each epoch. However, the evaluations of the test sets to determine the performance in terms of loss are not included. As optimizers such as Adam typically performs a larger number of update steps including these evaluations would have put these optimizers at an unnecessary disadvantage.
B.1 TOY EXAMPLE (SECTION 2.1)
For the toy example, the target function is given by f̂(x) = (sin(6x), cos(9x)). We used a dense neural network consisting of one hidden layer with 7 neurons and tanh activation, and an output layer with 2 neurons and linear activation. For training, we use 1024 data points uniformly sampled from the [−1, 1] interval, and a batch size of 256. The test set consists of 1024 new samples. For the optimizers, the following hyperparameters were used for both the well-conditioned loss and the ill-conditioned loss: Adam η = 0.3; has no learning rate (equivalent to η = 1), τ = 10−4; HIG η = 1.0, τ = 10−6.
B.2 CONTROL OF NONLINEAR OSCILLATORS (SECTION 3.1)
The Hamiltonian function given in equation 7 leads to the following equations of motions:
ẍi = −xi + 4α(xi − xi−1)3 − 4α(xi − xi+1)3 − u(t) · ci (15)
The simulations of the nonlinear oscillators were performed for two mass points and a time interval of 12 units with a time step ∆t = 0.125. This results in 96 time steps via 4th order Runge-Kutta per learning iteration. We generated 4096 data points for a control vector c = (0.0, 3.0), and an interaction strength α = 1.0 with randomized conjugate variables x and p. The test set consists of 4096 new data points. For the neural network, we set up a fully-connected network with ReLU activations passing inputs through three hidden layers with 20 neurons in each layer before being mapped to a 96 output layer with linear activation.
For the comparison with other optimizers (figure 2b) we performed a broad hyperparameter search for each method, as outlined above, to determine suitable settings. The parameters for Adagrad (Duchi et al., 2011), Adadelta (Zeiler, 2012), Adam (Kingma & Ba, 2015), RMSprop (Hinton et al., 2012), Gauss-Newton (Gill & Murray, 1978), HIGs, and stochastic gradient descent (Curry, 1944) are summarized in table 2. For figure 2c the following hyperparameters were used: η = 3 · 10−4 for Adam, and η = 1.0, τ = 10−6 for HIG.
Further Experiments. Figure 5 and figure 6 contain additional runs with different hyperparameters for the method comparison of figure 2b in the main paper. The graphs illustrate that all five method do not change their behavior significantly for the different batch sizes in each plot, but become noticeably unstable for larger learning rates η (plots on the right sides of each section).
Details on the memory footprint and update durations can be found in table 3. Since our simulations were not limited by memory, we used an implementation for the Jacobian computation of HIGs, which scales quadratically in the batch size. Should this become a bottleneck, this scaling could potentially be made linear by exploiting that the Jacobian of the physical solver for multiple data points is blockdiagonal.
B.3 POISSON PROBLEM (SECTION 3.2)
We discretize Poisson’s equation on a regular grid for a two-dimensional domain Ω = [0, 8]× [0, 8] with a grid spacing of ∆x = 1. Dirichlet boundary conditions of φ = 0 are imposed on all four sides of Ω. The Laplace operator is discretized with a finite difference stencil (Ames, 2014).
For the neural network, we set up a fully-connected network with tanh activation functions. The 8x8 inputs pass through three hidden layers with 64, 256 and 64 neurons, respectively, before being mapped to 8x8 in the output layer. For training, source distributions ρ are sampled from random frequencies in Fourier space, and transformed to real space via the inverse Fourier transform. The mean value is normalized to zero. We sample data on-the-fly, resulting in an effectively infinite data set. This makes a separate test set redundant as all training data is previously unseen.
Table 4: Poisson problem: memory requirements, update duration and duration of the derivative computation for Adam and HIG
Optimizer Adam HIG Batch size 64 64 Memory (MB) 1.3 3560 Update duration (sec) 0.011 13.8
Derivative duration (sec) 0.010 0.0035
101 103 105 Epochs
10 2
100
102
Lo ss
Adam = 1e-03 Adam = 3e-04 Adam = 1e-04 HIG = 0.02
a)
105 1.1 105 1.2 105 1.3 105 Wall clock time
6 10 3
10 2
2 10 2
3 10 2
4 10 2 5 10 2
Lo ss
Adam pretrained HIG Adam after HIG b)
Figure 7: Poisson problem: a) Loss curves for Adam and HIG per epoch for different learning rates, b) Loss curves of Adam (η =1e-04), of HIG (η = 0.02) pretrained with Adam, and of Adam (η =1e-04) pretrained with the HIGs.
Further Experiments. Figure 7a shows Adam and HIG runs from figure 3b over epochs. The HIG runs converge faster per iteration, which indicates that HIGs perform qualitatively better updates.
Additionally, we use the pretrained HIG run from figure 3c as a starting point for further Adam training. The results are shown in 7b. We observe that the network quickly looses the progress the HIGs have made, and continues with a loss value similar to the orginal Adam run. This again supports our intuition that Adam, in contrast to HIGs, cannot harness the full potential of the physics solver.
Details on the memory footprint and update durations can be found in table 4
B.4 QUANTUM DIPOLE (SECTION 3.3)
For the quantum dipole problem, we discretize the Schrödinger equation on a spatial domain Ω = [0, 2] with a spacing of ∆x = 0.133 resulting in 16 discretization points. In time, we simulate up to a time of 19.2 with a time step of ∆t = 0.05, which yields 384 time steps. Spatial and temporal discretization use a modified Crank-Nicolson scheme (Winckel et al., 2009) which is tailored to quantum simulations. The training data set consists of 1024 randomized superpositions of the first and second excited state, while the test set contains a new set of 1024 randomized superpositions. For the neural network, we set up a fully-connected network with tanh activations passing the inputs through three hidden layers with 20 neurons in each layer before being mapped to a 384 output layer with linear activation. Overall, the network contains 9484 trainable parameters.
Experimental details. For the training runs in figure 4b, Adam used b = 16, while for HIG b = 16, and τ = 10−5 were used. For the training runs in figure 4c, Adam used b = 16, η = 0.0001, while HIGs used b = 16, τ = 10−5, and η = 0.5. Details on the memory footprint and update durations can be found in table 5
Figure 8 and figure 9 show the performance of both methods for a broader range of τ settings for HIGs, and η for Adam. For Adam, a trade-off between slow convergence and oscillating updates exists. The HIGs yield high accuracy in training across a wide range of values for τ , ranging from 10−5 to 10−3. This supports the argumentation in the main text that the truncation is not overly critical for HIGs. As long as numerical noise is suppressed with τ > 10−6, and the actual information about scaling of network parameters and physical variables is not cut off. The latter case is visible for an overly large τ = 0.01 in the last graph on the right.
Note that many graphs in figure 9 contain a small plateau at the start of each training run. These regions with relatively small progress per wall clock time are caused by the initialization overhead of the underlying deep learning framework (TensorFlow in our case). As all graphs measure wall clock time, we include the initialization overhead of TensorFlow, which causes a noticeable slow down of the first iteration. Hence, the relatively slow convergence of the very first steps in figure 9 are not caused by conceptual issues with the HIGs themselves. Rather, they are a result of the software frameworks and could, e.g., be alleviated with a pre-compilation of the training graphs. In contrast, the initial convergence plateaus of Adam with smaller η in Figure 8 are of a fundamentally
different nature: they are caused by an inherent problem of non-inverting optimizers: their inability to appropriately handle the combination of large and small scale components in the physics of the quantum dipole setup (as outlined in section 3.3).
Loss Functions. While training is evaluated in terms of the regular inner product as loss function: L(Ψa,Ψb) = 1 − |〈Ψa,Ψb〉|2, we use the following modified losses to evaluate low- and highenergy states for figure 4c. Let Ψ1 be the first excited state, then we define the low-energy loss as:
L(Ψa,Ψb) = (|〈Ψa,Ψ1〉| − |〈Ψ1,Ψb〉|)2
Correspondingly, we define the high-energy loss with the second excited state Ψ2:
L(Ψa,Ψb) = (|〈Ψa,Ψ2〉| − |〈Ψ2,Ψb〉|)2

heading: C ABLATION STUDY
text: In this last section, we investigate how the HIG-hyperparameters affect the outcome. This includes ablation experiments with respect to κ, β, and τ defined in section 2.2. We use the nonlinear oscillator example as the basis for these comparisons.
Exponent of the diagonal singular value matrix κ: We test κ for various values between 1.0, stochastic gradient descent, and −1, Gauss-Newton. The results are shown in figure 10a. For positive values, curves stagnate early, while for negative κ, the final loss values are several orders of magnitude better. The HIG curve corresponding to β = −0.5 achieves the best result. This supports our argumentation that a strong dependence on this parameter exists, and that a choice of κ = −0.5 is indeed a good compromise for scale-correcting updates of reasonable size. The strong improvement as soon as κ becomes negative indicates that the collective inversion of the feedback of different data points of the mini-batch is an important ingredient in our method.
Scalar prefactor term β: We test β-values between 0, no scale correction, and−0.5, which fully normalizes the effect of the largest singular value. The results are shown in figure 10b. Compared to the other hyperparameters, we observe that β has only little influence on the outcome.
Truncation parameter τ : To understand the effect of this parameter, we consider the singular value decomposition (SVD) of the network-solver Jacobian, which is determined by the SVDs of the network Jacobian and the solver Jacobian. The singular values of a matrix product AB depend non-trivially on the singular values of the matrices A and B. In the simplest case, the singular values of the matrix product are received by multiplication of the individual singular values of both matrix factors. In the general case, this depends on how the singular vectors of A and B overlap with each other. However, it is likely that singular vectors with a small singular value of A or B overlap significantly with singular vectors with a small singular value of AB. For this reason, it is important not to truncate too much as this might remove the small-scale physics modes that we are ultimately trying to preserve in order to achieve accurate results. On the other hand, less truncation leads to large updates of network weights on a scale beyond the validation of the linear approximation by first-order derivatives. These uncontrolled network modifications can lead to over-saturated neurons and prevent further training progress.
From a practical point of view, we choose τ according to the accuracy of the pure physics optimization problem without a neural network. For the quantum dipole training, this value was set to 10−5. Trying to solve the pure physics optimization with far smaller values leads to a worse result or no convergence at all. The network training behaves in line with this: Figure 9 shows that the network does not learn to control the quantum system with τ -values far smaller than 10−5 . For the nonlinear oscillator system, the pure physics optimization is stable over a large range of τ -values with similarly good results. For the network training, we chose τ to be 10−6. We conducted further experiments for the network training with different τ from 10−5 to 10−10 presented in figure 11, which show that HIGs have a similar tolerance in τ . For a comparison, we also plotted Gauss-Newton curves for different τ . We observe that GN curves become more unstable for smaller truncation values τ and diverge in the case 10−9 and 10−10 while HIG curves achieve overall better loss values and start to converge in this parameter.

Review summaries:
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.

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.

The paper considers the training of neural networks for physical simulations. By distributing the burden equally between network and physics, the paper presents the half-inverse gradients (HIGs) method. Experiments show its advantage for achieving a faster and more accurate minimization.