假设 F ( x ; w ) F(x;w) F(x;w)是一个输出标量的深度神经网络，其中 x x x是输入， w w w表示权重。假设 F F F关于 w w w连续可微，并且对于训练数据 { x j , y j } j = 1 m \{x_{j},y_{j}\}{j=1}^{m} {xj,yj}j=1m过参数化：即存在 w ∗ w^* w∗使得对所有 j j j满足 F ( x j ; w ∗ ) = y j F(x{j};w^*)=y_{j} F(xj;w∗)=yj。为了研究训练神经网络时在 w ∗ w^* w∗的局部优化动力学，我们考虑线性化神经网络 F ^ ( x ; w ) = F ( x ; w ∗ ) + ( w − w ∗ ) ⊤ ∇ F ( x ; w ∗ ) \widehat{F}(x;w)=F(x;w^*)+(w-w^*)^{\top}\nabla F(x;w^*) F (x;w)=F(x;w∗)+(w−w∗)⊤∇F(x;w∗)，其损失函数为
L o s s ( w ) : = 1 2 m ∑ j = 1 m ( y j − F ^ ( x j ; w ) ) 2 Loss(w):=\frac{1}{2m}\sum_{j=1}^{m}(y_{j}-\widehat{F}(x_{j};w))^{2} Loss(w):=2m1j=1∑m(yj−F (xj;w))2。
令 s s s表示学习率，梯度下降法为 w i + 1 = w i − s ∇ L o s s ( w i ) w_{i+1}=w_{i}-s\nabla Loss(w_{i}) wi+1=wi−s∇Loss(wi)，而随机梯度下降法为 w i + 1 = w i − s ( ∇ L o s s ( w i ) + ϵ i ) w_{i+1}=w_{i}-s(\nabla Loss(w_{i})+\epsilon_{i}) wi+1=wi−s(∇Loss(wi)+ϵi)，其中噪声项 ϵ i \mathcal{\epsilon}{i} ϵi满足 E ϵ i = 0 \mathbb{E}\mathcal{\epsilon}{i}=0 Eϵi=0和 E ϵ i ϵ i ⊤ = M ( w i ) / b \mathbb{E}\mathcal{\epsilon}{i}\mathcal{\epsilon}{i}^{\top}=M(w_{i})/b Eϵiϵi⊤=M(wi)/b, b b b是mini-batch的大小。假设协方差矩阵 M M M与
Σ = 1 m ∑ j = 1 m ∇ F ( x j ; w ∗ ) ∇ F ( x j ; w ∗ ) ⊤ \Sigma=\frac{1}{m}\sum_{j=1}^{m}\nabla F(x_{j};w^*)\nabla F(x_{j};w^*)^{\top} Σ=m1∑j=1m∇F(xj;w∗)∇F(xj;w∗)⊤
在以下意义上对齐:
T r ( M ( w ) Σ ) 2 ( L o s s ( w ) ) 3 2 ∥ Σ ∥ F 2 ≥ δ \frac{Tr(M(w)\Sigma)}{2(Loss(w))^{\frac{3}{2}}\|\Sigma\|_F^2}\geq\delta 2(Loss(w))23∥Σ∥F2Tr(M(w)Σ)≥δ，
对于 δ > 0 \delta>0 δ>0和所有 w w w成立。这里 ∥ ⋅ ∥ F \lVert\cdot\rVert_F ∥⋅∥F表示Frobenius范数。
(1) 对于梯度下降，证明如果 Σ \Sigma Σ的谱范数满足
∥ Σ ∥ 2 ≤ 2 s ， \lVert\Sigma\rVert_2\leq\frac{2}{s}， ∥Σ∥2≤s2，则梯度下降是局部稳定的（即对所有t，Loss ( w t ) (w_t) (wt)是有界的）。（注意，这蕴含了一个依赖维度的界: ∥ Σ ∥ F ≤ 2 d s \lVert\Sigma\rVert_F\leq\frac{2\sqrt{d}}{s} ∥Σ∥F≤s2d ,其中 d d d是 w w w的维度。)
(2) 对于随机梯度下降，如果 E L o s s ( w t ) \mathbb{E}Loss(w_t) ELoss(wt)对所有 t t t都有界，则以独立于维度的不等式必须成立:
∥ Σ ∥ F ≤ b / δ s \lVert\Sigma\rVert_F\leq\frac{\sqrt{b/\delta}}{s} ∥Σ∥F≤sb/δ 。
证：
(1)梯度下降的局部稳定性
我们需要证明在使用梯度下降时，损失函数 Loss ( w t ) \text{Loss}(w_t) Loss(wt)是有界的。
考虑梯度下降的更新规则: w i + 1 = w i − s ∇ Loss ( w i ) w_{i+1}=w_i-s \nabla \text{Loss}(w_i) wi+1=wi−s∇Loss(wi)
首先，我们计算损失函数的梯度:
∇ Loss ( w ) = 1 m ∑ j = 1 m ( F ^ ( x j ; w ) − y j ) ∇ F ^ ( x j ; w ) \nabla \text{Loss}(w) = \frac{1}{m} \sum_{j=1}^{m} \left( \widehat{F}(x_j; w)-y_j \right) \nabla \widehat{F}(x_j; w) ∇Loss(w)=m1∑j=1m(F (xj;w)−yj)∇F (xj;w)
由于 F ^ ( x ; w ) = F ( x ; w ∗ ) + ( w − w ∗ ) ⊤ ∇ F ( x ; w ∗ ) \widehat{F}(x; w) = F(x; w^*) + (w- w^*)^\top \nabla F(x; w^*) F (x;w)=F(x;w∗)+(w−w∗)⊤∇F(x;w∗)。
我们有: ∇ F ^ ( x j ; w ) = ∇ F ( x j ; w ∗ ) \nabla \widehat{F}(x_j; w) = \nabla F(x_j;w^*) ∇F (xj;w)=∇F(xj;w∗)
因此:
∇ Loss ( w ) = 1 m ∑ j = 1 m ( ( w − w ∗ ) ⊤ ∇ F ( x j ; w ∗ ) ∇ F ( x j ; w ∗ ) ⊤ ) \nabla \text{Loss}(w) = \frac{1}{m} \sum_{j=1}^{m}\left( (w - w^*)^\top \nabla F(x_j; w^*) \nabla F(x j; w^*)^\top \right) ∇Loss(w)=m1∑j=1m((w−w∗)⊤∇F(xj;w∗)∇F(xj;w∗)⊤)
定义矩阵(\Sigmal) 为:
Σ = 1 m ∑ j = 1 m ∇ F ( x j ; w ∗ ) ∇ F ( x j ; w ∗ ) ⊤ \Sigma = \frac{1}{m} \sum_{j=1}^{m}\nabla F(x_j;w^*) \nabla F(x_j; w^*)^\top Σ=m1∑j=1m∇F(xj;w∗)∇F(xj;w∗)⊤
于是:
∇ Loss ( w ) = Σ ( w − W ∗ ) \nabla \text{Loss}(w) = \Sigma (w-W^*) ∇Loss(w)=Σ(w−W∗)
现在考虑梯度下降的更新:
w i + 1 − w ∗ = w j − w ∗ − s ∇ Loss ( w i ) w_{i+1} - w^* =w_j-w^*- s \nabla \text{Loss}(w_i) wi+1−w∗=wj−w∗−s∇Loss(wi) = w i − w ∗ − s Σ ( w i − w ∗ ) =w_i-w^*-s\Sigma(w_i-w^*) =wi−w∗−sΣ(wi−w∗)
= ( l − s Σ ) ( w i − w ∗ ) =(l-s \Sigma) (w_i - w^*) =(l−sΣ)(wi−w∗)
取范数:
∥ w i + 1 − w ∗ ∥ 2 = ∥ 1 − s Σ ∥ 2 ∥ w i − w ∗ ∥ 2 \|w_{i+1} - w^*\|_2 =\|1 - s \Sigma\|_2\|w_i -w^*\|_2 ∥wi+1−w∗∥2=∥1−sΣ∥2∥wi−w∗∥2
由于 I − s Σ ∣ 2 ≤ 1 I-s\Sigma|2\leq 1 I−sΣ∣2≤1当且仅当 s ≤ 2 λ max ( Σ ) s \leq \frac{2}{\lambda{\max}(\Sigma)} s≤λmax(Σ)2，即 Σ 2 ≤ 2 s \Sigma_2\leq \frac{2}{s} Σ2≤s2，我们可以得到:
∥ w i − w ∗ ∥ 2 ≤ ∥ w 0 − w ∗ ∥ 2 \|w_i -w^*\|_2\leq \|w_0 - w^*\|_2 ∥wi−w∗∥2≤∥w0−w∗∥2
这意味着 ∥ w i − w ∗ ∥ 2 \|w_i - w^*\|_2 ∥wi−w∗∥2 是有界的，因此 Loss ( w t ) \text{Loss}(w_t) Loss(wt)也是有界的。
(2)随机梯度下降的有界性
对于随机梯度下降，我们需要证明如果 E Loss ( w t ) \mathbb{E}\text{Loss}(w_t) ELoss(wt)对所有 t t t都有界，则必须满足独立于维度的约束 ∥ Σ ∥ F ≤ b / δ s \|\Sigma\|_F\leq\frac{\sqrt{b/\delta}}{s} ∥Σ∥F≤sb/δ 。
考虑随机梯度下降的更新规则:
w i + 1 = w i − s ( ∇ Loss ( w i ) + ϵ i ) w_{i+1}=w_i-s (\nabla \text{Loss}(w_i) + \epsilon_i) wi+1=wi−s(∇Loss(wi)+ϵi)
其中 ϵ i \epsilon_i ϵi是噪声项，满足 E [ ϵ i ] = 0 \mathbb{E}[\epsilon_i]= 0 E[ϵi]=0 和 E [ ϵ i ϵ i ⊤ ] = M ( w i ) b \mathbb{E}[\epsilon_i \epsilon_i^\top] = \frac{M(w_i)}{b} E[ϵiϵi⊤]=bM(wi)。
我们需要分析 E [ Loss ( w i + 1 ) ] \mathbb{E} [\text{Loss}(w_{i+1})] E[Loss(wi+1)]。令 w i − w ∗ = z i w_i-w^*=z_i wi−w∗=zi。则： z i + 1 = z i − s ( Σ z i + ϵ i ) = ( I − s Σ ) z i − s ϵ i z {i+1}=z_i- s (\Sigma z_i+\epsilon_i)=(I-s\Sigma)zi-s\epsilon_i zi+1=zi−s(Σzi+ϵi)=(I−sΣ)zi−sϵi
取范数的平方并取期望:
E [ ∥ z i + 1 ∥ 2 2 ] = E [ ∥ ( l − s Σ ) z i − s ϵ − i _ 2 2 ] \mathbb{E}[\|z_{i+1}\|_2^2] = \mathbb{E}[\|(l- s\Sigma)z_i- s \epsilon-i\_2^2] E[∥zi+1∥22]=E[∥(l−sΣ)zi−sϵ−i_22]
= E [ ∥ ( 1 − s Σ ) z i _ 2 2 ] + s 2 E [ ∥ ϵ i ∥ 2 2 ] = \mathbb{E}[\|(1 -s \Sigma)z_i\_2^2]+ s^2\mathbb{E}[\|\epsilon_i\|_2^2] =E[∥(1−sΣ)zi_22]+s2E[∥ϵi∥22]
由于 ϵ i \epsilon_i ϵi的协方差为 M ( w i ) b \frac{M(w_i)}{b} bM(wi)。我们有:
E [ ∥ ϵ i ∥ 2 2 ] = Tr ( M ( w i ) b ) \mathbb{E}[\|\epsilon_i\|_2^2] = \text{Tr}\left(\frac{M(w_i)}{b}\right) E[∥ϵi∥22]=Tr(bM(wi))
并且:
E [ ∥ z i + 1 ∥ 2 2 ] = ∥ I − s Σ ∥ 2 2 E [ ∥ z i ∥ 2 2 ] + s 2 b Tr ( M ( w i ) ) \mathbb{E}[\|z_{i+1}\|_2^2]=\|I - s \Sigma\|_2^2\mathbb{E}[\|z_i\|_2^2] +\frac{s^2}{b} \text{Tr}(M(w_i)) E[∥zi+1∥22]=∥I−sΣ∥22E[∥zi∥22]+bs2Tr(M(wi))
为了确保 E [ ∥ z i ∥ 2 2 ] \mathbb{E}[\|z_i\|_2^2] E[∥zi∥22]的有界性，我们需要: ∥ I − s Σ ∥ 2 2 ≤ 1 \|I- s \Sigma\|_2^2\leq 1 ∥I−sΣ∥22≤1
即:
∥ Σ ∥ 2 ≤ 2 s \|\Sigma\|_2\leq \frac{2}{s} ∥Σ∥2≤s2
并且我们需要控制噪声项:
s 2 b Tr ( M ( w i ) ) ≤ C \frac{s^2}{b} \text{Tr}(M(w_i)) \leq C bs2Tr(M(wi))≤C
根据题目中的对齐条件:
Tr ( M ( w ) Σ ) 2 ( Loss ( w ) ) 3 / 2 ∥ Σ ∥ F 2 ≥ δ \frac{\text{Tr}(M(w) \Sigma)}{2 (\text{Loss}(w))^{3/2}\|\Sigma\|_F^2}\geq \delta 2(Loss(w))3/2∥Σ∥F2Tr(M(w)Σ)≥δ
我们有:
Tr ( M ( w ) Σ ) ≥ 2 δ ( Loss ( w ) ) 3 / 2 Σ ∥ F 2 \text{Tr}(M(w) \Sigma) \geq 2 \delta (\text{Loss}(w))^{3/2}\Sigma\|_F^2 Tr(M(w)Σ)≥2δ(Loss(w))3/2Σ∥F2
由于 Tr ( M ( w ) ) = Tr ( M ( w ) l ) ≤ ∥ Σ ∥ 2 Tr ( M ( w ) ) \text{Tr}(M(w)) = \text{Tr}(M(w) l)\leq \|\Sigma\|_2 \text{Tr}(M(w)) Tr(M(w))=Tr(M(w)l)≤∥Σ∥2Tr(M(w))，我们有: Tr ( M ( w ) ) ≤ Tr ( M ( w ) Σ ) ∥ Σ ∥ 2 \text{Tr}(M(w)) \leq \frac{\text{Tr}(M(w) \Sigma)}{\|\Sigma\|_2} Tr(M(w))≤∥Σ∥2Tr(M(w)Σ)
因此:
Tr ( M ( w ) ) ≤ 2 δ ( Loss ( w ) ) 3 / 2 ∥ Σ ∥ F 2 ∥ Σ ∥ 2 \text{Tr}(M(w)) \leq \frac{2 \delta (\text{Loss}(w))^{3/2} \|\Sigma\|_F^2}{\|\Sigma\|_2} Tr(M(w))≤∥Σ∥22δ(Loss(w))3/2∥Σ∥F2
为了确保 E [ Loss ( w t ) ] \mathbb{E}[\text{Loss}(w_t)] E[Loss(wt)]有界，我们需要满足:
s 2 b ⋅ 2 δ ( Loss ( w ) ) 3 / 2 ∥ Σ ∥ F 2 ∥ Σ ∥ 2 ≤ C \frac{s^2}{b} \cdot \frac{2 \delta (\text{Loss}(w))^{3/2}\|\Sigma\|_F^2}{\|\Sigma\|_2}\leq C bs2⋅∥Σ∥22δ(Loss(w))3/2∥Σ∥F2≤C
简化并得出:
∥ Σ ∥ F ≤ b / δ s \|\Sigma\|_F \leq \frac{\sqrt{b/\delta}}{s} ∥Σ∥F≤sb/δ
综上，我们得到了独立于维度的界，这证明了随机梯度下降的有界性条件。