Gradient Centralization: 簡單的梯度中心化，一行程式碼加速訓練並提升泛化能力 | ECCV 2020 Oral

曉飛的演算法工程筆記發表於2020-08-03

梯度中心化GC對權值梯度進行零均值化，能夠使得網路的訓練更加穩定，並且能提高網路的泛化能力，演算法思路簡單，論文的理論分析十分充分，能夠很好地解釋GC的作用原理

來源：曉飛的演算法工程筆記公眾號

論文: Gradient Centralization: A New Optimization Technique for Deep Neural Networks

論文地址：https://arxiv.org/abs/2004.01461
論文程式碼：https://github.com/Yonghongwei/Gradient-Centralization

Introduction

優化器(Optimizer)對於深度神經網路在大型資料集上的訓練是十分重要的，如SGD和SGDM，優化器的目標有兩個：加速訓練過程和提高模型的泛化能力。目前，很多工作研究如何提高如SGD等優化器的效能，如克服訓練中的梯度消失和梯度爆炸問題，有效的trick有權值初始化、啟用函式、梯度裁剪以及自適應學習率等。而一些工作則從統計的角度對權值和特徵值進行標準化來讓訓練更穩定，比如特徵圖示準化方法BN以及權值標準化方法WN。。

與在權值和特徵值進行標準化方法不同，論文提出作用於權值梯度的高效能網路優化演算法梯度中心化(GC, gradient centralization)，能夠加速網路訓練，提高泛化能力以及相容模型fine-tune。如圖a所示，GC的思想很簡單，零均值化梯度向量，能夠輕鬆地嵌入各種優化器中。論文主要貢獻如下：

提出新的通用網路優化方法，梯度中心化(GC)，不僅能平滑和加速訓練過程，還能提高模型的泛化能力。
分析了GC的理論屬性，表明GC能夠約束損失函式，標準化權值空間和特徵值空間，提升模型的泛化能力。另外，約束的損失函式有更好的Lipschitzness(抗擾動能力，函式斜率恆定小於一個Lipschitze常數)，讓訓練更穩定、更高效。

Gradient Centralization

Motivation

BN和WS使用Z-score標準化分別操作於特徵值和權重，實際是間接地對權值的梯度進行約束，從而提高優化時損失函式的Lipschitz屬性。受此啟發，論文直接對梯度操作，首先嚐試了Z-score標準化，但實驗發現並沒有提升訓練的穩定性。之後，嘗試計算梯度向量的均值，對梯度向量進行零均值化，實驗發現能夠有效地提高損失函式的Lipschitz屬性，使網路訓練更穩定、更具泛化能力，得到梯度中心化(GC)演算法。

Notations

定義一些基礎符號，使用$W \in \mathbb{R}^{M \times N}$統一表示全連線層的權值矩陣$W_{fc} \in \mathbb{R}^{C_{in}\times C_{out}}$和卷積層的權值張量$W_{conv} \in \mathbb{R}^{(C_{in} k_1 k_2)\times C_{out}}$，$w_i \in \mathbb{R}^{M$為權值矩陣$W$的第$i$列，$\mathcal{L}$為目標函式，$\nabla_{W}\mathcal{L}$和$\nabla_{w_i}\mathcal{L}$為$\mathcal{L}$對$W$和$w_i$的梯度，$W$與$\nabla_{W}\mathcal{L}$的大小一樣。定義$X$為輸入特徵圖，則$W}T X$為輸出特徵圖，$e=\frac{1}{\sqrt{M}}1$為$M$位單位向量(unit vector)，$I\in\mathbb{R}^{M\times M}$為單位矩陣(identity matrix)。

Formulation of GC

對於卷積層或全連線層的權值向量$w_i$，通過反向傳播得到其梯度$\nabla_{w_i}\mathcal{L}$，然後如圖b所示計算其均值$\mu\nabla_{w_i}\mathcal{L}=\frac{1}{M}{\sum}^M_{j=1} \nabla_{w_{i,j}\mathcal{L}}$，GC操作$\Phi$定義如下：

