最近看深度學習二階優化方面的工作，順便做一些整理。以下內容翻譯自
"Xiaoxin et al., He Large-Scale Deep Learning Optimizations: A Comprehensive Survey, 2021"(https://arxiv.org/abs/2111.00856)
深度學習二階優化算法進展
深度學習的優化在理論和經驗上主要由一階梯度方法主導，二階優化方法由於涉及二階求導或數據的二階統計遠沒有被流行起來，盡管其有着很強的理論性質，主要是受阻於其計算量、內存消耗和通信成本。
1. 二階優化基礎
許多機器學習問題都可以簡單描述為最小化關於變量$x\in \mathbb{R}^d$的損失函數：
\begin{equation}
\min_{x\in \mathbb{R}^d} F(x)
\end{equation}
當訓練一個神經網絡的權重時，我們嘗試盡可能降低誤差曲面。在大多數情況下，我們常用SGD來更新參數向量去解決該優化問題：
\begin{equation}
x_{t+1}=x_t-\eta_t g_t
\end{equation}
在實踐中，另外一種非常流行的算法是自適應優化算法（如AdaGrad，Adadelta、RMSProp、Adam等）。有一些基本算法對參數向量中每個元素進行更新。每個元素有其自己的步長，用過去的梯度進行自適應更新。
\begin{equation}
x_{t+1}^{(i)}=x_t^{(i)}-\eta_t^{(i)} g_t^{(i)}, i=1,\cdots,d
\end{equation}
這些方法也有一些動量變形有一些不同的更新的規則。 潛在的且非常強大的一類算法被稱為預條件(precondition)算法，其用一些矩陣稱為預條件子 (preconditioner)在取一個步長時對梯度進行變換。一般來說，二階優化的思想是通過一個局部優化對目標函數進行建模
\begin{equation}
f(x+\delta) \approx M(x) = f(x) + \nabla f(x)^\top \delta + \frac{1}{2}\delta^\top B(x) \delta
\end{equation}
這里$B$是對稱預條件子(preconditioner)，$\delta$ 是參數的改變量。在牛頓法里面，$B=H$, 或 $B=H+\lambda I$. 由於關於$\delta$的$M(x_t)$得到
\begin{equation}
\delta^*={\rm{arg}}\min_\delta M(x_t) = - B^{-1}\nabla f(x),
\end{equation}
然后應用這個更新
\begin{equation}
x_{t+1} = x_t + \delta^*.
\end{equation}
這類算法包括AdaGrad, Adam, 其中預條件子是對角的。但是，如果這些度量不是對角的而是全矩陣預條件子，這些算法將更加強大，例如，完全AdaGrad、Natural Gradient以及如牛頓方法、擬牛頓法等。 在優化中預條件方法通常可以在許多場景下得到更快的收斂速度或更好的``條件數"。但是它也有明顯的問題：假設參數量為$n$, 我們需要：1）最少二次方的空間復雜度$\Omega(n^2)$來存儲預條件子。2）$n^3$的時間復雜度計算預條件子的逆以應用於梯度向量。 一般來說， 這些二階方法並不實際，因為其計算成本和內存對於深度學習的設置中是難以承受的。可替換的預條件子是對角化近似或用SGD。
最近，二階方法取得了顯著進展，通過探索全矩陣和對角矩陣之間的折中。這些方法通常以模塊化的方式對梯度進行預條件，模塊化可以實現與全矩陣方式的相同或近似的效果，而且可以在存儲和訓練時間上實現與對角情況相同。受自然梯度方法的思想啟發，Martens等人提出了利用Fisher矩陣的Kronecker因子近似作為預條件矩陣，其可以應用到多層感知器（MLP）,之后又被擴展到其他架構，如卷積神經網絡和循環神經網絡中。之后，基於Hessian和擬牛頓方法的Kronecker因子近似也被提出。
2. 牛頓法
一個函數的梯度被定義為偏導數的向量
\begin{equation}
g_t = \nabla f(x) = [\frac{\partial f}{x_1},\ldots, \frac{\partial f}{x_n} ]
\end{equation}
這意味着我們能假設神經網絡的誤差曲面局部表現類似一個圓。我們將忽略曲面的所有曲率，這可能會導致我們的訓練會變得非常緩慢。 為了修正該問題， 我們可以用來自函數的二階導數的信息。牛頓法的思想是應用線性變換將橢圓變成圓。如果我們應用這個變換到梯度向量，就好像我們在圓誤差曲面的下坡。形式上，牛頓法用Hessian矩陣作為預條件子，
\begin{equation}
H = [H_{ij}],\quad H_{ij} = \frac{\partial^2 f}{\partial x_i\partial x_j}
\end{equation}
該Hessian是參數的函數，我們需要取其逆並乘以梯度，然后朝該方向走一定距離
\begin{equation}
\Delta = -\eta H(x)^{-1} \nabla f(x)
\end{equation}
如果其是一個真正的二次曲面且選擇正確的學習率，我們將可以通過一步到達曲面的最小點。但是，該單步會涉及一些復雜內容，如Hessian矩陣的求逆。 假設我們神經網絡僅僅有1M的參數，該Hessian矩陣將有1T(1M$\times$1M)的項,其求逆是完全做不到的。
曲率矩陣. 曲率矩陣中每個元素給出的是當我們向其他方向移動時，一個方向的梯度如何變換的。曲率矩陣中的非對角線項對應於誤差面中的“扭轉”。一個扭轉意味着當你朝着一個方向移動時，梯度在另外一個方向的變化。 如果我們有一個漂亮的圓形曲面，非對角線的所有元素將變成0。當我們朝着一個方向移動，梯度在其他方向上將不會改變。 但是當我們有一個橢圓的誤差面，如果我們朝着一個放移動，其他方向的梯度將會改變。這實際上是梯度方向的問題所在。 當梯度方向更新了權重，與此同時，其也更新了其他權重，引起第一個權重的梯度變化。這意味着當我們更新權重我們可能會導致事情變得更加糟糕。由於所有其他權重的變化，梯度實際上可能有反向符號。所以我們得到的權重越多，我們需要改變他們就越謹慎，因為所有其他權重的同時變化可能會改變權重的梯度。
如何避免求解大矩陣的逆. 曲率的密集型計算限制了二階優化方法在深度學習中的應用。 為了解決該問題，有許多方法被提出。 一種十分流行的工作是采用對角近似，如AdaGrad、RMSProp、Adadelta等。但是這些對角元素僅僅由很少比例的交互所組成，其忽略了曲率矩陣中大多數交互項。實驗表明這種做法在實踐中相對於調優的SGD是基本沒有改善的。這些對角線方法的好處似乎主要在於易於選擇學習率，但可能不會提供除此之外的任何基本好處。我們可以做的另外一種做法是用低秩矩陣來近似曲率矩陣，捕捉它的主要方面。這里代表性的方法是Limited-Memory BFGS(L-BFGS)。
最近，許多關注主要在全矩陣和對角矩陣的折中，尋求兩種極端之間的平衡。一些研究者尋找方法其與全矩陣情況相同或相似效果的，且在實際應用中能夠與對角矩陣在存儲和運行時間相似。 一些最近的近似全矩陣的方法有K-FAC、Shampoo等。其他方法還有自動植入Hessian操作算子，如Hessian-Free方法和信賴域方法。
3. Hessian-Free方法
Hessian-Free方法是一個擬牛頓方法，其用無低秩近似。稱為“Free”的原因是其沒有顯式的計算預條件子B取而代之的是二次模型$M(\delta)$的近似最小化。該Hessian-Free方法主要受兩個觀察啟發。第一個是計算對於任意向量$v$的矩陣向量乘法$Hv$是相對容易的，例如利用有限差去近似極限。
\begin{equation}
Hv = \lim_{\epsilon\longrightarrow 0} \frac{\nabla f(x+\epsilon v)-\nabla f(x) }{\epsilon}
\end{equation}
第二個觀測是線性共軛梯度僅用矩陣向量乘法最小化正定二次成本函數，其相對容易計算的。 共軛梯度沒有像牛頓方法達通過一步到最小點。其通過梯度方向開始，在這個方向達到最小值，可能需要重新評估梯度或重新評估誤差，以找到這個方向的最小值。一旦做完，共軛梯度方法找另外一個方向，按照第二個方向達到最小點。這項技術選擇第二個方向時不會打亂它在第一個方向上已經做過的最小化，這被稱為共軛方向。 “共軛”是指當我們走向新的方向時，我們不會改變前面方向上的梯度。共軛梯度方法對於$n$-維二次曲面可以在$n$迭代步到達全局最小點。更加重要的是，在一個典型的二次曲面上的許多小於$n$的步驟中，它將減少的誤差與在最小點基本接近。如果執行全部$n$步，其計算成本將於求整個矩陣逆相差無異。阻尼Hessian-Free優化的一個簡單變形如算法所示：
但是，Hessian-free的常見變形並不能在神經網絡模型中表現並不是很好。 一些強化技術被提出，如對Hessian的高斯-牛頓近似、共軛梯度的early stop、阻尼方法等。 最近研究表明，深度神經網絡學習比以前認為的使用簡單的方法更容易。仔細調優的動量方法足以處理深度和循環神經網絡目標中的曲率問題，而不需要復雜的二階方法。盡管SGD有或沒有動量仍然是在大多數情況下最廣泛使用和最好的方法，事實上Hessian-Free方法用比SGD的100-1000倍更少的迭代，原則上二階方法可以幫助更多，只要我們可以便宜的使用這些迭代計算。
4. Kronecker因子近似曲率算法(K-FAC)
Kronecker因子近似曲率算法(Kronecker-Factored Approximate Curvature, K-FAC)是一種自然梯度近似方法，其預條件子(Preconditioner)是一種Fisher信息矩陣的高質量近似。下面首先介紹自然梯度下降，然后解釋K-FAC的算法框架。
自然梯度下降(NGD)，是建立在信息幾何的二階優化方法。NGD利用Fisher信息矩陣作為損失函數的曲率，可以正確的獲得損失函數全貌，並且在迭代方面比一種簡單的一階方法收斂得更快。關於網絡分布$p(y|x;\theta)$的Fisher信息矩陣為
\begin{equation}
F=E[\nabla\log(p(y|x;\theta))\nabla \log(p(y|x;\theta))^\top].
\end{equation}
重要的是，$F$的一個特征是，它可以被解釋為模型的對數似然的負期望Hessian，即
\begin{equation}
F=-E_{p(y|x;\theta)}[H_{\log p(x|\theta)}]
\end{equation}
從該結論中可以看到$F$作為對數似然函數曲率度量的作用。因此，$F$的直接應用是在二階優化方法中作為$H$的直接替代。利用KL散度來衡量兩種模型的不同程度，自然梯度下降NGD的更新規則為
\begin{equation}
\theta_{t+1} = \theta_t - \eta_t F^{-1} \nabla f(\theta_t).
\end{equation}
K-FAC近似。 K-FAC近似Fisher信息矩陣可以使很容易計算逆矩陣。首先，K-FAC將$F$近似為$\hat F$, 一種對角塊矩陣，其中每個塊對應神經網絡的一層。
\begin{equation}
\hat F = {\rm{diag}} (\hat F_1, \ldots, \hat F_l,\ldots, \hat F_L).
\end{equation}
下一步，每個對角塊矩陣$\hat F_l$ 被近似成Kronecker乘積,即
\begin{equation}
\hat F_l \approx A_{l-1} \otimes G_l.
\end{equation}
這被稱為Kronecker的因式分解，$G_l, A_{l-1}$也稱為Kronecker因子，分別表示第$l$層輸出的梯度和第$l-1$層的激活輸出。根據Kronecker矩陣乘積的重要性質$(A \otimes B)^{-1}=A^{-1}\otimes B^{-1}$, $\hat F_l$的逆可以通過下式子進行計算
\begin{equation}
\hat F_l^{-1} = A_{l-1}^{-1}\otimes G_l^{-1}.
\end{equation}
第$l$層的$w_l$最終迭代步如下
\begin{equation}
w_l^{(t+1)} =w_l^{(t)} - \eta^{(t)}G_l^{-1}\nabla L_l^{(t)}A_{l-1}^{-1}.
\end{equation}
在多數實現中，Tikhonov正則項被用於避免病態條件數矩陣逆，通過在$\hat F$的對角線上添加一個阻尼系數$\gamma$,
\begin{align}
V_1 &= Q_G^\top L_i(w_i^{(k)})Q_A\\
V_2 &= V_1/(v_G v_A^\top+\lambda)\\
(\hat F_l +\gamma I)^{-1} \nabla L_i(w_i^{(k)}) &= Q_G V_2 Q_A^\top.
\end{align}
在實踐中，可以通過減少這些因素及其特征分解的計算頻率來避免重要的計算和通信，但代價是引入隱藏的信息。
5. Shampoo算法
Shampoo是另外一種預條件算法，其是全矩陣AdaGrad的近似版本，其首先通過獨立處理每一層（塊對角線）來近似。它使用了小矩陣， 其 Kronecker乘積近似全AdaGrad統計。這兩種近似使Shampoo比對角線預條件具有更強的表達力，而且可以避免大規模存儲和計算。從數學上來講，所尋找的預條件子Shampoo可以寫成兩個較小的矩陣$L$和$R$的Kronecker乘積，即
\begin{equation}
\mathrm{arg}\min_{H=L\otimes R,\\L,R>0} H^{-1}\bar G_t + Tr(H).
\end{equation}
雖然她不能解決精確的優化問題，但它有一個很好的上界，
\begin{equation}
\frac{1}{r} (\sum_{t=1}^T g_t(g_t)^\top )^{1/2} \preceq (\sum_{t=1}^T G_t G_t^\top)^{1/4}\otimes (\sum_{t=1}^T G_t^\top G_t )^{1/4}
\end{equation}
左側是全AdaGrad預條件子，其上界是兩個小矩陣$L_t, R_t$的Kronecker乘積（右邊的$\otimes$的左側項$L_t$，右側項為$R_t$）。Shampoo的更新如下：
\begin{align}
L_t &= L_{t-1} + G_t G_t^\top,\\
R_t &= R_{t-1} + G_t^\top G_t,\\
W_{t+1} &= W_{t} - L_t^{-1}G_t R_t^{-1}.
\end{align}
Shampoo和K-FAC都采用了一種近似於Fisher信息矩陣的預條件方案，盡管他們在結構上相似，但在幾個重要方面卻有所不同。一個重要的不同是，Shampoo構建與層類型無關，而K-FAC十分依賴於前饋神經網絡中的反向傳播梯度的結構。相比之下，Shampoo實際上忽略了特定的模型結構，而只依賴於標准的梯度信息。最近，Anil等人以多種方式擴展了Shampoo算法，使其適用於更大范圍的深度架構。
盡管一階方法在近十年里都占據了主導地位，但最近二階方法，如K-FAC和Shampoo顯示出了一定的前景，它們降低了全矩陣二階算法的空間和運行成本，並已適用於更大范圍的深度架構。