Accelerated PDEs for Construction and Theoretical Analysis of an SGD Extension
Abstract
We introduce a recently developed framework ( PDE acceleration ), which is a variational approach to accelerated optimization with partial differential equations (PDE), in the context of optimization of deep networks.
We derive the PDE evolution equations for optimization of general loss functions using this variational approach.
We propose discretizations of these PDE based on numerical PDE discretizations, and establish a mapping between these discretizations and stochastic gradient descent (SGD).
We show that our framework can give rise to new PDEs that can be mapped to new optimization algorithms, and thus theoretical insights from the PDE domain can be used to analyze optimization algorithms.
We show an example by introducing a new PDE with diffusion that naturally arises from the viscosity solution, which translates to a novel extension of SGD.
We analytically analyze the stability and convergence using Von-Neumann analysis.
We apply the proposed extension to optimization of convolutional neural networks (CNNs).
We empirically validate the theory and evaluate our new extension on image classiﬁcation showing empirical improvement over SGD.
Introduction
Acceleration (momentum) and Nesterov acceleration are now standard in optimizing deep networks by stochastic gradient descent (SGD) and its variants, although the underlying principles of these methods are still under development.
Understanding the principles behind these methods could lead to new improved optimization methods.
To this end, [27] showed that Nesterov’s accelerated optimization [15] can be formulated as a discretization of an ordinary differential equation (ODE) [23] arising from a variational principle, i.e., the Principle of Least Action in classical physics.
Several new optimization schemes have since arisen from different discretizations of the same ODE [22].
An extension of this approach from the ODE framework into the PDE framework was done in [24, 3], which allowed for the accelerated optimization of multi-dimensional problems.
New optimization methods were formulated arising in particular from various discretizations of the underlying PDE [3].
In this work, we study SGD with momentum within the framework of Accelerated PDE and PDE numerical discretization.
We establish a mapping from SGD with momentum to Accelerated PDE.
This allows one to derive new extensions of SGD by translating new PDE naturally arising within the Accelerated PDE framework into numerical discretizations, which map back to extensions of
SGD.
We propose a novel extension of SGD in this way, and apply it to optimization of CNNs.
The derivation of the method from numerical discretization of Accelerated PDE allows us to establish stability and convergence results.
More speciﬁcally, our contributions are
We show how SGD with momentum can be related to discretizations of an Accelerated PDE. 2.
Within this Accelerated PDE framework, we introduce a diffusion term arising from the viscosity solution (2) of the PDE.
35th
We show how the discretization of the viscosity solution of this Accelerated PDE translates to a simple modiﬁcation of the usual SGD with momentum, leading to a new optimization scheme.
We exploit the discretizations of the Accelerated PDE to show that our new optimization scheme is stable/convergent provided that the original SGD with momentum scheme is stable/convergent.
1.1 Related Work
Many extensions of SGD (e.g. [7, 29, 10, 10, 2, 13]) have been proposed (see [4] for a survey) and there has been rigorous analysis of some of these methods (e.g. [21]).
A framework to derive methods could lead to further understanding and new optimization schemes.
Our goal is to take a small step in this direction by applying the Accelerated PDE framework [24] to study SGD with momentum under discretizations of PDE.
In this regard, Wibisono, Wilson, and Jordan [27] showed that all
Nesterov type accelerated descent methods [15, 18, 19, 16, 17] can be realized as discretizations of equations of motion (ODE) arising from an Action Integral of a generalized Lagrangian.
Following the Lagrangian formulation, [24, 3] developed an accelerated PDE framework and applied it to computer vision problems.
[24] introduced the use of diffusion motivated from Viscosity Theory [5], allowing the use of entropy discretizations to address shocks/fans in the resulting PDE.
The PDE acceleration framework was further developed in [3] in the context of calculus of variations problems deﬁned for functions on R n , including stability analysis for various discretization schemes.
We apply discretization methods considered there to the novel accelerated PDE that we consider in this paper.
In application to CNNs, our method induces smoothness into optimization by diffusing the velocity, which has an effect similar to gradient smoothing.
[20] considers smoothing the stochastic gradient by applying a Laplacian smoothing operator.
[11] smooths along certain 1D directions of convolutional weight tensors using Sobolev gradients.
While both methods [20, 11] are 1D, our diffusion can induce multi-dimensional smoothing with little implementation effort and computational cost.
Continuum PDE Acceleration Theory
We consider minimizing a loss function of the form U : U → R where U
= { u : Ω ⊂ R N → R } is a set of N -dimensional functions.
In subsequent sections, U will represent a continuum version of the space of weight tensors at a particular layer of a network.
The argument of u , denoted x , will be called the "spatial" argument.
We derive the continuum evolution equations for accelerated optimization, which are PDE, and discretize the PDE in the next section for numerical implementation.
The formulation below to derive the optimization equations of a loss function is derived from
Hamilton’s principle of least action in classical physics [1, 14].
Following [24], we start with an action integral, deﬁned on possible optimization paths u t : Ω → R of U , where t denotes the parameter (time) along the path, and its velocity (in the tangent space of U )
V t : Ω → R .
We suppress t from now on for simplicity of notation.
The action integral is deﬁned as
[Equation]
where the path and its velocity are related as ∂ t u = V , ˜ a, ˜ b are functions only of time that give rise to dissipative forces and will be speciﬁed, T ( V ) is the kinetic energy, and U is the loss function to be minimized, analogous to the potential.
For the purpose of this paper, the kinetic energy is deﬁned as
T ( V ) =
(cid:82) Ω 12 ρ ( x )
| V ( x ) | 2 dx , where ρ ( x ) is the mass density, which for this paper is chosen to be 1.
The stationary conditions are derived by computing the ﬁrst variation of A and result in the equations:
[Equation]
where ∂ t denotes the time derivative, ∇ is the spatial gradient, ∇ 2 is the spatial Laplacian, a = ∂ t ˜ a/ ˜ a is chosen to be a positive constant, ˜ b is chosen to be 1, and b > 0 .
Note that the diffusion term (Laplacian term) results from the PDE viscosity solution to ensure existence of a solution [5].
The above system of equations is a forward (in time) evolution that starts with initial conditions on u and V .
We may write the above ﬁrst order PDE system as a second order in time PDE, which results in
[Equation]
where a, b > 0 are considered hyperparameters of the optimization, with a interpretated as a damping coefﬁcient that leads to convergence, and b interpretated as a diffusive coefﬁcient.
The above equation resembles a wave equation, except for the fact that the spatial Laplacian term is on the velocity rather than u ; this induces regularity on the velocity or the optimization in contrast to regularity within the loss function, which would have led to the standard wave equation.
Discretization of Accelerated PDE, Relation to SGD, and Analysis 3.1 Semi-Implicit Euler Discretization of the Accelerated PDE
We discretize the ﬁrst two time derivatives in (3) using central differences, and the last time derivative with a backward difference.
The last time derivative uses a backward difference as a central difference with the spatial derivative approximations would require an implicit scheme, requiring solution of a linear system, which we would like to avoid for simplicity.
To obtain systems which more closely resemble the classic two-part Nesterov recursion, we propose a semi-implicit Euler style discretization as considered in [3].
To do this, we replace the explicit discretization ∇ U n of the gradient with a "predicted estimate”, (cid:92) ∇ U n +1 , of ∇ U n +1 evaluated at a projected location of u n , denoted w n , as follows:
[Equation]
(4) where ∆ u n = u n − u n − 1 , and ∆ t is the time-step of the time discretizations.
The spatial Laplacian above is discretized using its standard discretization as
[Equation]
where x is the (spatial) argument of the tensor, e k is the unit vector in the k th direction, and ∆ x is the spatial discretization amount (assumed to be uniform in all N directions).
3.2 Relation to SGD With Nesterov Momentum
We establish a link between the Semi-Implicit Euler discretization (4) of the Accelerated PDE and
SGD, which shows how quantities in the PDE discretization (e.g., step size, damping) translate to the quantities in SGD (e.g., learning rate, momentum), and vice-versa.
This shows how the discretization of the Accelerated PDE (3), which gives rise to a novel extension of SGD, can be implemented with conventional deep learning frameworks.
Theorem 1.
The semi-implicit scheme (4) can be written as the following modiﬁcation of SGD with
Nesterov momentum 1 (SGD-Nesterov):
[Equation]
where γ, η, β are functions of the time step ∆ t and damping a .
γ is the momentum, and η is the learning rate.
The mapping of SGD-Nesterov to our Semi-Implicit Euler Scheme (4) is done by setting θ n =
[Equation]
n ,
[Equation]
and J = U .
Notice in the case that the diffusion coefﬁcient is set to b = 0 ( β = 0 ), (6) , (7) reduce to standard SGD with Nesterov momentum.
The mapping from SGD-Nesterov to the Semi-Implicit Euler Scheme (4) as
[Equation]
3.3 Stability and Convergence Analysis
Using the mapping between the Accelerated PDE and SGD, we are able to use methods of analysis in numerical PDE to analyze properties of our new SGD extension, in particular stability and convergence of the new SGD extension.
We compute the Courant–Friedrichs–Lewy (CFL) condition of the schemes, which are necessary conditions for convergence, to analyze the stability and convergence of the discretizations proposed in the previous section by applying Von-Neuman analysis [26], which involves analysis in the Fourier domain.
By doing so, we can show the following theorem.
Theorem 2 (CFL Conditions for Accelerated Discretizations) .
The Semi-Implicit Euler (4) Scheme with addition of diffusion term will maintain stability and converge provided that the time step ∆ t , the spatial step ∆ x , and the diffusion coefﬁcient b satisfy the relation:
[Equation]
where N is the dimension of the functions in U .
The maximum diffusion b ( β ) for a given time step that yields a stable/convergent scheme is
[Equation]
where we have chosen ∆ x = 1 .
We experiment with varying amounts of diffusion relative to β max .
Experiments 4.1 Empirical Validation of Stability Analysis
We empirically validate the condition in Theorem 2.
We use the Cifar10 dataset and ResNet56 [8] with diffusion added to the BN weight tensor velocity (other tensors and architectures show similar behavior) and with the rest of the network trained on standard SGD with Nesterov.
As shown in Figure 1, the training is stable with the maximum stable diffusion, β max .
However, if we raise the diffusion amount to over β max , the training becomes unstable and fails to converge.
The experiment shows the accuracy of the bound in Theorem 2 and the validity of PDE analysis for analyzing a deep learning optimization algorithm.
4.2 Empirical Analysis of Diffusion Effect
We analyze the behavior of the diffusion term on various choices of velocity tensors in optimization of a CNN.
We study diffusion over various dimensions of 4D conv tensors, i.e., spatial dimensions of kernels, input/output directions, batch norm tensors.
To isolate the effect of diffusion from empirically chosen hyperparameters (e.g., learning rate schedules) that are optimized for standard SGD, we study the case of ﬁxed learning rates, which is also more common in PDE-driven optimization methods.
Test accuracy (over 10 trials) against SGD-Nesterov for Cifar-10 on ResNet-56 are shown in Table(diffusion amounts are in terms of β max and amounts above β max are achieved by lowering the step size) and Table 2.
For diffusion coefﬁcients which exceed the maximum stable diffusion for a speciﬁc time step, the time must be reduced to ensure the stability of the system.
For 2 and 4 times of the maximum stable diffusion, we divide the time step by 2 and 4 respectively and recalculate the corresponding learning rate and momentum with equation(4).
Other hyperparameters (e.g., momentum) are set to current practice.
All diffusion amounts improve performance for batch norm layers over all learning rates and conv directions except at small learning rates.
We are currently studying how to exploit the larger boosts for large learning rates to construct schedules for the learning rate and diffusion.
In the meantime, with current schedules and training techniques optimized for standard optimizers, we do see statistically improved performance for diffusion to batch norm layers (see appendix), though at a diminished level.
[Table] [Table]
A Proof for theorem
Proof.
First we rewrite the semi-implicit scheme (4) in a Nesterov form.
Replacing u n in the ﬁrst equation with second equation we obtain:
[Equation]
Substituting w n in the second equation into the ﬁrst equation, we obtain:
[Equation]
We now start to map the semi-implicit scheme to SGD with Nesterov momentum using the standard form in PyTorch.
This is in following form:
[Equation]
where θ is the parameter and v is the momentum buffer, η for learning rate and γ for momentum.
Compared to the equation (12) and (13) [Equation] we formulate the Nesterov momentum with diffusion as
[Equation]
Note we can also obtain the relationship between the variables in the two different schemes:
[Equation]
B Proof for theorem
Proof.
We base our analysis on the homogeneous part of the PDE, and neglect the inhomegenous part of the PDE, i.e., the gradient component ∇ U , by setting it to zero 2 .
In this way, the PDE becomes linear and we may apply the Discrete Fourier Transform (DFT) to the update schemes, which in the absence of the gradient are identical.
Taking the DFT with respect to x on both sides of the homogeneous part yields
[Equation]
where ω = ( ω 1 , . . . , ω N ) represents the spatial frequency, and ˆ u n represents the DFT of u n .
Assuming a complex ampliﬁer ξ ( ω ) , we may substitute ˆ u n +1 ( ω ) =
ξ ( ω )ˆ u n ( ω ) for all n , and thus ﬁnd the following equation for ξ :
[Equation]
Noting that the coefﬁcients of this quadratic polynomial in ξ are all real (and that the leading coefﬁcient is positive) we may use the Root Amplitude Lemma [3] to show that the stability condition | ξ | ≤ 1 is obtained if and only if the following two inequalities are satisﬁed:
[Equation]
The upper inequality is satisﬁed as a, b, ∆ t, ∆ x ≥ 0 .
The bottom inequality is satisﬁed if and only if its left hand side is non-negative, leaving us with
[Equation]
which, in turn, must be satisﬁed for all digital frequencies ω k .
This yields the result speciﬁed in the
Theorem.
C PDE Acceleration Applied to Weight Tensors in CNNs
We now specify how Accelerated PDE are applied to CNNs.
For a given layer, θ will denote the convolutional ﬁlters or batch norm weight tensors.
In the former case, θ is a 4D tensor, i.e., θ is a function of h, w (the spatial dimensions of the ﬁlters) and i, o (the input and output channel directions).
In the case of batch norm weights, θ is a 1D tensor of the scalings or a 1D tensor of the bias.
See Figure 2.To implement our scheme as a modiﬁcation of SGD-Nesterov, we specify the computation of the diffusion term ∇ 2 v , as the other terms are the same as SGD-Nesterov.
In the case of the convolutional ﬁlter weight tensor, we can compute the Laplacian in all or a subset of directions.
In experiments, we explore the 2D spatial diffusion, and 1D output and input as well as the 1D diffusions of the batch norm tensors.
The Laplacian is approximated in these directions with central differences, using (5), but only in the desired subset of the N directions.
A Neumann boundary condition is used.
Our implementation is shown in Algorithm 1, which over standard SGD Nesterov requires just an extra diffusion computation, as well as a few formulas to compute β to ensure a stable scheme. Algorithm
Input: momentum γ, learning rate η, max iterations M, dimensions of diffusion N
Initialize: θ 0 , v1: Compute L 0 = ∇ 2 v 0 using (5) 2: while n < M do 3:
Compute ∆ t = (cid:113) 2 η γ +1 , and a = 2 − 2 γ √ 2 η ( γ +1) 4:
Compute b max = 1 2 N ∆ t 5:
Compute β
= 2 b max ∆ t 2+ a ∆ t 6:
Update v n +1 = γv n + ∇ θ J ( θ n ) + βL n 7:
Compute L n +1 = ∇ 2 v n +1 using (5) 8:
Update θ n +
=
θ n − η [ ∇ J ( θ n ) + γv n +1 + βL n +1 ] 9:
Update n ← n +10: end while
D Image classiﬁcation experiments
Experiments are shown here on CNNs (DenseNet [9], WideResNet[28]) and datasets (Cifar10,
Cifar100, SVHN, Fashion-MNIST) to evaluate the beneﬁt of adding diffusion.
We use the cosine annealing schedule with warm restarts [12] for all experiments and Auto Augment [6] is also tested.
The diffusion is simply added to the weight of Batch Normalization layer, which is a 1-D tensor.
So the diffusion term in equation (7) is a 1-D spatial Laplacian
and we use periodic padding for boundary condition.
And for the ﬁrst epoch, we use the maximum stable diffusion computed from equation 11 and we reduce the diffusion linearly to zero until the last epoch to achieve the best performance.
Other parameters still use the standard SGD update scheme.
As seen in Table 4, diffusion outperforms standard SGD with statistical signiﬁcance in these cases.
The experimental results represent the average of 5-10 independent trials.
This is just the preliminary result we have on the practical training schedule. However, we have not done enough research on how to ﬁt the diffusion in the practical case. There is more potential for the beneﬁt of diffusion and we will further explore it [Table]
[Table]
In the ofﬁcial document, the authors claim they implement a modiﬁed version of Nesterov, which has the following form: [Equation] Deﬁne θ (cid:48)
n = θ n − γηv n Then [Equation] θ
[Equation]
And γηv
[Equation]
So the original form could be written as [Equation] [Equation] which is the algorithm implemented in the PyTorch SGD optimizer.