也可以將公式1轉換為矩陣形式：

$P$由單位矩陣以及單位向量形成矩陣構成，分別負責保留原值以及求均值。

Embedding of GC to SGDM/Adam

GC能夠簡單地嵌入當前的主流網路優化演算法中，如SGDM和Adam，直接使用零均值化的梯度$\Phi_{GC}(\nabla_w \mathcal{L})$進行權值的更新。

演算法1和演算法2分別展示了將GC嵌入到SGDM和Adam中，基本上不需要對原優化器演算法進行修改，僅需加入一行梯度零均值化計算即可，大約僅需0.6sec。

Properties of GC

下面從理論的角度分析GC為何能提高模型的泛化能力以及加速訓練。

Improving Generalization Performance

GC有一個很重要的優點是提高模型的泛化能力，主要得益於權值空間正則化和特徵值空間正則化。

Weight space regularization

首先介紹$P$的物理意義，經過推算可以得到：

即$P$可以看作對映矩陣，將$\nabla_W \mathcal{L}$對映到空間向量中法向量為$e$的超平面，$P\nabla_W \mathcal{L}$為對映梯度。

以SGD優化為例，權值梯度的對映能夠將權值空間約束在一個超平面或黎曼流形(Riemannian manifold)中，如圖2所示，梯度首先對映到$e^T(w-wt)=0$的超平面中，然後跟隨對映梯度$-P\nabla_{w^{t}\mathcal{L}$的方向進行更新。從$e}T(w-w^{t)=0$可以得到$e}Tw^{t+1}=eTw^t=\cdots=eTw^0$，目標函式實際變為：

這是一個權值空間$w$的約束優化問題，正則化$w$的解空間，降低了過擬合的可能性(過擬合通常是學習了複雜的權值來適應訓練資料)，能夠提升網路的泛化能力，特別是當訓練樣本較少的情況下。
WS對權值進行$e^{Tw=0$的約束，當初始權值不滿足約束時，會直接修改權值來滿足約束條件。假設進行fine-tune訓練，WS則會完全丟棄預訓練模型的優勢，而GC可以適應任何初始權值$e}T(w^0-w0)=0$。

Output feature space regularization

以SGD優化方法為例，權值更新$w^{t+1}=wt-\alpha^{tP\nabla_{w_t}\mathcal{L}$，可以推導得到$w}t=w^0-P{\sum}{t-1}_{i=0}\alpha^{{(i)}\nabla_{w}{(i)}}\mathcal{L}$。對於任何輸入特徵向量$x$，有以下定理：

相關證明可以看原文附錄，定理4.1表明輸入特徵的常量變化會造成輸出的變化，而輸出的變化量僅與標量$\gamma$和$1^Tw0$相關，與當前權值$w^{t$無關。$\gamma1}Tw^{0$為初始化權值向量縮放後的均值，假設$w}0$接近0，則輸入特徵值的常量變化將幾乎不會改變輸出特徵值，意味著輸出特徵空間對訓練樣本的變化更魯棒。

對ResNet50的不同初始權值進行視覺化，可以看到權值都非常小(小於$e^{-7}$)，這說明如果使用GC來訓練，輸出特徵不會對輸入特徵的變化過於敏感。這個屬性正則化輸出特徵空間，並且提升網路訓練的泛化能力。

Accelerating Training Process

Optimization landscape smoothing

前面提到BN和WS都間接地對權值梯度進行約束，使損失函式滿足Lipschitz屬性，$||\nabla_w\mathcal{L}||_2$和$||\nabla^2_w\mathcal{L}||_2$($w$的Hessian矩陣)都有上界。GC直接對梯度進行約束，也有類似於BN和WS的屬性，對比原損失函式滿足以下定理：