第3章:结构化二阶方法
在大规模优化问题中,完整的Hessian矩阵往往因其$O(n^2)$的存储需求和$O(n^3)$的求逆复杂度而变得不切实际。本章探讨如何通过识别和利用Hessian的特殊结构来设计高效的二阶优化算法。我们将深入分析四种主要的结构化方法:Kronecker因子分解、块对角近似、低秩加对角结构,以及稀疏模式的自动发现。这些方法不仅在理论上优雅,更在深度学习、推荐系统等现代应用中展现出卓越的实用价值。
3.1 Kronecker因子分解:K-FAC及其变体
3.1.1 数学基础
考虑一个具有分层结构的神经网络,其中第$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: $$\mathbf{F}(\boldsymbol{\theta}) = \nabla^2 A(\boldsymbol{\eta}(\boldsymbol{\theta}))$$ 这揭示了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积,实现了计算效率与表达能力的平衡。
关键性质:
- 存储效率:从$O(m^2n^2)$降至$O(m^2 + n^2)$
- 计算效率:利用Kronecker积的性质,$(\mathbf{A} \otimes \mathbf{B})^{-1} = \mathbf{A}^{-1} \otimes \mathbf{B}^{-1}$
- 谱分析:若$\lambda_i$和$\mu_j$分别是$\mathbf{A}$和$\mathbf{B}$的特征值,则$\lambda_i\mu_j$是$\mathbf{A} \otimes \mathbf{B}$的特征值
高级性质与理论扩展:
-
矩阵范数的保持:对于任意相容的矩阵范数$|\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近似基于以下关键假设:
- 激活与梯度的统计独立性:假设前向激活$\mathbf{a}_{l-1}$与反向梯度$\mathbf{g}_l$在数据分布上独立
- 层间独立性:不同层的参数更新可以独立处理
- 同质性假设:batch内的样本对Fisher矩阵贡献相似
这些假设虽然在实践中并不完全成立,但经验表明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在量子机器学习中的潜在应用。
3.1.2 K-FAC算法详解
K-FAC通过以下步骤近似Fisher信息矩阵:
-
激活统计收集: - 输入激活:$\mathbf{a}_{l-1} \in \mathbb{R}^n$ - 预激活梯度:$\mathbf{g}_l \in \mathbb{R}^m$
-
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]$$
-
自然梯度计算: $$\text{vec}(\Delta\mathbf{W}_l) = (\mathbf{A}_l \otimes \mathbf{B}_l)^{-1} \text{vec}(\nabla_{\mathbf{W}_l} \mathcal{L})$$ 实际实现中,通常采用运行平均来估计期望值,并定期更新逆矩阵以平衡计算成本。
算法的精确描述与变体:
算法 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:所有节点同步更新因子 - 异步K-FAC:因子更新可以延迟,减少通信开销 - 局部K-FAC:每个节点维护局部因子,周期性同步
-
内存优化技术: - 因子共享:相似层(如ResNet块)共享Kronecker因子 - 低精度存储:因子用FP16存储,计算时转FP32 - 稀疏表示:对于ReLU网络,利用激活稀疏性
高级实现技巧:
- 动量机制的整合: 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):$\mathbf{v}_t = (\mathbf{F}_t)^{-1}\mathbf{m}_t$
- 后条件动量(AM):$\mathbf{m}_t = \beta\mathbf{m}_{t-1} + (1-\beta)(\mathbf{F}_t)^{-1}\mathbf{g}_t$
- 混合动量(HM):$\mathbf{v}_t = \alpha(\mathbf{F}_t)^{-1}\mathbf{m}_t + (1-\alpha)(\mathbf{F}_t)^{-1}\mathbf{g}_t$
研究表明,PM在非凸优化中更稳定,而AM在凸优化中收敛更快。
- 分块更新策略: 为了降低计算开销,可以采用分块更新策略:
- 将层分组,每次只更新一组的Kronecker因子
- 使用循环调度或基于重要性的调度
- 关键层(如输出层附近)更频繁更新
层重要性的量化指标:
- 梯度范数比:$\gamma_l = |\nabla_{\mathbf{W}_l}\mathcal{L}|_F / \sum_j |\nabla_{\mathbf{W}_j}\mathcal{L}|_F$
- Fisher迹比:$\tau_l = \text{tr}(\mathbf{F}_l) / \sum_j \text{tr}(\mathbf{F}_j)$
- 有效秩:$r_{\text{eff},l} = (\text{tr}(\mathbf{F}_l))^2 / \text{tr}(\mathbf{F}_l^2)$
自适应调度算法: ``` 每隔T步:
- 计算所有层的重要性指标
- 选择top-k层进行因子更新
-
其余层使用过时的因子或对角近似 ```
-
数值稳定性保证:
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$的选择策略:
- 基于迹的自适应:$\lambda = \max(\epsilon, \alpha \cdot \text{tr}(\mathbf{A})/m)$
- 基于条件数的调整:监控$\kappa(\mathbf{A})$,当超过阈值时增大$\lambda$
谱截断: 对于病态矩阵,可以使用谱截断: $$\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因子
内存池化技术: ```python 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)
```
量化压缩:
- 对称量化:利用协方差矩阵的对称性
- 块量化:将矩阵分块,每块使用不同的量化级别
- 动态范围:基于历史统计调整量化范围
3.1.3 K-FAC的变体与扩展
K-BFGS:将BFGS的拟牛顿思想与Kronecker结构结合
K-BFGS维护Kronecker结构的同时进行拟牛顿更新: $$\mathbf{B}_{k+1} = \mathbf{B}_k + \Delta\mathbf{B}_k$$ 其中$\Delta\mathbf{B}_k$保持Kronecker形式。关键创新:
- 使用秩一修正的Kronecker分解:$\Delta\mathbf{B}_k = \mathbf{u}\mathbf{u}^T \otimes \mathbf{v}\mathbf{v}^T$
- 满足拟牛顿条件:$\mathbf{B}_{k+1}\mathbf{s}_k = \mathbf{y}_k$
- 保持正定性的同时降低存储需求
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})}$$
-
校正策略: - 全局缩放:$\tilde{\mathbf{F}} = \alpha(\mathbf{A} \otimes \mathbf{B})$,其中$\alpha = \frac{\text{tr}(\mathbf{F})}{\text{tr}(\mathbf{A})\text{tr}(\mathbf{B})}$ - 谱对齐:通过匹配前$k$个特征值进行校正 - 自适应阻尼:$\tilde{\mathbf{F}}_l = \mathbf{F}_l + \lambda(t)\mathbf{I}$,$\lambda(t)$随训练进程衰减
-
计算优化: - 使用随机化方法估计谱校正因子 - 周期性更新校正参数(如每100步)
分布式K-FAC:
大规模分布式训练中的K-FAC实现面临独特挑战:
-
通信模式优化: - All-reduce优化:Kronecker因子的分片聚合 - 梯度压缩:利用K-FAC的低秩结构进行梯度压缩 - 异步通信:计算与通信重叠,隐藏延迟
-
负载均衡: - 动态层分配:根据计算复杂度分配层到不同节点 - 混合并行:数据并行与模型并行的结合 - 弹性计算:支持节点动态加入/退出
-
容错机制: - Byzantine-robust聚合:使用中位数或修剪均值 - 检查点机制:定期保存Kronecker因子 - 增量恢复:故障后快速恢复训练状态
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$$ 关键特性:
- 自适应步长控制,避免过大更新
- 二阶模型的质量保证
- 与line search的对比优势
研究方向与开放问题:
-
自适应Kronecker分解: - 动态选择分解的秩 - 多重Kronecker分解:$\mathbf{F} \approx \sum_{i=1}^r \mathbf{A}_i \otimes \mathbf{B}_i$ - 与张量分解理论的联系
-
理论保证: - 非凸优化中的收敛性分析 - 近似误差对收敛速度的影响 - 最优更新频率的理论指导
-
新架构的适应: - Transformer中的多头注意力K-FAC - 图神经网络的邻域聚合K-FAC - 3D卷积和视频模型的高阶K-FAC
3.1.4 收敛性分析
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的收敛性需要考虑:
- 梯度估计的方差:$\mathbb{E}[|\nabla f_i(\mathbf{w}) - \nabla f(\mathbf{w})|^2] \leq \sigma^2$
- Kronecker因子估计的方差:$\text{Var}(\hat{\mathbf{A}}_t), \text{Var}(\hat{\mathbf{B}}_t)$
收敛速率:使用递减步长$\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}|$是相对近似误差。
开放问题:
- 非凸情况下的全局收敛保证
- 自适应Kronecker分解的理论分析
- 与其他二阶方法的统一理论框架
- 在线学习设置下的遗憾界分析
- 分布式K-FAC的一致性与收敛性权衡
3.2 Block对角近似:Shampoo算法解析
3.2.1 块结构的数学动机
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-FAC:仅处理矩阵(2阶张量),使用两个Kronecker因子
- Shampoo:可处理任意阶张量,每个模式一个因子
- 完整二阶方法:需要存储和操作完整的Hessian或Fisher矩阵
计算与存储优势:
对于$k$阶张量$\mathcal{W} \in \mathbb{R}^{d_1 \times \cdots \times d_k}$:
- 完整方法:$O((\prod_i d_i)^2)$存储,$O((\prod_i d_i)^3)$计算
- Shampoo:$O(\sum_i d_i^2)$存储,$O(\sum_i d_i^3 + k\prod_i d_i)$计算
当$d_i$相近且$k$较大时,Shampoo的优势尤为明显。
3.2.2 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$可以基于:
- 梯度信噪比:$p_i \propto |\mathbf{G}_i|_F/\text{tr}(\mathbf{H}_i)$
- 有效秩:$p_i \propto \text{tr}(\mathbf{H}_i)/|\mathbf{H}_i|_2$
- 学习率适应:通过线搜索或trust region调整
- 分块策略:
对于超大维度$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}$$
3.2.3 计算复杂度分析
存储复杂度:
- 标准全矩阵预条件:$O((\prod_i d_i)^2)$
- Shampoo:$O(\sum_i d_i^2)$
计算复杂度(每次迭代):
- 矩阵求逆:$O(\sum_i d_i^3)$
- 张量收缩:$O(k \prod_i d_i)$
通过使用矩阵函数的低秩近似(如Lanczos方法计算$\mathbf{H}^{-1/4}$),可进一步降低计算成本。
3.2.4 分布式Shampoo实现
在大规模分布式训练中,Shampoo的实现面临独特挑战:
-
预条件矩阵的分片存储: - 使用分布式矩阵分解技术 - 每个节点负责部分模式的预条件
-
通信优化: - 梯度聚合与预条件计算的流水线化 - 使用梯度压缩技术减少通信量
-
异步更新策略: - 允许不同模式的预条件矩阵异步更新 - 使用版本控制确保一致性
研究方向:
- 自适应块大小选择
- 稀疏Shampoo变体
- 与量化技术的结合
3.3 低秩加对角结构的利用
3.3.1 Woodbury恒等式的核心作用
对于形如$\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)$。
3.3.2 内存高效的实现策略
增量秩更新:
- 初始化:$\mathbf{H}_0 = \mathbf{D}_0$(对角矩阵)
- 秩一更新:$\mathbf{H}_{t+1} = \mathbf{H}_t + \mathbf{u}_t\mathbf{v}_t^T$
- 维护紧凑表示:$\mathbf{U}_t = [\mathbf{u}_1, \ldots, \mathbf{u}_t]$, $\mathbf{V}_t = [\mathbf{v}_1, \ldots, \mathbf{v}_t]$
内存管理:
- 固定秩窗口:保持最近$r$个秩一更新
- 基于重要性的选择:使用奇异值阈值
- 周期性重构:定期吸收低秩部分到对角部分
3.3.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}$$
3.3.4 与谱分析的联系
低秩加对角结构与矩阵的谱分解有深刻联系。考虑谱分解: $$\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}$是尾部特征值的代表值(如平均值或中位数)。
研究方向:
- 在线谱分解算法
- 随机化低秩近似
- 与神经网络剪枝的理论联系
3.4 稀疏Hessian模式的自动发现
3.4.1 基于图的稀疏性检测
神经网络的计算图结构天然诱导出Hessian的稀疏模式。关键观察是:如果参数$w_i$和$w_j$在计算图中没有共同影响任何输出,则$\frac{\partial^2 \mathcal{L}}{\partial w_i \partial w_j} = 0$。
稀疏模式发现算法:
-
构建影响图: - 节点:参数 - 边:若两参数共同影响某输出,则连边
-
图着色问题: - 使用图着色算法确定Hessian-vector乘积的计算组 - 最小化所需的前向-反向传播次数
-
模式压缩: - 识别重复的稀疏模式 - 使用模板匹配减少存储
3.4.2 动态模式适应
训练过程中Hessian的稀疏模式可能变化,需要动态适应机制:
自适应阈值策略: $$\mathbf{H}_{ij} = \begin{cases} \mathbf{H}_{ij}, & \text{if } |\mathbf{H}_{ij}| > \tau_t \ 0, & \text{otherwise} \end{cases}$$ 其中$\tau_t$根据以下准则动态调整:
- 保持固定稀疏度:$|\mathbf{H}|_0 = k$
- 基于误差的调整:$\tau_{t+1} = \tau_t \cdot (1 + \alpha \cdot \text{approx_error})$
模式迁移学习:
- 从预训练模型继承稀疏模式
- 逐步精炼以适应特定任务
3.4.3 稀疏因子分解技术
对于发现的稀疏模式,选择合适的因子分解至关重要:
Cholesky分解的稀疏变体:
- 使用最小度排序减少填充
- 符号分解预计算非零模式
- 数值分解的并行化
不完全LU分解(ILU):
- 控制填充级别ILU(k)
- 自适应选择丢弃阈值
- 与预条件共轭梯度法结合
多级方法:
- 构建稀疏矩阵的层次表示
- 使用代数多重网格(AMG)技术
- 在不同尺度上求解
3.4.4 在神经架构中的应用
卷积网络:
- 利用权重共享导致的块稀疏结构
- 通道间稀疏性的自动发现
- 与网络剪枝的协同优化
Transformer架构:
- 注意力机制的稀疏模式
- 多头结构的块对角性质
- 长序列的分块处理策略
图神经网络:
- 邻接矩阵诱导的稀疏性
- 消息传递的局部性
- 分层图结构的利用
开放问题:
- 稀疏模式的理论保证
- 与网络架构搜索(NAS)的结合
- 硬件友好的稀疏格式设计
本章小结
本章深入探讨了四种主要的结构化二阶优化方法:
-
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模式:通过计算图分析自动发现稀疏结构,并使用专门的稀疏线性代数技术加速计算。
这些方法的共同特点是在保持二阶信息主要特征的同时,大幅降低计算和存储开销,使得二阶优化在大规模问题中变得可行。
练习题
练习3.1(基础)
证明对于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积的定义,逐元素验证。
答案
(a) 设$\mathbf{A} \in \mathbb{R}^{m \times n}$,$\mathbf{B} \in \mathbb{R}^{p \times q}$。 Kronecker积的第$(i,j)$块为$a_{ij}\mathbf{B}$。 转置后第$(j,i)$块为$(a_{ij}\mathbf{B})^T = a_{ij}\mathbf{B}^T = a_{ji}^T\mathbf{B}^T$。 这正是$\mathbf{A}^T \otimes \mathbf{B}^T$的定义。
(b) 利用分块矩阵乘法, $(\mathbf{A} \otimes \mathbf{B})$的第$(i,j)$块为$a_{ij}\mathbf{B}$, $(\mathbf{C} \otimes \mathbf{D})$的第$(k,l)$块为$c_{kl}\mathbf{D}$。 乘积的第$(i,l)$块为$\sum_k a_{ik}c_{kl}\mathbf{BD} = (\mathbf{AC})_{il}\mathbf{BD}$。
练习3.2(基础)
对于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$。
答案
(a) Shampoo存储需求:$d_1^2 + d_2^2 + d_3^2$
(b) 完整协方差矩阵需求:$(d_1d_2d_3)^2$
内存节省比例: $$\frac{(d_1d_2d_3)^2 - (d_1^2 + d_2^2 + d_3^2)}{(d_1d_2d_3)^2} = 1 - \frac{d_1^2 + d_2^2 + d_3^2}{(d_1d_2d_3)^2}$$ 例如,当$d_1 = d_2 = d_3 = 100$时,节省比例约为$99.97\%$。
练习3.3(基础)
使用Woodbury恒等式计算以下矩阵的逆: $$\mathbf{H} = \mathbf{I} + \mathbf{u}\mathbf{v}^T$$ 其中$\mathbf{u}, \mathbf{v} \in \mathbb{R}^n$。
提示:这是秩一更新的特殊情况。
答案
应用Woodbury恒等式,其中$\mathbf{D} = \mathbf{I}$,$\mathbf{U} = \mathbf{u}$,$\mathbf{V} = \mathbf{v}$: $$\mathbf{H}^{-1} = \mathbf{I} - \frac{\mathbf{u}\mathbf{v}^T}{1 + \mathbf{v}^T\mathbf{u}}$$ 验证: $$\mathbf{H}\mathbf{H}^{-1} = (\mathbf{I} + \mathbf{u}\mathbf{v}^T)(\mathbf{I} - \frac{\mathbf{u}\mathbf{v}^T}{1 + \mathbf{v}^T\mathbf{u}})$$ $$= \mathbf{I} + \mathbf{u}\mathbf{v}^T - \frac{\mathbf{u}\mathbf{v}^T}{1 + \mathbf{v}^T\mathbf{u}} - \frac{\mathbf{u}\mathbf{v}^T\mathbf{u}\mathbf{v}^T}{1 + \mathbf{v}^T\mathbf{u}}$$ $$= \mathbf{I} + \mathbf{u}\mathbf{v}^T - \frac{\mathbf{u}\mathbf{v}^T(1 + \mathbf{v}^T\mathbf{u})}{1 + \mathbf{v}^T\mathbf{u}} = \mathbf{I}$$
练习3.4(挑战)
考虑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}$。
提示:考虑将问题转化为矩阵形式,利用迹的性质。
答案
设$\mathbf{G} = \text{mat}(\mathbf{g})$为梯度的矩阵形式,则: $$\mathbf{F} = \mathbb{E}[\text{vec}(\mathbf{G})\text{vec}(\mathbf{G})^T]$$ 最优化问题: $$\min_{\mathbf{A},\mathbf{B}} |\mathbf{F} - \mathbf{A} \otimes \mathbf{B}|_F^2$$ 利用Kronecker积的性质和迹的循环性质,可以证明最优解为: $$\mathbf{B} = \mathbb{E}[\mathbf{G}^T\mathbf{A}^{-1}\mathbf{G}] / \text{tr}(\mathbf{A}^{-1})$$ $$\mathbf{A} = \mathbb{E}[\mathbf{G}\mathbf{B}^{-1}\mathbf{G}^T] / \text{tr}(\mathbf{B}^{-1})$$ 这导出了K-FAC中使用的近似: $$\mathbf{A} \approx \mathbb{E}[\mathbf{g}_{\text{out}}\mathbf{g}_{\text{out}}^T], \quad \mathbf{B} \approx \mathbb{E}[\mathbf{a}_{\text{in}}\mathbf{a}_{\text{in}}^T]$$
练习3.5(挑战)
设计一个自适应算法,根据当前梯度信息动态选择使用K-FAC、Shampoo还是低秩加对角近似。给出选择准则和切换策略。
提示:考虑计算效率、内存限制和近似质量的权衡。
答案
自适应选择算法:
-
计算统计量: - 梯度的有效秩:$r_{\text{eff}} = \frac{(\sum_i \sigma_i)^2}{\sum_i \sigma_i^2}$ - 模式相关性:$\rho = \frac{|\mathbf{G}|_F^2}{|\mathbf{G}|_1 |\mathbf{G}^T|_1}$ - 稀疏度:$s = \frac{|\text{vec}(\mathbf{G})|_0}{n^2}$
-
选择准则: - 若$\rho > 0.8$且参数为矩阵形式:使用K-FAC - 若$r_{\text{eff}} < 0.1n$:使用低秩加对角 - 若参数为高阶张量:使用Shampoo - 若$s < 0.01$:使用稀疏方法
-
平滑切换: $$\mathbf{H}_{\text{new}} = \alpha\mathbf{H}_{\text{method1}} + (1-\alpha)\mathbf{H}_{\text{method2}}$$ 其中$\alpha$从0逐渐增加到1。
练习3.6(挑战)
证明当Hessian具有块对角结构时,Shampoo算法给出的更新方向与精确Newton方向一致。讨论这一性质的实际意义。
提示:考虑块对角矩阵的Kronecker积表示。
答案
设Hessian为块对角: $$\mathbf{H} = \begin{pmatrix} \mathbf{H}_1 & & \ & \ddots & \ & & \mathbf{H}_k \end{pmatrix}$$
若参数可表示为张量$\mathcal{W} \in \mathbb{R}^{d_1 \times \cdots \times d_k}$,且每个块$\mathbf{H}_i$对应第$i$个模式,则:
-
Shampoo预条件:$\mathcal{P} = \mathcal{G} \times_1 \mathbf{H}_1^{-1/2} \times_2 \mathbf{H}_2^{-1/2} \cdots$
-
精确Newton方向:$\mathbf{d} = \mathbf{H}^{-1}\mathbf{g}$
当Hessian恰好可分解为$\mathbf{H} = \mathbf{H}_1 \otimes \cdots \otimes \mathbf{H}_k$时,两者等价。
实际意义:
- Shampoo自然捕获了参数的张量结构
- 当不同模式间相互作用较弱时,Shampoo近似质量高
- 提供了一种隐式的结构发现机制
练习3.7(开放思考)
考虑将结构化二阶方法应用于联邦学习场景,其中数据分布在多个客户端。设计一个通信高效的协议,使得客户端可以协作计算结构化的Hessian近似,同时保护隐私。
提示:考虑差分隐私、安全聚合和梯度压缩技术。
答案
联邦结构化二阶优化协议:
-
本地计算: - 每个客户端计算本地Kronecker因子:$\mathbf{A}_i^{(c)}, \mathbf{B}_i^{(c)}$ - 添加差分隐私噪声:$\tilde{\mathbf{A}}_i^{(c)} = \mathbf{A}_i^{(c)} + \mathcal{N}(0, \sigma^2\mathbf{I})$
-
安全聚合: - 使用同态加密或秘密分享聚合因子 - 服务器计算:$\bar{\mathbf{A}}_i = \frac{1}{C}\sum_{c=1}^C \tilde{\mathbf{A}}_i^{(c)}$
-
压缩通信: - 只传输因子的主要特征向量 - 使用量化和稀疏化技术
-
隐私分析: - Kronecker因子比完整Hessian泄露更少信息 - 差分隐私预算分配:$\epsilon_{\text{total}} = \epsilon_A + \epsilon_B$
-
自适应策略: - 根据通信带宽动态调整更新频率 - 异构客户端的负载均衡
练习3.8(研究方向)
探讨如何将结构化二阶方法与现代硬件加速器(如TPU、GPU)的特性相结合。特别关注:
- 矩阵乘法单元的高效利用
- 内存层次结构的优化
- 混合精度计算的应用
提示:考虑硬件的并行特性和内存带宽限制。
答案
硬件感知的结构化二阶优化:
-
张量核心利用: - 将Kronecker积运算映射到张量核心操作 - 使用混合精度:FP16计算,FP32累加 - 批量矩阵乘法(BLAS-3)优化
-
内存优化: - 因子矩阵的分块存储,适配缓存大小 - 预条件计算与梯度计算的融合 - 使用环形缓冲区减少内存分配
-
并行策略: - 不同Kronecker因子的并行计算 - 模式并行与数据并行的结合 - 异步因子更新减少同步开销
-
数值稳定性: - 关键操作(如矩阵求逆)使用高精度 - 动态缩放防止溢出/下溢 - Kahan求和算法提高精度
-
性能模型: - 建立roof-line模型预测性能 - 自动调优选择最佳分块大小 - 运行时自适应调整计算策略
常见陷阱与错误 (Gotchas)
1. Kronecker因子分解的陷阱
问题:盲目应用K-FAC到所有层
- 错误表现:批归一化层、残差连接处的近似质量差
- 解决方案:对特殊层使用不同的近似策略,如对BN层使用对角近似
问题:更新频率选择不当
- 错误表现:频繁更新导致计算开销过大;更新太少导致过时的二阶信息
- 解决方案:使用自适应更新频率,基于梯度变化率调整
问题:数值不稳定
- 错误表现:矩阵求逆时出现奇异或病态
- 解决方案:添加适当的阻尼项$\lambda$,使用Tikhonov正则化
2. Shampoo实现的常见错误
问题:张量模式顺序混淆
- 错误表现:预条件应用到错误的维度
- 解决方案:明确定义并文档化张量维度顺序,使用命名维度
问题:内存爆炸
- 错误表现:对高维张量,预条件矩阵占用过多内存
- 解决方案:使用分组策略,将大维度分割成小块
问题:矩阵幂运算的数值误差
- 错误表现:$\mathbf{H}^{-1/4}$计算不准确
- 解决方案:使用稳定的特征分解或迭代方法(如Newton-Schulz迭代)
3. 低秩近似的误区
问题:秩选择过于激进
- 错误表现:丢失重要的曲率信息
- 解决方案:监控近似误差,保守地选择初始秩
问题:忽视更新的时序相关性
- 错误表现:新旧更新权重不当
- 解决方案:使用指数衰减权重,重视近期更新
问题:Woodbury公式的错误应用
- 错误表现:公式应用条件不满足(如$\mathbf{I} + \mathbf{V}^T\mathbf{D}^{-1}\mathbf{U}$奇异)
- 解决方案:添加正则化项或使用伪逆
4. 稀疏模式的挑战
问题:过度稀疏化
- 错误表现:重要的二阶信息被错误地置零
- 解决方案:使用自适应阈值,保留足够的非零元素
问题:稀疏模式频繁变化
- 错误表现:预处理开销超过计算节省
- 解决方案:使用滞后策略,只在模式显著变化时更新
问题:填充现象
- 错误表现:稀疏分解产生大量填充元素
- 解决方案:使用合适的重排序算法(如AMD、METIS)
5. 通用调试技巧
梯度检查:
验证:g^T H^{-1} g > 0 (正定性)
监控:||Hp - g|| / ||g|| (近似质量)
条件数监控:
- 跟踪$\kappa(\mathbf{H})$的变化
- 条件数爆炸时增加正则化
增量验证:
- 先在小模型上验证实现
- 逐步增加模型规模和复杂度
可视化诊断:
- 绘制特征值分布
- 可视化稀疏模式
- 监控收敛曲线
最佳实践检查清单
算法选择
- [ ] 分析参数结构(矩阵/张量/稀疏)
- [ ] 评估内存预算和计算资源
- [ ] 考虑硬件特性(GPU/TPU/CPU)
- [ ] 测试不同方法的pilot运行
实现要点
- [ ] 使用数值稳定的矩阵运算库
- [ ] 实现自适应正则化机制
- [ ] 添加异常检测和恢复机制
- [ ] 分离统计收集和预条件更新
性能优化
- [ ] 批量化矩阵运算
- [ ] 重用中间计算结果
- [ ] 并行化独立的计算
- [ ] 使用混合精度where appropriate
监控与调试
- [ ] 记录关键指标(近似误差、条件数等)
- [ ] 实现checkpoint和恢复机制
- [ ] 设置合理的默认超参数
- [ ] 提供详细的错误信息
扩展性考虑
- [ ] 设计模块化的接口
- [ ] 支持分布式训练
- [ ] 考虑异构计算环境
- [ ] 预留自定义结构的接口
通过遵循这些最佳实践,可以有效避免常见陷阱,实现高效稳定的结构化二阶优化。记住,没有一种方法适用于所有场景,需要根据具体问题特性和计算环境做出明智选择。