技术摘要：Hermitian 矩阵函数双线性形式的快速稳定梯度近似

问题陈述

在科学计算和概率机器学习中，许多目标函数涉及形式为 $u^\top f(A(\theta))v$ 的双线性形式，其中 $A(\theta)$ 是一个大型参数化 Hermitian 矩阵，$f$ 是一个标量函数（例如矩阵指数、对数）。这些形式出现在高斯过程的 log-行列式估计、网络分析和量子哈密顿量学习等应用中。

虽然 Lanczos 方法通过将大型矩阵 $A$ 投影到小型三对角矩阵 $T$ 上来高效地近似这些形式的正向评估，但计算关于参数 $\theta$ 的梯度仍然具有挑战性。现有方法面临显著的权衡：

1. 反向模式微分（反向传播）：对 Lanczos 递推进行微分需要存储所有中间向量，并且通常需要昂贵的重正交化以维持数值稳定性。如果没有重正交化，在有限精度下梯度计算可能变得不稳定。
2. 增强块方法：诸如将 Arnoldi 方法应用于 $2n \times 2n$ 块矩阵（用于计算 Fréchet 导数）之类的方法会使问题规模加倍，从而增加计算和内存成本。

核心问题在于推导一种梯度近似方法，该方法在计算上高效、无需完全重正交化即可保持数值稳定，并避免通过迭代 Lanczos 过程进行反向传播的开销。

方法论

作者提出了一种仅前向的近似梯度，该方法重用了前向 Lanczos 传递。该方法依赖于仅对小型投影问题进行微分，而不是对整个递推过程进行微分。

核心推导

设二次型近似为 $\hat{\phi}(\theta; u) = \|u\|^2 e_1^\top f(T) e_1$，其中 $T = V^\top A V$ 是通过 $m$ 步 Lanczos 迭代（基为 $V$）得到的 $m \times m$ 三对角矩阵。

1. 投影灵敏度：利用矩阵函数的 Fréchet 导数计算标量目标关于小型矩阵 $T$ 的灵敏度。这产生了一个对称矩阵 $G$，使得微分形式为 $d\hat{\phi} = \text{tr}(G^\top dT)$。
2. 忽略基的变化：$T$ 的全微分为 $dT = V^\top dA V + dV^\top A V + V^\top A dV$。作者近似 $dT \approx V^\top dA V$，有效地忽略了涉及 Lanczos 基 $V$ 关于 $\theta$ 变化的项。
3. 误差分析：本文证明，忽略基变化项所引入的误差与最终的 Lanczos 残差系数 $\beta_m$ 成正比。具体而言，误差项为 $2\beta_m e_m^\top G \eta$。因此，当 Krylov 子空间不变（即 $\beta_m = 0$）时，该近似是精确的，并且随着 Lanczos 近似的收敛，误差将消失。
4. 梯度计算：最终梯度计算为：
   $$\nabla_\theta \hat{\phi} \approx \sum_{j=1}^m w_j^\top \left[ D_\theta A(\theta, v_j) \right]$$
   其中 $W = VG$。这只需要$m$ 个向量 - 雅可比积（VJPs）的矩阵 - 向量映射，重用了前向传递中使用的相同原语。

算法

提出的算法 1 包括：

1. 前向传递：运行 $m$ 步 Lanczos 以获得 $V$ 和 $T$。
2. 灵敏度构建：计算 $T$ 的特征分解，形成差商矩阵 $F$，并构建灵敏度矩阵 $G$。
3. 后向传递：计算 $W = VG$并求和收缩项$w_j^\top [D_\theta A(\theta, v_j)]$。

主要贡献

1. 仅前向梯度推导：作者推导了一种近似梯度，该梯度不需要对 Lanczos 递推进行反向传播，仅需前向传递和对小型投影矩阵的操作。
2. 理论误差界：他们证明了该近似的误差由 Lanczos 残差 $\beta_m$ 控制，这是控制前向近似精度的同一量。这意味着只要前向估计准确，梯度就是准确的。
3. 数值稳定性：实验结果表明，该方法继承了前向 Lanczos 求积的稳定性。与基于伴随的方法不同，它在无需完全重正交化的情况下仍然可靠，避免了有限精度下朴素反向传播所固有的不稳定性。
4. 效率：该方法避免了每一步 $O(nk^2)$ 的重正交化成本，以及块 Arnoldi 方法的 $2n \times 2n$ 矩阵开销，从而实现了显著的速度提升。

实验结果

该方法在三个基准测试中进行了评估：

1. 矩阵对数（高斯过程）：

   • 任务：计算 $u^\top \log(K(\theta))u$ 的梯度。
   • 比较：与基于伴随的 Lanczos 方法 [22] 进行比较。
   • 结果：一旦前向估计收敛，所提出的方法与伴随方法的精度相匹配。关键在于，在没有重正交化的情况下，伴随方法变得不稳定，而所提出的方法保持稳定且准确。

2. 网络灵敏度（矩阵指数）：

   • 任务：计算图通信能力和子图中心性（$u^\top \exp(A)v$）的灵敏度。
   • 比较：与改进的块 Arnoldi 方法 [24] 进行比较。
   • 结果：所提出的方法相比块 Arnoldi 基线实现了 1.42$\times$ 到 2.64$\times$ 的速度提升。在 Krylov 深度 $m=32$ 时，两种方法都达到了接近机器精度的梯度误差，但所提出的方法显著更快。

3. 量子哈密顿量学习：

   • 任务：学习横向场 Ising 哈密顿量的参数。
   • 比较：与密集 Fréchet 导数基线进行比较。
   • 结果：该方法在 $m=50$ 时实现了 10.4$\times$ 的速度提升，同时保持了 $1.05 \times 10^{-7}$ 的参数误差。优化轨迹紧密跟踪了密集解。

意义与主张

本文声称，这种方法为 Hermitian 矩阵函数的双线性形式微分提供了一种“快速且稳定”的替代方案。其主要意义在于将梯度稳定性与通常用于稳定 Lanczos 微分的昂贵重正交化过程解耦。通过证明梯度误差由控制前向近似的同一残差控制，作者表明，只需增加 Krylov 子空间的维度即可获得高精度的梯度，而无需承担完全重正交化或增强矩阵方法的计算惩罚。

作者指出了局限性，具体而言，该方法目前仅限于 Hermitian 矩阵和标量目标，并且 Lanczos 向量的存储可能会成为超大规模问题的瓶颈。未来的工作建议包括非 Hermitian 扩展和减少存储需求。