On Regularization of Gradient Descent,
Layer Imbalance and Flat Minima

Following Hochreiter et al [5], we are interested in flat minima – “a region in weight space with the property that each weight from that region has similar small error". In contrast, sharp minima are regions where the function can increase rapidly. Let’s compute the loss gradient $\nabla L({\bm{w}})$:

Here we denote $\mathbf{W/w_{i}:=w_{1}\cdots w_{i-1}\cdot w_{i+1}\cdots w_{d}}$ for brevity. The minima of loss $L$ are located on hyperbola $w_{1}\cdots w_{d}=1$ (Fig. 1). Our interest in flat minima is related to training robustness. Training in the flat area is more stable than in the sharp area: the gradient $\dfrac{\partial L}{\partial w_{i}}$ vanishes if $|w_{i}|$ is very large, and the gradient explodes if $|w_{i}|$ is very small.

It was suggested by Hochreiter et al [6] that flat minima have smaller generalization errors than sharp minima. Keskar et al. [7] observed that large-batch training tends to converge towards a sharp minima with a significant number of large positive eigenvalues of Hessian. They suggested that sharp minima generalize worse than flat minima, which have smaller eigenvalues. In contrast, Dinh et al. [4] argued that flatness of minima can’t be directly applied to explain generalization; since both flat and sharp minima represent the same function, they perform equally on a validation set.

The question of how minima flatness is related to good generalization is out of scope of this paper.

In this section we define a new metric related to the flatness of the minimizer – layer imbalance.

Dinh [4] showed that minima flatness is defined by the largest eigenvalue of Hessian $H$:

The eigenvalues of the Hessian $H({\bm{w}})$ are $\{0,\dots,0,\sum{\dfrac{1}{w_{i}^{2}}}\}$. Minima close to the axes are sharp. Minima close to the origin are flat. Note that flat minima are balanced: $|w_{i}|\approx 1$ for all layers.

In the spirit of [1, 9], let’s define layer imbalance for a deep linear network:

Minima with low layer imbalance are flat, and minima with high layer imbalance are sharp.

We start with an analysis of training dynamics for continuous GD. By taking a time limit for gradient descent: $w_{i}(t+1)=w_{i}(t)-\lambda\cdot\nabla L({\bm{w}})$, we obtain the following DEs [10]:

For continuous GD, the loss function monotonically decreases:

The trajectory for continuous GD is hyperbola: $w_{i}^{2}(t)-w_{j}^{2}(t)=$ const (see Fig. 2(a)) [10] . The layer imbalance remains constant during training. So if training starts close to the origin, then a final point will also have a small layer imbalance and a minimum will be flat.

Let’s turn from continuous to regular gradient descent:¹¹1We omit $t$ in the right part for brevity, so $w_{i}$ means $w_{i}(t)$.

We would like to find conditions, which would guarantee that the loss monotonically decreases. For any fixed learning rate, one can find a point ${\bm{w}}$, such that the loss will increase after the GD step.²²2For example, consider the network with 2 layers. The loss $L$ after GD step is: $\displaystyle L(t+1)=\big{(}w_{1}(t+1)w_{2}(t+1)-1\big{)}^{2}=\big{(}(w_{1}-2\lambda(w_{1}w_{2}-1)w_{2})(w_{2}-2\lambda(w_{1}w_{2}-1)w_{1})-1\big{)}^{2}$ $\displaystyle=(w_{1}w_{2}-1)^{2}\big{(}1-2\lambda(w_{1}^{2}+w_{2}^{2})+4\lambda^{2}(w_{1}w_{2}-1)\big{)}^{2}=L(t)\Big{(}1-2\lambda(w_{1}^{2}+w_{2}^{2})+4\lambda^{2}(w_{1}w_{2}-1)\Big{)}^{2}$ For any fixed $\lambda$, one can find $(w_{1},w_{2})$ with $w_{1}\cdot w_{2}\approx 1$ and large enough $(w_{1}^{2}+w_{2}^{2})$ to make $|1-2\lambda(w_{1}^{2}+w_{2}^{2})+4\lambda^{2}(w_{1}w_{2}-1)|\gg 1$, and therefore the loss will increase: $L(t+1)>L(t)$. But we can define an adaptive learning rate $\lambda({\bm{w}})$ which guarantees that the loss decreases.

