在大规模优化问题中,完整的Hessian矩阵往往因其$O(n^2)$的存储需求和$O(n^3)$的求逆复杂度而变得不切实际。本章探讨如何通过识别和利用Hessian的特殊结构来设计高效的二阶优化算法。我们将深入分析四种主要的结构化方法:Kronecker因子分解、块对角近似、低秩加对角结构,以及稀疏模式的自动发现。这些方法不仅在理论上优雅,更在深度学习、推荐系统等现代应用中展现出卓越的实用价值。
考虑一个具有分层结构的神经网络,其中第$l$层的权重矩阵为$\mathbf{W}_l \in \mathbb{R}^{m \times n}$。在许多情况下,该层的Fisher信息矩阵(作为Hessian的近似)可以表示为:
\[\mathbf{F}_l = \mathbb{E}[\text{vec}(\nabla_{\mathbf{W}_l} \mathcal{L}) \text{vec}(\nabla_{\mathbf{W}_l} \mathcal{L})^T]\]K-FAC的核心洞察是假设这个Fisher矩阵可以近似为两个较小矩阵的Kronecker积:
\[\mathbf{F}_l \approx \mathbf{A}_l \otimes \mathbf{B}_l\]其中$\mathbf{A}_l \in \mathbb{R}^{m \times m}$捕获输出之间的相关性,$\mathbf{B}_l \in \mathbb{R}^{n \times n}$捕获输入之间的相关性。
数学推导的深层含义:
| 从信息几何的角度,Fisher信息矩阵定义了参数空间的黎曼度量。对于指数族分布$p(\mathbf{y} | \mathbf{x}, \boldsymbol{\theta}) = h(\mathbf{y})\exp(\boldsymbol{\eta}(\boldsymbol{\theta})^T\mathbf{T}(\mathbf{y}) - A(\boldsymbol{\eta}))$,Fisher矩阵等于对数配分函数的Hessian: |
这揭示了Fisher矩阵与曲率的本质联系。在深度网络中,层级结构诱导了特殊的几何性质:
局部因子化假设:对于全连接层,若记$\mathbf{a}{l-1}$为输入激活,$\mathbf{g}_l = \nabla{\mathbf{z}_l}\mathcal{L}$为预激活梯度,则: \(\nabla_{\mathbf{W}_l}\mathcal{L} = \mathbf{g}_l \mathbf{a}_{l-1}^T\)
这自然导出了Kronecker积结构,因为: \(\text{vec}(\mathbf{g}_l \mathbf{a}_{l-1}^T) = \mathbf{a}_{l-1} \otimes \mathbf{g}_l\)
统计独立性的几何解释:K-FAC假设$\mathbf{a}_{l-1}$与$\mathbf{g}_l$统计独立,这在几何上意味着输入流形与输出流形正交。虽然这在深度网络中并不严格成立,但经验表明这种近似在优化景观的大部分区域足够准确。
分层度量的复合:整个网络的Fisher矩阵具有块对角结构(忽略层间相关性): \(\mathbf{F} = \text{diag}(\mathbf{F}_1, \mathbf{F}_2, \ldots, \mathbf{F}_L)\)
K-FAC进一步将每个块分解为Kronecker积,实现了计算效率与表达能力的平衡。
关键性质:
高级性质与理论扩展:
矩阵范数的保持:对于任意相容的矩阵范数$|\cdot|$: \(\|\mathbf{A} \otimes \mathbf{B}\| = \|\mathbf{A}\| \cdot \|\mathbf{B}\|\)
这保证了K-FAC近似不会引入数值不稳定性。
条件数的乘积性: \(\kappa(\mathbf{A} \otimes \mathbf{B}) = \kappa(\mathbf{A}) \cdot \kappa(\mathbf{B})\)
这意味着需要同时控制两个因子的条件数以保证数值稳定性。
梯度流的分解:在连续时间极限下,K-FAC对应的梯度流可以分解为: \(\frac{d\mathbf{W}}{dt} = -(\mathbf{A}^{-1} \otimes \mathbf{B}^{-1})\text{vec}(\nabla_{\mathbf{W}}\mathcal{L})\)
这可以解释为在两个不同时间尺度上的耦合动力学。
信息论解释:Kronecker积分解最小化了以下信息论准则: \(\min_{\mathbf{A},\mathbf{B}} D_{KL}(\mathcal{N}(0, \mathbf{F}) \| \mathcal{N}(0, \mathbf{A} \otimes \mathbf{B}))\)
其中$D_{KL}$是Kullback-Leibler散度。这提供了K-FAC的信息几何基础。
与张量分解的联系:K-FAC可以看作是对四阶张量$\mathcal{F}{ijkl} = \mathbf{F}{(i-1)m+j,(k-1)m+l}$的特殊Tucker分解: \(\mathcal{F} \approx \mathcal{G} \times_1 \mathbf{U}_1 \times_2 \mathbf{U}_2 \times_3 \mathbf{U}_3 \times_4 \mathbf{U}_4\)
其中核心张量$\mathcal{G}$具有特殊的对角结构。
理论依据与假设:
K-FAC近似基于以下关键假设:
这些假设虽然在实践中并不完全成立,但经验表明K-FAC仍能提供高质量的曲率近似。近期研究表明,通过引入高阶修正项可以放松这些假设:
\[\mathbf{F}_l = \mathbf{A}_l \otimes \mathbf{B}_l + \mathbf{R}_l\]其中$\mathbf{R}_l$是捕获非Kronecker结构的残差项。
假设的数学刻画与放松:
独立性假设的量化:定义相关性度量: \(\rho(\mathbf{a}, \mathbf{g}) = \frac{\|\mathbb{E}[\mathbf{a}\mathbf{a}^T \otimes \mathbf{g}\mathbf{g}^T] - \mathbb{E}[\mathbf{a}\mathbf{a}^T] \otimes \mathbb{E}[\mathbf{g}\mathbf{g}^T]\|_F}{\|\mathbb{E}[\mathbf{a}\mathbf{a}^T]\|_F \cdot \|\mathbb{E}[\mathbf{g}\mathbf{g}^T]\|_F}\)
当$\rho \approx 0$时,K-FAC近似质量高。实践中,$\rho$通常在0.1-0.3范围内。
层间相关性的处理:通过引入三因子分解捕获相邻层的相关性: \(\mathbf{F}_{l,l+1} \approx \mathbf{A}_l \otimes \mathbf{C}_{l,l+1} \otimes \mathbf{B}_{l+1}\)
其中$\mathbf{C}_{l,l+1}$捕获层间的传递结构。
非同质性的自适应处理:使用样本特定的权重: \(\mathbf{F}_l = \sum_{i=1}^B w_i \cdot \text{vec}(\nabla_{\mathbf{W}_l}\mathcal{L}_i)\text{vec}(\nabla_{\mathbf{W}_l}\mathcal{L}_i)^T\)
其中$w_i$基于梯度范数或损失值自适应调整。
高阶修正的具体形式: \(\mathbf{R}_l = \sum_{k=1}^K \sigma_k (\mathbf{u}_k \mathbf{v}_k^T) \otimes (\mathbf{p}_k \mathbf{q}_k^T)\)
这是一个低秩修正,可以通过随机SVD高效计算。
理论保证的最新进展:
定理(K-FAC近似误差界):在温和的正则性条件下,存在常数$C$使得: \(\|\mathbf{F}_l - \mathbf{A}_l \otimes \mathbf{B}_l\|_F \leq C \cdot \sqrt{mn} \cdot \mathbb{E}[\|\mathbf{a}_{l-1}\|^2\|\mathbf{g}_l\|^2] \cdot \rho(\mathbf{a}_{l-1}, \mathbf{g}_l)\)
这个界说明了网络宽度、激活/梯度的大小以及相关性如何影响近似质量。
与Natural Gradient的深层联系:
K-FAC可视为Natural Gradient在深度网络中的高效实现。考虑参数空间的黎曼度量:
\[ds^2 = d\boldsymbol{\theta}^T \mathbf{F}(\boldsymbol{\theta}) d\boldsymbol{\theta}\]Natural Gradient沿着该度量定义的最陡下降方向更新参数:
\[\boldsymbol{\theta}_{t+1} = \boldsymbol{\theta}_t - \eta \mathbf{F}^{-1}(\boldsymbol{\theta}_t) \nabla_{\boldsymbol{\theta}} \mathcal{L}\]K-FAC通过Kronecker分解使得$\mathbf{F}^{-1}$的计算变得可行。更深入地,这种分解隐含了对网络激活流的马尔可夫假设,即信息在层间的传播具有无记忆性。
信息几何的深层视角:
参数化不变性:Natural Gradient具有参数化不变性,即对于可逆变换$\boldsymbol{\phi} = f(\boldsymbol{\theta})$: \(\tilde{\nabla}_{\boldsymbol{\phi}}\mathcal{L} = \mathbf{J}_f^{-T} \tilde{\nabla}_{\boldsymbol{\theta}}\mathcal{L}\)
其中$\tilde{\nabla}$表示自然梯度,$\mathbf{J}_f$是Jacobian矩阵。这保证了优化轨迹不依赖于参数化方式。
最小长度原理:Natural Gradient更新对应于参数空间中的测地线,最小化了以下泛函: \(\mathcal{S}[\boldsymbol{\theta}(t)] = \int_0^1 \sqrt{\dot{\boldsymbol{\theta}}^T(t) \mathbf{F}(\boldsymbol{\theta}(t)) \dot{\boldsymbol{\theta}}(t)} dt\)
K-FAC通过Kronecker近似保持了这种几何结构的主要特征。
KL散度的二阶Taylor展开:Fisher矩阵是KL散度的局部二次近似: \(D_{KL}(p(\cdot|\boldsymbol{\theta}) \| p(\cdot|\boldsymbol{\theta} + \Delta\boldsymbol{\theta})) \approx \frac{1}{2}\Delta\boldsymbol{\theta}^T \mathbf{F}(\boldsymbol{\theta}) \Delta\boldsymbol{\theta}\)
K-FAC保持了这种近似的主要结构,同时大幅降低了计算复杂度。
与量子信息的联系:Fisher信息矩阵与量子保真度(quantum fidelity)的关系: \(F(\rho_{\boldsymbol{\theta}}, \rho_{\boldsymbol{\theta}+d\boldsymbol{\theta}}) = 1 - \frac{1}{8}d\boldsymbol{\theta}^T \mathbf{F}_{Q}(\boldsymbol{\theta}) d\boldsymbol{\theta} + O(\|d\boldsymbol{\theta}\|^3)\)
其中$\mathbf{F}_Q$是量子Fisher信息矩阵。这暗示了K-FAC在量子机器学习中的潜在应用。
K-FAC通过以下步骤近似Fisher信息矩阵:
Kronecker因子估计: \(\mathbf{A}_l = \mathbb{E}[\mathbf{g}_l \mathbf{g}_l^T], \quad \mathbf{B}_l = \mathbb{E}[\mathbf{a}_{l-1} \mathbf{a}_{l-1}^T]\)
实际实现中,通常采用运行平均来估计期望值,并定期更新逆矩阵以平衡计算成本。
算法的精确描述与变体:
算法 3.1:K-FAC with Adaptive Damping
输入:学习率η,动量参数β₁,阻尼参数λ₀,更新频率T_inv
初始化:A_l = εI, B_l = εI for all layers l
for t = 1, 2, ... do
// 前向传播,收集激活
for l = 1 to L do
a_{l-1} = 层l的输入激活
存储 a_{l-1}
// 反向传播,收集梯度
for l = L to 1 do
g_l = 层l的预激活梯度
∇W_l = g_l ⊗ a_{l-1}^T
// 更新Kronecker因子(指数移动平均)
for l = 1 to L do
A_l ← β₁ A_l + (1-β₁) E[g_l g_l^T]
B_l ← β₁ B_l + (1-β₁) E[a_{l-1} a_{l-1}^T]
// 周期性更新逆矩阵
if t mod T_inv == 0 then
for l = 1 to L do
// 自适应阻尼
λ_l = λ₀ * (tr(A_l)/m + tr(B_l)/n) / 2
Ã_l = A_l + λ_l I
B̃_l = B_l + λ_l I
// 计算逆或伪逆
A_l^{-1} = pinv(Ã_l)
B_l^{-1} = pinv(B̃_l)
// 应用预条件更新
for l = 1 to L do
vec(ΔW_l) = (A_l^{-1} ⊗ B_l^{-1}) vec(∇W_l)
W_l ← W_l - η ΔW_l
关键实现细节:
批处理中的高效计算: 对于批量大小B,激活和梯度的协方差估计: \(\mathbf{A}_l = \frac{1}{B}\sum_{i=1}^B \mathbf{g}_l^{(i)} (\mathbf{g}_l^{(i)})^T\) \(\mathbf{B}_l = \frac{1}{B}\sum_{i=1}^B \mathbf{a}_{l-1}^{(i)} (\mathbf{a}_{l-1}^{(i)})^T\)
可以通过矩阵乘法高效实现: \(\mathbf{A}_l = \frac{1}{B}\mathbf{G}_l\mathbf{G}_l^T, \quad \mathbf{B}_l = \frac{1}{B}\mathbf{A}_{l-1}\mathbf{A}_{l-1}^T\)
其中$\mathbf{G}_l = [\mathbf{g}_l^{(1)}, \ldots, \mathbf{g}_l^{(B)}]$。
高级实现技巧:
动量机制的整合: K-FAC可以与动量方法结合,形成预条件动量更新: \(\mathbf{m}_{t+1} = \beta_1 \mathbf{m}_t + (1-\beta_1) \mathbf{g}_t\) \(\mathbf{v}_{t+1} = (\mathbf{A}_t \otimes \mathbf{B}_t)^{-1} \mathbf{m}_{t+1}\)
这种结合既保留了动量的加速效果,又利用了二阶信息的方向校正。
动量的不同组合方式:
研究表明,PM在非凸优化中更稳定,而AM在凸优化中收敛更快。
层重要性的量化指标:
自适应调度算法:
每隔T步:
1. 计算所有层的重要性指标
2. 选择top-k层进行因子更新
3. 其余层使用过时的因子或对角近似
数值稳定性保证:
Tikhonov正则化: \(\tilde{\mathbf{A}}_l = \mathbf{A}_l + \lambda_A \mathbf{I}, \quad \tilde{\mathbf{B}}_l = \mathbf{B}_l + \lambda_B \mathbf{I}\)
其中$\lambda_A$和$\lambda_B$的选择策略:
谱截断: 对于病态矩阵,可以使用谱截断: \(\mathbf{A} = \mathbf{U}\mathbf{\Lambda}\mathbf{U}^T \rightarrow \tilde{\mathbf{A}}^{-1} = \mathbf{U}\text{diag}(\min(\lambda_i^{-1}, \tau))\mathbf{U}^T\)
高级正则化技术:
自适应谱正则化:基于有效秩的动态调整 \(\lambda(t) = \lambda_0 \cdot \exp\left(-\beta \cdot \frac{r_{\text{eff}}(t)}{n}\right)\)
其中$r_{\text{eff}} = (\text{tr}(\mathbf{A}))^2 / \text{tr}(\mathbf{A}^2)$是有效秩。
分层正则化:不同层使用不同的正则化强度 \(\lambda_l = \lambda_{\text{base}} \cdot \left(\frac{\text{width}_l}{\text{width}_{\text{max}}}\right)^\alpha\)
鲁棒伪逆:使用Moore-Penrose伪逆处理奇异情况 \(\mathbf{A}^{\dagger} = \lim_{\epsilon \to 0^+} (\mathbf{A}^T\mathbf{A} + \epsilon\mathbf{I})^{-1}\mathbf{A}^T\)
内存高效的实现:
低精度存储:Kronecker因子可以用FP16存储,计算时转换为FP32
稀疏化策略:对于稀疏激活(如ReLU后),利用稀疏矩阵格式
共享因子:相似层(如ResNet中的重复块)可以共享Kronecker因子
内存池化技术:
class KroneckerFactorPool:
def __init__(self, max_memory):
self.pool = {}
self.usage_count = {}
self.memory_limit = max_memory
def get_or_create(self, layer_key, shape):
# 检查是否存在可复用的因子
for key, factor in self.pool.items():
if self.is_compatible(factor, shape):
self.usage_count[key] += 1
return factor
# 创建新因子,必要时逐出最少使用的
if self.current_memory > self.memory_limit:
self.evict_lru()
return self.create_new_factor(layer_key, shape)
量化压缩:
K-BFGS:将BFGS的拟牛顿思想与Kronecker结构结合
K-BFGS维护Kronecker结构的同时进行拟牛顿更新:
\[\mathbf{B}_{k+1} = \mathbf{B}_k + \Delta\mathbf{B}_k\]其中$\Delta\mathbf{B}_k$保持Kronecker形式。关键创新:
EKFAC (Eigenvalue-corrected K-FAC):
EKFAC通过特征值校正解决K-FAC的系统性偏差:
谱分析:计算真实Fisher与K-FAC近似的谱差异 \(\text{spec}(\mathbf{F}) \neq \{\lambda_i\mu_j : \lambda_i \in \text{spec}(\mathbf{A}), \mu_j \in \text{spec}(\mathbf{B})\}\)
分布式K-FAC:
大规模分布式训练中的K-FAC实现面临独特挑战:
Kronecker-factored Trust Region (KF-TR):
将信赖域方法与K-FAC结合:
\(\min_{\Delta\mathbf{w}} \mathcal{L}(\mathbf{w}) + \nabla\mathcal{L}^T\Delta\mathbf{w} + \frac{1}{2}\Delta\mathbf{w}^T(\mathbf{A} \otimes \mathbf{B})\Delta\mathbf{w}\) \(\text{s.t. } \|\Delta\mathbf{w}\|_{(\mathbf{A} \otimes \mathbf{B})} \leq \delta\)
关键特性:
研究方向与开放问题:
K-FAC的收敛性分析涉及两个关键方面:
近似误差界: \(\|\mathbf{F}_l - \mathbf{A}_l \otimes \mathbf{B}_l\|_F \leq \epsilon(\mathbf{F}_l)\)
其中$\epsilon$依赖于激活和梯度的统计独立性假设。
全局收敛速率: 在强凸假设下,K-FAC可达到超线性收敛: \(\|\mathbf{w}_{t+1} - \mathbf{w}^*\| \leq \rho^t \|\mathbf{w}_0 - \mathbf{w}^*\|\)
其中$\rho < 1$取决于Kronecker近似的质量。
深入的理论分析:
近似误差的精确刻画:
考虑真实Fisher矩阵的分解: \(\mathbf{F} = \mathbb{E}[\text{vec}(\mathbf{g})\text{vec}(\mathbf{g})^T]\)
K-FAC近似误差可以表示为: \(\mathbf{E} = \mathbf{F} - \mathbf{A} \otimes \mathbf{B} = \mathbb{E}[(\mathbf{g} \otimes \mathbf{a})(\mathbf{g} \otimes \mathbf{a})^T] - \mathbb{E}[\mathbf{g}\mathbf{g}^T] \otimes \mathbb{E}[\mathbf{a}\mathbf{a}^T]\)
当$\mathbf{g}$和$\mathbf{a}$统计独立时,$\mathbf{E} = 0$。实际中,可以证明: \(\|\mathbf{E}\|_F \leq C \cdot \text{Cov}(\|\mathbf{g}\|^2, \|\mathbf{a}\|^2)\)
其中$C$是与网络结构相关的常数。
非凸情况下的局部收敛性:
在非凸优化中,K-FAC在满足以下条件时具有局部线性收敛性:
条件A(局部强凸性):存在邻域$\mathcal{N}(\mathbf{w}^*)$使得: \(\lambda_{\min}(\mathbf{F}(\mathbf{w})) \geq \mu > 0, \quad \forall \mathbf{w} \in \mathcal{N}(\mathbf{w}^*)\)
条件B(Lipschitz连续性):Fisher矩阵满足: \(\|\mathbf{F}(\mathbf{w}_1) - \mathbf{F}(\mathbf{w}_2)\| \leq L_F\|\mathbf{w}_1 - \mathbf{w}_2\|\)
定理:在条件A和B下,若初始点$\mathbf{w}_0$足够接近$\mathbf{w}^*$,且K-FAC近似误差满足$|\mathbf{E}| < \mu/2$,则: \(\|\mathbf{w}_{k+1} - \mathbf{w}^*\| \leq \left(1 - \frac{\mu}{2L}\right)\|\mathbf{w}_k - \mathbf{w}^*\| + O(\|\mathbf{E}\|)\)
随机优化框架下的分析:
在随机小批量设置下,K-FAC的收敛性需要考虑:
收敛速率:使用递减步长$\eta_t = \eta_0/\sqrt{t}$时: \(\mathbb{E}[\|\mathbf{w}_T - \mathbf{w}^*\|^2] \leq O\left(\frac{1}{\sqrt{T}}\right) + O(\epsilon_{\text{approx}})\)
其中$\epsilon_{\text{approx}}$是K-FAC近似引入的偏差。
与一阶方法的比较:
加速区域:当Hessian的条件数$\kappa(\mathbf{H}) \gg 1$时,K-FAC相比SGD的加速比约为: \(\frac{T_{\text{SGD}}}{T_{\text{K-FAC}}} \approx \sqrt{\kappa(\mathbf{H})} \cdot \frac{1}{1 + \epsilon_{\text{rel}}}\)
其中$\epsilon_{\text{rel}} = |\mathbf{E}|/|\mathbf{F}|$是相对近似误差。
开放问题:
Shampoo算法将参数张量的不同模式(modes)分别预条件,实现了比K-FAC更灵活的结构化近似。对于张量$\mathcal{W} \in \mathbb{R}^{d_1 \times d_2 \times \cdots \times d_k}$,Shampoo维护$k$个预条件矩阵$\mathbf{H}_i \in \mathbb{R}^{d_i \times d_i}$。
核心思想是将高阶张量的预条件问题分解为多个低维矩阵的预条件问题,每个矩阵对应张量的一个模式。
理论基础:
张量梯度的协方差分解:
对于张量参数$\mathcal{W}$,其梯度$\mathcal{G}$的完整协方差矩阵大小为$(\prod_i d_i)^2$。Shampoo假设这个协方差可以近似为: \(\text{Cov}(\text{vec}(\mathcal{G})) \approx \mathbf{H}_1 \otimes \mathbf{H}_2 \otimes \cdots \otimes \mathbf{H}_k\)
这比K-FAC的双因子分解更一般化,允许处理高阶张量。
模式独立性假设:
Shampoo隐含假设不同模式之间的统计相关性可以通过各自的协方差矩阵充分捕获。数学上: \(\mathbb{E}[\mathcal{G} \times_i \mathcal{G}^T] \approx \mathbb{E}_i[\mathcal{G}_{(i)} \mathcal{G}_{(i)}^T]\)
其中$\mathcal{G}_{(i)}$是张量$\mathcal{G}$沿第$i$个模式的展开。
与黎曼优化的联系:
Shampoo可以解释为在乘积流形$\mathcal{M} = \mathbb{R}^{d_1} \times \mathbb{R}^{d_2} \times \cdots \times \mathbb{R}^{d_k}$上的黎曼优化,其中每个因子空间配备由$\mathbf{H}_i$诱导的度量: \(g_i(\mathbf{u}, \mathbf{v}) = \mathbf{u}^T \mathbf{H}_i \mathbf{v}\)
这提供了Shampoo的几何解释:在每个模式的自然几何下进行最陡下降。
与其他方法的区别:
计算与存储优势:
对于$k$阶张量$\mathcal{W} \in \mathbb{R}^{d_1 \times \cdots \times d_k}$:
当$d_i$相近且$k$较大时,Shampoo的优势尤为明显。
算法步骤:
梯度张量化:将梯度$\mathbf{g}$重塑为与参数相同的张量形式$\mathcal{G}$
模式展开与统计收集: 对每个模式$i$,计算: \(\mathbf{G}_i = \text{unfold}_i(\mathcal{G}), \quad \mathbf{H}_i \leftarrow \beta\mathbf{H}_i + (1-\beta)\mathbf{G}_i\mathbf{G}_i^T\)
预条件应用: \(\mathcal{P} = \mathcal{G} \times_1 \mathbf{H}_1^{-1/4} \times_2 \mathbf{H}_2^{-1/4} \cdots \times_k \mathbf{H}_k^{-1/4}\)
其中$\times_i$表示沿第$i$个模式的张量积。
参数更新: \(\mathcal{W} \leftarrow \mathcal{W} - \eta \mathcal{P}\)
高级实现细节:
矩阵幂的高效计算:
计算$\mathbf{H}_i^{-1/4}$是Shampoo的计算瓶颈。常用方法:
特征分解法: \(\mathbf{H}_i = \mathbf{U}_i\mathbf{\Lambda}_i\mathbf{U}_i^T \Rightarrow \mathbf{H}_i^{-1/4} = \mathbf{U}_i\mathbf{\Lambda}_i^{-1/4}\mathbf{U}_i^T\)
Newton-Schulz迭代: \(\mathbf{X}_{k+1} = \frac{1}{2}\mathbf{X}_k(3\mathbf{I} - \mathbf{X}_k^4\mathbf{H}_i), \quad \mathbf{X}_0 = \frac{\mathbf{I}}{\|\mathbf{H}_i\|^{1/4}}\)
收敛到$\mathbf{H}_i^{-1/4}$,适合GPU并行计算。
随机近似: 使用随机SVD或Nyström方法近似主要特征空间: \(\mathbf{H}_i^{-1/4} \approx \mathbf{V}_r\mathbf{\Lambda}_r^{-1/4}\mathbf{V}_r^T + \epsilon^{-1/4}(\mathbf{I} - \mathbf{V}_r\mathbf{V}_r^T)\)
自适应预条件强度:
不同模式可能需要不同的预条件强度。引入模式特定的指数$p_i$: \(\mathcal{P} = \mathcal{G} \times_1 \mathbf{H}_1^{-p_1} \times_2 \mathbf{H}_2^{-p_2} \cdots\)
其中$p_i$可以基于:
分块策略:
对于超大维度$d_i$,可以进一步分块: \(\mathbf{H}_i = \begin{pmatrix} \mathbf{H}_{i,1} & & \\ & \ddots & \\ & & \mathbf{H}_{i,B} \end{pmatrix}\)
这牺牲了一些模式内的相关性,但大幅降低了计算成本。
动量与Shampoo的结合:
预条件动量(PM-Shampoo): \(\mathbf{m}_{t+1} = \beta_1\mathbf{m}_t + (1-\beta_1)\text{vec}(\mathcal{G}_t)\) \(\mathcal{P}_t = \text{reshape}(\mathbf{m}_{t+1}) \times_1 \mathbf{H}_{1,t}^{-1/4} \times_2 \cdots\)
后条件动量(AM-Shampoo): \(\mathcal{P}_t = \mathcal{G}_t \times_1 \mathbf{H}_{1,t}^{-1/4} \times_2 \cdots\) \(\mathbf{v}_{t+1} = \beta_1\mathbf{v}_t + (1-\beta_1)\text{vec}(\mathcal{P}_t)\)
数值稳定性技巧:
条件数控制: \(\tilde{\mathbf{H}}_i = \mathbf{H}_i + \lambda_i\mathbf{I}, \quad \lambda_i = \max(\epsilon, \alpha\cdot\text{median}(\text{diag}(\mathbf{H}_i)))\)
梯度裁剪集成: 在预条件前后都可以应用梯度裁剪: \(\tilde{\mathcal{G}} = \begin{cases} \mathcal{G}, & \|\mathcal{G}\|_F \leq c \\ c\cdot\mathcal{G}/\|\mathcal{G}\|_F, & \text{otherwise} \end{cases}\)
异常检测与恢复: 监控预条件后的更新范数,异常时回退到一阶更新: \(\mathcal{P} = \begin{cases} \mathcal{P}_{\text{Shampoo}}, & \|\mathcal{P}_{\text{Shampoo}}\|_F \leq \tau\|\mathcal{G}\|_F \\ \mathcal{G}, & \text{otherwise} \end{cases}\)
存储复杂度:
计算复杂度(每次迭代):
通过使用矩阵函数的低秩近似(如Lanczos方法计算$\mathbf{H}^{-1/4}$),可进一步降低计算成本。
在大规模分布式训练中,Shampoo的实现面临独特挑战:
研究方向:
对于形如$\mathbf{H} = \mathbf{D} + \mathbf{U}\mathbf{V}^T$的矩阵,其中$\mathbf{D}$是对角阵,$\mathbf{U}, \mathbf{V} \in \mathbb{R}^{n \times r}$且$r \ll n$,Woodbury恒等式给出:
\[\mathbf{H}^{-1} = \mathbf{D}^{-1} - \mathbf{D}^{-1}\mathbf{U}(\mathbf{I}_r + \mathbf{V}^T\mathbf{D}^{-1}\mathbf{U})^{-1}\mathbf{V}^T\mathbf{D}^{-1}\]这将$O(n^3)$的求逆降至$O(nr^2 + r^3)$。
增量秩更新:
内存管理:
动态调整低秩近似的秩$r$是实现计算效率与近似精度平衡的关键:
谱分析方法: 监控近似误差: \(\epsilon_r = \frac{\sum_{i=r+1}^n \sigma_i^2}{\sum_{i=1}^n \sigma_i^2}\)
当$\epsilon_r < \tau$时,认为秩$r$足够。
基于梯度的方法: 使用梯度信息自适应调整: \(r_{t+1} = \begin{cases} r_t + 1, & \text{if } \|\mathbf{g}_t - \mathbf{H}_t^{-1}\mathbf{g}_t\|/\|\mathbf{g}_t\| > \delta \\ r_t - 1, & \text{if } \|\mathbf{g}_t - \mathbf{H}_t^{-1}\mathbf{g}_t\|/\|\mathbf{g}_t\| < \delta/2 \\ r_t, & \text{otherwise} \end{cases}\)
低秩加对角结构与矩阵的谱分解有深刻联系。考虑谱分解: \(\mathbf{H} = \sum_{i=1}^n \lambda_i \mathbf{u}_i \mathbf{u}_i^T\)
低秩加对角近似可视为: \(\mathbf{H} \approx \sum_{i=1}^r \lambda_i \mathbf{u}_i \mathbf{u}_i^T + \bar{\lambda} \sum_{i=r+1}^n \mathbf{u}_i \mathbf{u}_i^T\)
其中$\bar{\lambda}$是尾部特征值的代表值(如平均值或中位数)。
研究方向:
神经网络的计算图结构天然诱导出Hessian的稀疏模式。关键观察是:如果参数$w_i$和$w_j$在计算图中没有共同影响任何输出,则$\frac{\partial^2 \mathcal{L}}{\partial w_i \partial w_j} = 0$。
稀疏模式发现算法:
训练过程中Hessian的稀疏模式可能变化,需要动态适应机制:
自适应阈值策略: \(\mathbf{H}_{ij} = \begin{cases} \mathbf{H}_{ij}, & \text{if } |\mathbf{H}_{ij}| > \tau_t \\ 0, & \text{otherwise} \end{cases}\)
其中$\tau_t$根据以下准则动态调整:
模式迁移学习:
对于发现的稀疏模式,选择合适的因子分解至关重要:
Cholesky分解的稀疏变体:
不完全LU分解(ILU):
多级方法:
卷积网络:
Transformer架构:
图神经网络:
开放问题:
本章深入探讨了四种主要的结构化二阶优化方法:
Kronecker因子分解(K-FAC):通过假设Fisher信息矩阵可分解为Kronecker积,将存储和计算复杂度大幅降低。关键公式:$\mathbf{F}_l \approx \mathbf{A}_l \otimes \mathbf{B}_l$
Block对角近似(Shampoo):将高维张量参数的预条件问题分解为多个低维问题,实现了灵活的结构化近似。核心操作:$\mathcal{P} = \mathcal{G} \times_1 \mathbf{H}_1^{-1/4} \times_2 \mathbf{H}_2^{-1/4} \cdots$
低秩加对角结构:利用Woodbury恒等式高效处理$\mathbf{H} = \mathbf{D} + \mathbf{U}\mathbf{V}^T$形式的矩阵,实现内存和计算的双重优化。
稀疏Hessian模式:通过计算图分析自动发现稀疏结构,并使用专门的稀疏线性代数技术加速计算。
这些方法的共同特点是在保持二阶信息主要特征的同时,大幅降低计算和存储开销,使得二阶优化在大规模问题中变得可行。
证明对于Kronecker积$\mathbf{A} \otimes \mathbf{B}$,有以下性质: (a) $(\mathbf{A} \otimes \mathbf{B})^T = \mathbf{A}^T \otimes \mathbf{B}^T$ (b) $(\mathbf{A} \otimes \mathbf{B})(\mathbf{C} \otimes \mathbf{D}) = \mathbf{AC} \otimes \mathbf{BD}$(当维度匹配时)
提示:使用Kronecker积的定义,逐元素验证。
对于Shampoo算法中的张量$\mathcal{W} \in \mathbb{R}^{d_1 \times d_2 \times d_3}$,若每个模式的预条件矩阵为$\mathbf{H}_i \in \mathbb{R}^{d_i \times d_i}$,计算: (a) 存储所有预条件矩阵所需的内存 (b) 与存储完整的$\text{vec}(\mathcal{W})$的协方差矩阵相比,内存节省了多少?
提示:完整协方差矩阵的大小为$(d_1d_2d_3)^2$。
使用Woodbury恒等式计算以下矩阵的逆: \(\mathbf{H} = \mathbf{I} + \mathbf{u}\mathbf{v}^T\) 其中$\mathbf{u}, \mathbf{v} \in \mathbb{R}^n$。
提示:这是秩一更新的特殊情况。
考虑K-FAC近似$\mathbf{F} \approx \mathbf{A} \otimes \mathbf{B}$,其中真实的Fisher矩阵为$\mathbf{F} = \mathbb{E}[\text{vec}(\mathbf{g})\text{vec}(\mathbf{g})^T]$,$\mathbf{g} = \nabla_{\mathbf{W}}\mathcal{L}$。推导K-FAC近似的最优性条件,即找到使得$|\mathbf{F} - \mathbf{A} \otimes \mathbf{B}|_F$最小的$\mathbf{A}$和$\mathbf{B}$。
提示:考虑将问题转化为矩阵形式,利用迹的性质。
设计一个自适应算法,根据当前梯度信息动态选择使用K-FAC、Shampoo还是低秩加对角近似。给出选择准则和切换策略。
提示:考虑计算效率、内存限制和近似质量的权衡。
证明当Hessian具有块对角结构时,Shampoo算法给出的更新方向与精确Newton方向一致。讨论这一性质的实际意义。
提示:考虑块对角矩阵的Kronecker积表示。
考虑将结构化二阶方法应用于联邦学习场景,其中数据分布在多个客户端。设计一个通信高效的协议,使得客户端可以协作计算结构化的Hessian近似,同时保护隐私。
提示:考虑差分隐私、安全聚合和梯度压缩技术。
探讨如何将结构化二阶方法与现代硬件加速器(如TPU、GPU)的特性相结合。特别关注:
提示:考虑硬件的并行特性和内存带宽限制。
问题:盲目应用K-FAC到所有层
问题:更新频率选择不当
问题:数值不稳定
问题:张量模式顺序混淆
问题:内存爆炸
问题:矩阵幂运算的数值误差
问题:秩选择过于激进
问题:忽视更新的时序相关性
问题:Woodbury公式的错误应用
问题:过度稀疏化
问题:稀疏模式频繁变化
问题:填充现象
梯度检查:
验证:g^T H^{-1} g > 0 (正定性)
监控:||Hp - g|| / ||g|| (近似质量)
条件数监控:
增量验证:
可视化诊断:
通过遵循这些最佳实践,可以有效避免常见陷阱,实现高效稳定的结构化二阶优化。记住,没有一种方法适用于所有场景,需要根据具体问题特性和计算环境做出明智选择。