Accelerating Training of Batch Normalization: A Manifold Perspective
Abstract
Batch normalization (BN) has become a critical component across diverse deep neural networks.
The network with BN is invariant to positively linear re-scale transformation, which makes there exist infinite functionally equivalent networks with different scales of weights.
However, optimizing these equivalent networks with the first-order method such as stochastic gradient descent will obtain a series of iterates converging to different local optima owing to their different gradients across training.
To obviate this, we propose a quotient manifold PSI manifold , in which all the equivalent weights of the network with BN are regarded as the same element.
Next, we construct gradient descent and stochastic gradient descent on the proposed
PSI manifold to train the network with BN.
The two algorithms guarantee that every group of equivalent weights (caused by positively re-scaling) converge to the equivalent optima.
Besides that, we give convergence rates of the proposed algorithms on the PSI manifold. The results show that our methods accelerate training compared with the algorithms on the Euclidean weight space.
Finally, empirical results verify that our algorithms consistently improve the existing methods in both convergence rate and generalization ability under various experimental settings.
INTRODUCTION
Batch normalization (BN) [Ioffe and Szegedy, 2015] is one of the most critical innovations in deep learning, which appears to help optimization as well as generalization [Ioffe and Szegedy, 2015, Santurkar et al., 2018].
Despite the success of BN, there is one adverse effect in terms of op- timization due to the positively scale-invariant (PSI) property brought by it.
The PSI property is explained as the weights in each layer with BN are invariant to positively linear re-scaling.
Due to this, there can be an infinite number of networks functionally equivalent to each other but with various scales of weights.
These networks can converge to different local optima owing to different gradients [Cho and Lee, 2017, Huang et al.].
Then the converged point can be sensitive to the scale of parameters, and some of them may have poor performances (see Section 4).
Hence, it is desirable to obviate such ambiguity of training.
To this end, we leverage the technique of optimization on manifold [Absil et al., 2009, Cho and Lee, 2017, Huang et al., Badrinarayanan et al., 2015].
We propose to constrain the scale-invariant weights of deep neural networks with BN in a quotient manifold i.e. PSI manifold defined in Section 4, in which all the positively scale-equivalent weights are viewed as the same element.
Constraining the scale-invariant weights on the PSI manifold maintains the representation ability of hypothesis space due to the PSI property of network with BN.
More importantly, optimizing on the manifold breaks the training ambiguity caused by the
PSI property.
By constructing the Riemannian metric and retraction function [Absil et al., 2009] on such PSI manifold [Boumal et al., 2019], we propose the gradient descent (GD), stochastic gradient descent (SGD), and SGD with momentum on such manifold.
We abbreviate the corresponding algorithms as PSI-GD, PSI-SGD, and PSI-SGDM, respectively.
Compared with vanilla GD and SGD on the Euclidean weight space, the proposed algorithms guarantee the equivalent scale-invariant weights (caused by positively re-scaling) converge to the equivalent local optima.
In contrast to the literature of optimizing network with BN while constraining parameters on manifold [Cho and Lee, 2017, Huang et al., Badrinarayanan et al., 2015,?], we give the convergence rates of our PSI-GD and PSI-SGD.
The convergence rates of the two methods are respectively in the order of O (1 /T ) and O (1 / √ T ) under the non-convex and smooth assumptions, where T is the number of iterations.
The results match the optimal one with fine-tuned learning rates [Ghadimi and Lan, 2013] under non-convex optimization problem.
Besides that, the proposed algorithms are actually training with adaptive learning rate, which is decided by the local smoothness of loss function.
The adaptive learning rate accelerates the training process as our proposed algorithms have improved constant dependence in the convergence rates, compared with their vanilla versions on the Euclidean space.
To empirically study the proposed algorithms, we compare
PSI-SGD and PSI-SGDM
with the other three baseline algorithms.
They are respectively vanilla SGD, adaptive learning rate algorithm Adam [Diederik P. Kingma, 2015], and another manifold based algorithm SGDG
[Cho and Lee, 2017].
The experiments are conducted on image classification task on datasets CIFAR and ImageNet [Krizhevsky and Hinton, 2009, Deng et al., 2009].
We empirical observe the proposed methods consistently improve the baseline algorithms in the sense of both convergence rate and generalization over various experimental settings.
The observed accelerated convergence rate and better generalization ability of
PSI-GD and PSI-SGD verify our theoretical justification that the two method are able to find local minima with less iterations.
RELATED WORKS
Optimization on Manifold.
Absil et al.
[2009] presents an introduction and summarization to this topic.
The foundations and some recent theoretical results in this topic refers to [Absil and Gallivan, 2006, Liu et al., 2017, Zhang and Sra, 2016, Boumal et al., 2019].
Roughly speaking, optimization on manifold converts the constrained problem into an unconstrained problem on a specific manifold. ?
Lezcano-Casado and Martınez-Rubio [2019], Casado [2019], Li et al.
[2019] leverage the technique to handle some explicit constraints to deep neural networks e.g., orthogonal constraints.
[Cho and Lee, 2017, Huang et al., Badrinarayanan et al., 2015] fix the ambiguity of optimizing network with BN by manifold based algorithms.
However, in contrast to this paper, some of the theoretical properties of their methods are missing, e.g., the convergence rate.
PSI property.
Our method in this paper is build upon PSI property brought by BN.
The PSI property properly defined in Section 4 are widely appears in modern deep neural networks.
For example, the network with ReLU activation [Neyshabur et al., 2015] and the network with normalization layer, e.g., batch normalization [Ioffe and Szegedy, 2015], layer normalization [Ba et al.] and group normalization [Wu and He, 2018].
Arora et al.
[2018], Wu et al.
[2018] use a similar theoretical framework as ours to show the network with PSI property allows a more flexible choice of the learning rate, and thus easier to be trained.
However, we design algorithm to obviate the optimization brought by PSI property, and show the proposed methods can have improved convergence rate.
Due to the PSI property, Li and Arora [2019] uses exponentially increased learning rate for the network with BN to handle implicit effect brought by weight decay.
In contrast to fixing the ambiguity of optimization by manifold based methods, Neyshabur et al.
[2015], Meng et al.
[2018], Zheng et al.
[2019] re-parameterize the network by “path value” and optimizing on the path value space.
BACKGROUND 3.1 PROBLEM SET UP
In this paper, we consider the deep neural network with the following structure
[Equation]
where ϕ ( · ) is the non-linear activation, and W l is the weight matrix of the l -th layer.
BN ( · ) operator in the equation represents batch normalization [Ioffe and Szegedy, 2015] layer, which normalizes every hidden output z = θ T x across a batch of training samples as
[Equation]
where γ and β are learnable affine parameters.
The loss is
[Equation]
n [Equation] across this paper, where { ( x i , y i ) } is the training set, ℓ ( x , y ) represents the loss function, and g are scalevariant parameters i.e. γ and β in BN layer.
Let W = ( w (1) · · · , w ( m ) ) , the w ( i ) is the weight vector of the i -th neuron.
Then L ( W , g ) is positively scale-invariant w.r.t.
W .
It means L ( W , g ) =
L ( T a ( W ) , g ) for any T a ( W ) =
( a 1 w (1) , · · · , a m w ( m ) ) with { a i > 0 } .
In the right hand side of Figure 1, we illustrate such PSI property (the network is invariant for rescaling operator T a ( · ) ) via a two layer network, where the number between two circles (neurons) is the weight between them.
Specially, let V be the normalized W which means V = ( w (1) / ∥ w (1) ∥ , · · · , w ( m ) / ∥ w ( m ) ∥ ) .
Then we have L ( V , g ) =
L ( W , g ) .
It worth noting that our discussions below can be applied on the top of any functions with PSI property e.g., ReLU neural networks
[Meng et al., 2018, Neyshabur et al., 2015].
However, we use NN with BN as a representation to illustrate our conclusions.
3.2 OPTIMIZATION ON MANIFOLD
We present a brief introduction to some definitions appeared in the optimization on manifold, more details are deferred to Appendix A or [Absil et al., 2009].
A Matrix Manifold M is a subset of Euclidean space that for any x ∈ M , there exists a neighborhood U x of x , such that U x is homeomorphic to a Euclidean space.
It has been pointed out that the PSI parameters W can be defined on a matrix manifold, e.g., Grassmann manifold or Oblique manifold without losing representation capacity of model [Cho and Lee, 2017, Huang et al.].
For any x ∈ M , there is a tangent space
T x M , such that for any f ( · ) defined on M , its Riemannian gradient grad f ( x ) ∈ T x M 1 .
With a given retraction function R x ( · ) : T x M → M , it was shown in [Boumal et al., 2019] that the gradient descent on manifold (refers to Figure 2)
[Equation]
has the same convergence rate with vanilla gradient descent in Euclidean space under some mild assumptions.
Roughly speaking, the Riemannian gradient descent creates a series of iterates moving along the direction of grad f ( x ) , while the moving is conducted via the pre-defined retraction function R x ( · ) .
PSI MANIFOLD
The optimization ambiguity of PSI parameters states as follows.
By chain rule,
[Equation]
so the two positively scale-equivalent points W and T a ( W ) with the same output converge to different local optima owing to the different gradients.
In other words, gradient-based algorithms in Euclidean space can not guarantee two iterates started from W and T a ( W ) positively scale equivalent with each other across the training.
This makes the converged point becomes sensitive the scale of weights.
Unfortunately, the sensitivity sometimes makes model converge to a local minima with poor performance.
For example, we train a simple two-layer neural network with BN on
MNIST [LeCun et al., 1998], which is a image classification task.
The network is trained with a unbalanced initialization where the unbalanced here means the weights in each layer are in different scales.
As can be seen, in the left hand side of
Figure 1, in contrast to our PSI-SGD (see Section 4.2), SGD performs very poorly under such initialization. However, the poor performance of SGD does adopted under balanced initialization [LeCun et al., 1998]. Thus the optimization ambiguity can break the success of training. To obviate this, we construct a quotient manifold, i.e., PSI manifold.
Constraining the model’s parameters in such manifold can fix the training ambiguity brought by the positively scale-invariant property.
4.1 CONSTRUCTION OF PSI MANIFOLD
A quotient manifold is induced by an equivalent relationship.
We formally formulate the positively scale-equivalence to derive the PSI manifold with PSI parameters defined in it.
Definition 1 (Positively Scale-Equivalent) .
For a pair of
W , W ′ , they are positively scale-equivalent with each other, if there exists a transformation T a ( · ) such that W ′ =
T a ( W ) .
W ∼ W ′ denotes the equivalence.
We can verify that the positively scale-equivalent “ ∼ ” is an equivalent relationship.
The Proposition 3.4.6 in [Absil et al., 2009] implies the following proposition which constructs the PSI manifold.
Proposition 1.
For PSI parameters W , the positively scaleequivalent relationship ∼ induces a quotient manifold M =
M / ∼ = R d 1 × · · · R d m / ∼ named PSI manifold.
Here d i is the dimension of w ( i ) .
We formally defined the PSI manifold.
Intuitively, all the parameters that are positively scale-equivalent with each other in the Euclidean space are viewed as the same element in the PSI manifold M .
Then, we can constrain the PSI parameters directly in the manifold without losing the representation ability of the model.
The PSI manifold is proposed to obviate the ambiguity of optimization brought by the PSI property.
To this end, we can optimize the PSI parameters directly in the PSI manifold, since all the positively scale-equivalent parameters are the same point in it.
But (4) implies that optimization on PSI manifold requires the formulation of Riemannian gradient and retraction function on it.
The following proposition defines a Riemannian metric in the PSI manifold, which formulates the Riemannian gradient (See the definition in
Appendix A).
Proposition 2.
For every W ∈ M 2 , the Riemannian metric ⟨· , ·⟩ W in the PSI manifold can be defined as
[Equation]
where Ξ 1 , Ξ 2 ∈ T W M , with Ξ k = ( ξ (1) k , · · · ξ ( m ) k ) .
The proof of this Proposition is in Appendix B.1.
With the
Riemannian metric, we can exactly compute the Riemannian gradient on the PSI manifold.
According to Definition of
Riemannian gradient in Appendix A, we have
[Equation]
for each w ( i ) .
The retraction function R W ( · ) is required to obtain gradient-based algorithms on the PSI manifold.
Proposition 3.
For every W ∈ M , Ξ ∈ T W M , a retraction function R W ( Ξ ) on the PSI manifold can be
[Equation]
The proof of this proposition appears in Appendix C.1.
Now we are ready to optimize PSI parameters in the PSI manifold directly, which obviates the ambiguity of training
4.2 OPTIMIZATION ON THE PSI MANIFOLD
In this subsection, we give the update rules of GD, SGD, and
SGD with momentum on the PSI manifold, abbreviated as
PSI-GD, PSI-SGD, and PSI-SGDM respectively.
We show that all of these update rules on the PSI manifold generate a unified optimization path for parameters that positively scale-equivalent with each other.
Combining (4) and Proposition 2, 3, we have the following update rule of GD
[Equation]
and SGD with momentum
[Equation]
for PSI parameters.
Here
[Equation]
for a batch of { ( x 1 , y 1 ) , · · · , ( x S , y S ) } is an unbiased estimation to the Riemannian gradient.
In addition, the update rules of non-scale-invariant parameters g follows the gradient-based algorithms in Euclidean space. Remark 1.
The proposed PSI-GD is similar to the one in [Badrinarayanan et al., 2015]. However, they do not prove that the update rule is induced by a retraction function. Besides that, our theoretical characterization of its convergence rate is absent in their work. One can justify that the proposed algorithms on PSI manifold are essentially using adaptive learning rates decided by ∥ w ( i ) t ∥ 2 to match the gradient scale brought by a variant of weight scale.
The following proposition shows the well-posedness of the algorithms on PSI manifold. Sk t
Theorem 1.
For two positively scale-equivalent weights
W 0 and T a ( W 0 ) , let W t , ˆ W t be t -th iterate of
PSI-SGDM respectively started from { W 0 , U 0 } and { T a ( W 0 ) , T a ( U 0 ) } .
Then we have ˆ W t =
T a ( W t ) .
This theorem can be easily generalized to PSI-GD and PSISGD; the proof of it is in Appendix C.1.
The conclusion shows that the iterates generated by the update rules on the PSI manifold are equivalent with respect to applying
T a ( · ) .
Hence, optimizing on the PSI manifold obviates the optimization ambiguity brought by the PSI property.
The complete algorithm of PSI-SGDM refers to Algorithm 1.
With ρ = 0 , the PSI-SGDM degenerates to PSI-SGD.
Remark 2.
The update rule of SGD with momentum on manifold [Cho and Lee, 2017, Liu et al., 2017] involves the parallel transformation to make sure the U t +1 locates in the tangent space T W t M .
However, it requires the proposed retraction in Proposition 3 to be an exponential retraction map [Absil et al., 2009], which may fail in-practical.
Hence, we heuristically propose the PSI-SGDM without verifying the retraction function.
Even though, our method is welldefined to positively scale-transformation.
OPTIMIZATION ON PSI MANIFOLD ACCELERATES TRAINING
In this section, we give the convergence rates of PSI-GD and
PSI-SGD, and show that they accelerate training compared with vanilla GD and SGD on the Euclidean space.
5.1 CONVERGENCE RESULTS
We give the convergence rates of PSI-GD and PSI-SGD (Algorithm 1) in this subsection.
For the PSI parameters, let
V t =
( w (1) t / ∥ w (1) t ∥ , · · · , w ( m ) t / ∥ w ( m ) t ∥ ) be the normalized parameters, we impose the following assumptions to characterize the smoothness of loss function.
[Equation]
Here ∥ · ∥ 2 is the spectral norm of matrix.
For PSI-SGD, we further impose the bounded variance assumption:
[Equation]
Remark 3.
Due to (5) , the upper bound (14) can not hold as ∥ w i ∥ goes to zero.
However, due to Lemma 1 below, the weights w i obtained by GD or SGD will never goes to zero.
Thus the bounded variance (14) is reasonable in this regime.
Let L ( W ∗ , g ∗ ) =
inf W , g L ( W , g ) , the following theorems give the convergence rates of PSI-GD and PSI-SGD.
Theorem 2.
Let { W t , g t } be the iterates of PSI-GD (9) , then we have
[Equation]
by choosing η ( i ) w t
= 1 / ˜ L v ( i ) and η g t = 1 / ˜ L g .
Here ˜ L = max { ˜ L v (1) , · · · , ˜ L v ( m ) , ˜ L g } , where ˜ L v ( i ) =
L vg i + (cid:80)
[Equation]
Theorem 3.
Let { W t , g t } be the iterates of PSI-SGD (10) , then we have
[Equation]
by choosing η w ( i ) t
= 1 ˜ L v ( i ) √ T and η g t = 1 ˜ L g √ T , where ˜ L, ˜
L v ( i ) and ˜ L g are defined in Theorem 2.
The theorems are respectively proved in Appendix D.3 and D.4.
The convergence rates of PSI-GD and PSI-SGD are respectively O (1 /T ) and O (1 / √ T ) .
The convergence rates of PSI-SGDM can be also obtained by combining the technique in [Yang et al., 2016] and our proof of Theorem 3.
One may note the convergence rate is computed on the gradient w.r.t. normalized parameters V t .
But (5) implies
[Equation]
if ∥ w ( i ) t ∥ 2 ≥ 1 for 1 ≤ i ≤ m .
Thus we can get the corresponded convergence rates of PSI-GD and PSI-SGD.
They match the optimal results with well tuned learning rates in the Euclidean space [Ghadimi and Lan, 2013].
The following lemma shows that the ∥ w ( i ) t ∥ 2 keeps increasing across training for gradient-based algorithms.
Lemma 1 (Lemma 2.4 in [Arora et al., 2018]) .
For any PSI weight w ( i ) , w ( i ) and ∇ w ( i ) ℓ ( f ( x k , W ) , y k ) are perpendicular for any ( x k , y k ) .
On the other hand
[Equation] Thus, ∥ w ( i ) t ∥ 2 ≥ 1 holds, if ∥ w ( i )
∥ 2 ≥ 1 .
Since the network is usually initialized by w ( i ) ∼ N (0 , 2 / √ d i )
[He et al., 2015], ∥ w ( i )
∥ 2 ≈ 2 with high probability.
Hence, we can conclude (cid:13)(cid:13) ∇ W t , g t L ( W t , g t ) (cid:13)(cid:13) 2 ≤ (cid:13)(cid:13) ∇ V t , g t L ( V t , g t ) (cid:13)(cid:13) 2 .
We will show the inequality and the increasing ∥ w ( i ) t ∥explain the acceleration of the algorithms on PSI manifold.
5.2 WHY OPTIMIZATION ON THE PSI MANIFOLD ACCELERATES TRAINING
Next, we show PSI-GD and PSI-SGD accelerate training compared with the algorithms on Euclidean space.
Generally, PSI parameters move towards a smoother region (smaller gradient Lipschitz constant) across training, since the gradient Lipschitz constant is in inverse ratio to the weight scale which keeps increasing according to Lemma 1.
As smoother region allows a larger learning rate, gradually increasing the learning rate accelerates training.
Fortunately, PSI-GD and PSI-SGD happen to be vanilla GD and SGD with adaptively increased learning rates according to (9), (10) and (18).
To begin with, we verify the convergence rates of vanilla GD and SGD.
We assume a global smoothness for L ( W , g ) ,
[Equation]
The assumption is stronger than (13) since (13) only involves the normalized parameters, thus L vv ij ≤
L ww ij ; L vg i ≤ L wg i .
We consider the update rule of GD
[Equation]
and SGD
[Equation]
where G w ( i ) t ( W t , g t ) and G g t ( W t , g t ) are respectively unbiased estimations of ∇ w ( i ) t L ( W t , g t ) and ∇ g t L ( W t , g t ) defined in (12).
For SGD, we assume the two estimators satisfy the bounded variance assumption (14).
The following two Theorems characterize the convergence rate of GD and SGD.
Theorem 4.
Let { W t , g t } updated by GD (20) , by choosing η ( i ) w t =
1 / ˜
L w ( i ) and η g t = 1 / ˜ L g ,
[Equation]
t ,
[Equation]
Here ˜ L = [Equation] [Equation]
[Equation]
˜ L
[Equation]
[Equation]
, ˜ L g } , where ˜ L w (
i
[Equation]
=
L
[Equation]
i
[Equation]
Theorem 5.
Let { W t , g t } updated by SGD (21) , then
[Equation]
by choosing η w ( i ) t = 1 ˜ L w [Equation] i ) √ T and η g t = 1 ˜ L g √ T , where ˜ L, ˜
L w ( i ) and ˜ L g are defined in Theorem 4.
The two convergence rates of ∥∇ W , g L ( W , g ) ∥ match the optimal results with well tuned learning rates [Ghadimi and
Lan, 2013].
Since L vv ij ≤ L ww ij ; L vg i ≤ L wg i , the right hand side in (22) is smaller than the one in (15).
To see the results for SGD, under (19), the upper bound in (16) can be replaced by the smaller one of the upper bounds in (16) and (23).
Thus the convergence rate of PSI-SGD is also sharper than the one of SGD.
On the other hand, (17) and the discussion in the above section implies [Equation]
∥ 2 .
Thus, from (17), one can verify the proposed algorithms accelerate training in a factor ∥ w ( i ) t ∥ 2 ≥ 1 compared with the versions on Euclidean space.
In addition, ∥ w ( i ) t ∥ 2 keeps increasing to a bounded constant across training, thus the acceleration is more significant after a number of iterations.
A detailed discussion to ∥ w ( i ) t ∥ 2 is in Appendix E.
The proofs of the two theorems are delegated to Appendix
D.1 and D.2.
In the proof, we see that scaling the learning rate with 1 / ˜ L w ( i ) and 1 / ˜ L g are respectively the optimal schedule of η w ( i ) t and η g t .
Since smaller ˜
L w ( i ) and ˜ L g corresponds with a smoother loss landscapes, it explains that
L ( W , g ) allows a larger learning rate in a smoother region to get the optimal convergence rate.
Now we are ready to illustrate the reason of PSI-GD and
PSI-SGD’s acceleration.
Due to L ( W , g ) =
L ( T a ( W ) , g ) ,
[Equation]
[Equation]
[Table]
The fact shows that the smoothness of PSI parameters increasing with its scale.
As larger a in the right hand side of (24) results in smaller ∥∇ 2 a i w ( i ) a j w ( j ) L ( T a ( W ) , g ) ∥ 2 and ∥∇ 2 a i w ( i ) g L ( T a ( W ) , g ) ∥ 2 .
Due to Lemma 1, the increasing ∥ w ( i ) t ∥ implies the loss landscape of PSI parameters becomes smoother across training.
Thus, gradually increasing the learning rate to update PSI parameters can accelerate training.
The proposed PSI-GD and PSI-SGD are proven to be vanilla GD and SGD with the increasing learning rate that is proportional to ∥ w ( i ) t ∥ 2 , which interprets the acceleration.
Finally, (24) shows that our algorithms are optimal in the manner of leveraging the smoothness to accelerate training.
The intuition is that the proposed methods have a more accurate estimation of local smoothness across training.
We point out that L vv ij ≤ L ww ij ;
L vg i ≤ L wg i in (19) also gives a sharper convergence rate of PSI-GD and PSI-SGD.
EXPERIMENTS 6.1 IMPROVED CONVERGENCE RATE
We use a toy example to verify that our algorithms have improved convergence rate compared with other methods.
Data.
We sample 1000 training samples { x i } from 10dimensional normal distribution with its label y i = µ ⊤ x i + ϵ i for ϵ i ∼ N (0 , 1) .
Setup.
We use a toy neural network f ( x ) =
w ⊤ 1 ϕ ( W 2 x ) with w 1 ∈ R 100 , W 2 ∈ R 100 × 10 , and the activation function ϕ ( · ) is ELU function.
The training loss is MES loss.
We compare the convergence rates in terms of gradient norm of the proposed algorithms PSI-SGD and PSI-SGDM with three baselines algorithms.
The benchmark optimization algorithm SGD with momentum (abbrev.
SGD); adaptive learning rate algorithm Adam [Diederik P. Kingma,
2015]; manifold based algorithm SGDG
[Cho and Lee, 2017] which is applied to the network with BN.
Main Results.
The convergence rates of gradient norm are in Figure 4.
As can be seen, the three manifold based methods PSI-SGD, PSI-SGDM, SGDG exhibit improved convergence rates compared other methods.
However, we still observe that PSI-SGD and PSI-SGDM are slightly better than SGDG.
More important is that in contrast to SGDG, our algorithms have proved convergence rate.
The toy example verifies that our algorithms have improved convergence rates compared with other methods.
6.2 EXPERIMENTS ON REAL-WORLD DATASET
In this section, we empirically study the proposed algorithms
PSI-SGD and PSI-SGDM on real-world dataset.
Data.
We consider the image classification task on three benchmark datasets.
CIFAR10 and CIFAR [Krizhevsky and Hinton, 2009] are respectively colorful images with 50K training samples and 10K validation samples from 10 and 100 categories.
ImageNet [Deng et al., 2009] are colorful images with 1M+ training samples from 1K object classes.
Setup.
The model is a unified structure ResNet with various structures.
As in the above section, we compare our methods with SGD, Adam, and SGDG.
For PSI-SGD or
PSI-SGDM, the PSI parameters are updated by the two algorithms, while the other parameters are updated by SGD.
We do not use the regularizer to the PSI weights e.g., l 2regularizer in the loss function since it breaks the PSI property of PSI weights.
More experiments conducted with regularizer are in Appendix F.
For CIFAR we conduct 200 epochs of training for each algorithm.
The learning rate starts from 0.1 and decays by a factor 0.2 at epochs 60, 120, and 160.
For ImageNet , the training is conducted for 100 epochs, and the learning rate starts from 0.1 and decays by a factor 0.2 at epochs 30, 60, and 90.
For the hyperparameters of baseline methods, we grid search the learning rates in the range of {0.01, 0.1, 1.0} and {0.0001, 0.001, 0.01} respectively for SGD and
Adam, and the hyperparameters of SGDG follow the one of [Cho and Lee, 2017].
The other hyperparameters of all these methods are summarized in Appendix F.2.
It worth noting that the PSI weights updated by PSI-SGD and PSI-SGDM may overflow after quite a number of iterations.
Hence, we normalize the PSI weights w ( i ) if ∥ w ( i ) ∥ is larger than 10000 during training.
This operation will not change the output of the model due to the PSI property.
Main Results.
We report the test accuracy of each method.
The results refers to Table 1 and Figure 3.
We have the following observations and conclusions from the experimental results.
In Table 1, the model trained by manifold based methods i.e., SGDG, PSI-SGD and, PSI-SGDM generalize better in most cases.
this is due to the manifold based algorithms can obviate quite a lot of local minima with poor generalization, since they have a unified optimization path for the parameters equivalent to each other.
Besides that, the proposed PSI-SGD and PSI-SGDM are significantly better than SGDG, and the performances of PSI-SGD and PSI-SGDM are comparable.
We speculate this is due to the sharper convergence rate of PSI-SGD and PSI-SGDM allow them to find local minima in less number of iterations.
Figure 3 show that the PSI-SGD and PSI-SGDM converge faster than the three baselines after a certain number of iterations i.e., after 120 epochs of update.
This justifies our theoretical results in Section 5, since the acceleration is linearly with ∥ w ( i ) t ∥ 2 which is large after a while of training (Lemma 1).
We present the results of training error instead of gradient norm here because evaluating the gradient norm requires implementing back propagations on all training data which brings a great extra computational effort, especially for large scale dataset e.g., ImageNet.
We have one remark about the efficiency of our methods.
They are simple and efficient compared with SGDG, since
SGDG involves the operators like trigonometric functions in the update rule.
For example, the elapsed time of ResNet with respect to one-epoch training of CIFAR10 under SGD,
Adam, SGDG, PSI-SGD, and PSI-SGDM are respectively 16.98, 18.1, 26.6, 18.2, and 18.5 seconds.
All the experiments are conducted on a server with single NVIDIA V
GPU.
CONCLUSION
In this paper, we fix the optimization ambiguity brought by the PSI property of the network with BN.
Our scenario is built upon optimization on manifold by constructing a specific manifold and optimizing the PSI weights of the network with BN in it.
The developed gradient-based algorithms on
PSI manifold are shown to have a well-defined optimization path with respect to positively equivalent rescaling.
We also give the convergence rates of the proposed methods.
Besides that, we theoretically justify that PSI-GD and PSISGD accelerate training by a clever schedule of adaptive learning rate.
Finally, we conduct various experiments to show that the proposed methods have better performance in the generalization and efficiency compared with the other three baselines.