曲率感知最佳化：從自然梯度到 Shampoo

為何一階步可能是一步壞棋

想像一個損失曲面像一條又長又窄的峽谷——沿著長邊平緩、沿著寬邊陡峭。梯度大多指向橫越的方向，於是固定步長要嘛沿著長邊爬行、要嘛在谷壁間彈跳。這個落差由曲率矩陣的條件數（condition number）刻畫：其最大特徵值與最小特徵值之比。當這個比值很大時，一階方法會慢到爬行，而沒有任何單一學習率對所有方向都恰當。

確切的解藥是牛頓法：用反海森矩陣（Hessian）做預條件，`w ← w − η H⁻¹ g`，它把峽谷局部變圓，使每個方向都一樣容易。對於擁有數十億參數的網路，這是幻想——`H` 是參數數目見方的方陣，大到無法成形、儲存或求逆。曲率感知最佳化，就是工程化地對 `H⁻¹ g` 做有用的近似，使其放得進記憶體、並在一小部分步驟預算內跑得完。

自然梯度：曲率來自機率，而非損失

自然梯度下降做預條件時用的不是損失的海森矩陣，而是費雪資訊矩陣（Fisher information matrix）——當你移動參數時，模型輸出分布的曲率。關鍵性質是不變性：無論你用變異數還是對數變異數來參數化一個高斯分布，自然梯度步都是同一個物理上的移動，因為費雪量的是預測分布本身的變化，並以 KL 散度作為局部的量尺。

\tilde{\nabla}\mathcal{L}=F^{-1}\nabla\mathcal{L},\qquad F=\mathbb{E}_{x\sim p_\theta}\!\left[\nabla_{\theta}\log p_\theta(x)\,\nabla_{\theta}\log p_\theta(x)^{\top}\right]

自然梯度用模型输出分布的费舍尔信息矩阵之逆来重新缩放损失梯度。

這個不變性正是自然梯度吸引人的深層理由：普通梯度下降的路徑取決於對單位的任意選擇，而自然梯度則在分布空間中沿著同一條軌跡前進。對指數族損失而言，費雪矩陣與一個半正定的曲率——稱為廣義高斯-牛頓（Generalized Gauss-Newton）矩陣——重合，這也是為何自然梯度與高斯-牛頓法在實務上常常相遇。

K-FAC：曲率的 Kronecker 因子分解

K-FAC（Kronecker 因子化近似曲率）以一個結構性假設讓費雪矩陣變得可處理：對單一層而言，曲率分解成兩個小矩陣的 Kronecker 乘積（Kronecker product）——一個由該層的輸入激活值構成，另一個由反向傳播回來的輸出梯度構成。你不必對一個龐大的區塊求逆，而是對兩個尺寸分別等於該層輸入寬度與輸出寬度的中等矩陣求逆，而且 Kronecker 乘積的逆，就是各自逆的 Kronecker 乘積。

F_\ell\approx A_{\ell-1}\otimes G_\ell\quad\Longrightarrow\quad F_\ell^{-1}\approx A_{\ell-1}^{-1}\otimes G_\ell^{-1}

K-FAC 的关键技巧:用克罗内克积分解某层的费舍尔矩阵,其逆就是两个小矩阵之逆的克罗内克积。

# K-FAC sketch for one linear layer (a = input, g = output-grad)
A = ema(A, a.T @ a)          # input covariance  (in x in)
G = ema(G, dg.T @ dg)        # grad covariance   (out x out)
# preconditioned update reshapes the gradient matrix dW:
dW_nat = inv(G) @ dW @ inv(A)   # = (A kron G)^-1 applied to vec(dW)

兩個小矩陣的逆取代一個巨大的逆；曲率統計以指數移動平均保存，並不常更新。

這些因子能容忍陳舊——你可以每隔幾百步才重新計算並重新求逆——因此 K-FAC 攤提後的額外開銷並不大。實務上，在某些視覺與強化學習工作上，它能以明顯更少的步數達到目標損失。代價是真實的工程複雜度與儲存因子的額外記憶體，這也是它從未成為通用預設值的原因。

Hessian-free：不成形矩陣的曲率

Hessian-free 最佳化走另一條路：根本不成形任何曲率矩陣。訣竅在於，即使從不把 `H` 實體化，你也能用自動微分以大約一次額外反向傳播的代價計算 Hessian-向量乘積（Hessian-vector product） `Hv`。手上有了便宜的 `Hv`，你就用共軛梯度（conjugate gradient）——一個只需把 `H` 乘上向量的迭代解法——近似求解牛頓系統 `H δ = −g`。

每一次共軛梯度迭代就是一次 Hessian-向量乘積，你在做了少數幾次後就截斷，以精確度換取速度。它之所以優雅，是因為直接從反向傳播繼承曲率資訊、不需任何因子分解假設——但現在每一個外層步驟都包含一個內層迴圈，因此每次更新的實際時間成本很高。Hessian-free 在每一步都珍貴（小型、評估昂貴的模型）時發光，在步驟便宜又充足時則黯淡。

Shampoo：依軸分解的全矩陣自適應

Adam 保有一個對角預條件子——每個參數一個純量——它忽略座標之間的所有相關性。Shampoo保有一個全矩陣預條件子，但沿著參數張量的每一個軸做分解。對一個權重矩陣，它維護一個作用在列上的左預條件子與一個作用在行上的右預條件子，各自從梯度的外積累積而成，並升到一個分數的反次方。更新時把梯度從兩側相乘，捕捉對角方法無法捕捉的列內與行內相關性。

L_t=L_{t-1}+G_tG_t^{\top},\quad R_t=R_{t-1}+G_t^{\top}G_t,\qquad W_t=W_{t-1}-\eta\,L_t^{-1/4}\,G_t\,R_t^{-1/4}

Shampoo 用每个轴上梯度协方差的逆四次方根来预条件权重矩阵:全矩阵自适应,按轴分解。

Shampoo 坐落在 Adam 與完整牛頓法之間：比對角法豐富，又遠比真實曲率便宜。分散式實作把昂貴的矩陣方根計算攤提到多個裝置與多個步驟上，而大規模的成果——包括競賽級的最佳化器基準——使它成為最可信的、帶有二階風味的實戰訓練方法之一。它也是你將在第 5 篇遇見的正交化最佳化器的概念之父。