Note. We proved only that the layer imbalance $D$ decreases, but not that $D$ converges to $0$. The layer imbalance may stay large, if the loss $L\rightarrow 0$ too fast or if $W\approx 0$, so the factor $k=1-4\lambda^{2}\cdot L\cdot W^{2}(1/(w_{i}w_{j}))^{2}\rightarrow 1$. To make the layer imbalance $D\rightarrow 0$, we should keep the loss in certain range, e.g. $\frac{1}{4}<|W-1|<\frac{1}{2}$. For this, we could increase the learning rate if the loss becomes too small, and decrease learning rate if loss becomes large.

As before, we consider the gradient descent for linear network $(w_{1},\dots,w_{d})$ with $d$ layers. Let’s add the weight decay (WD) term to the loss: $\bar{L}({\bm{w}})=(w_{1}\cdots w_{d}-1)^{2}+\mu(w_{1}^{2}+\dots+w_{d}^{2})$.

The continuous GD with weight decay is described by the following DEs:

Accordingly, the loss dynamics for continuous GD with weight decay is:

The loss decreases when $k=(W-1)\big{(}W-(1-\mu\dfrac{d}{W(\sum 1/w_{i}^{2})})\big{)}>0$, outside the weight decay band: $1-\mu\dfrac{d}{W(\sum 1/w_{i}^{2})}\leq W\leq 1$. The width of this band is controlled by the weight decay $\mu$.

We can divide GD training with weight decay into two phases: (1) optimization and (2) regularization. During the first phase, the loss decreases until the trajectory gets into the WD-band. During the second phase, the loss $L$ can oscillate, but the trajectory stays inside the WD-band (Fig. 2(b)) and goes toward a flat minima area. The layer imbalance monotonically decreases:

Bishop [3] showed that for shallow networks, training with noise is equivalent to Tikhonov regularization. We extend this result to DLNs.

Let’s augment the training data with noise: $\tilde{x}=x\cdot(1+\eta)$, where the noise $\eta$ has $0-$mean and is bounded: $|\eta|\leq\delta<\frac{1}{2}$. The DLN with noise augmentation can be written in the following form:

This model also describes continuous dropout [11] when layer outputs $h_{i}$ are multiplied with the noise: $\tilde{h}_{i}=(1+\eta)\cdot h_{i}$. This model can be also used for continuous drop-connect [8, 12] when the noise is applied to weights: $\tilde{w}_{i}=(1+\eta)\cdot w_{i}$.

The GD with noise augmentation is described by the following stochastic DEs:

Let’s consider loss dynamics:

The loss decreases while the factor $k=(W-1)(W-\dfrac{1}{1+\eta})=(W-1)(W-1-\dfrac{\eta}{1+\eta})>0$, outside of the noise band $1-\dfrac{\delta}{1+\delta}<W<1+\dfrac{\delta}{1-\delta}$. The training trajectory is the hyperbola $w_{i}^{2}(t)-w_{j}^{2}(t)=$ const. When the trajectory gets inside the noise band, it oscillates around the minima manifold, but the layer imbalance remains constant for continuous GD.

Consider now discrete GD with noise augmentation:

For discrete GD, noise augmentation works similarly to weight decay. Training has two phases: (1) optimization and (2) regularization (Fig. 2(c)). During the optimization phase, the loss decreases until the trajectory hits the noise band. Next, the trajectory oscillates inside the noise band, and the layer imbalance decreases. The noise variance $\sigma^{2}$ defines the band width, similarly to the weight decay $\mu$.

Note. We can prove that the layer imbalance $E[D]\rightarrow 0$ if we also assume that all layers are uniformly bounded $|w_{i}|<C$. This implies that there is $\epsilon>0$ such that for all ${\bm{w}}$ the adaptive learning rate $\lambda({\bm{w}})>\epsilon$, and we can prove that the expectation $E(k)<1$:

This proves that the layer imbalance $D\rightarrow 0$ with rate $\big{(}1-\dfrac{\lambda^{2}\sigma^{2}}{C^{4}}\big{)}$.