★ FEATURED ARTICLE 很多人第一次碰到“共轭梯度算法PCG”是在做岭回归或者高斯过程回归时数据量稍微一大np.linalg.solve直接卡死或者要构造一个上万阶的稠密矩阵内存先爆了。这时候才意识到原来工程上解线性方程组不是只有“直接求逆”这一条路。共轭梯度法Conjugate GradientCG和带预条件的共轭梯度法Preconditioned Conjugate GradientPCG就是解决这类“大规模、稀疏或隐式线性系统”最常用的迭代武器。这篇文章我会从CG背后的几何直觉讲起再到PCG里预条件子怎么选、怎么实现最后给出一套能在机器学习任务里直接用的Python模板并附上我在实际项目中踩过的坑。适合做机器学习算法落地、想搞懂数值优化底层原理、或者正在写大作业仿真代码的读者内容偏实践但数学上的“为什么”我也会讲透。1. 共轭梯度算法到底在解决什么问题1.1 直接法与迭代法的分水岭矩阵规模决定一切在数值计算里求解线性方程组 Ax b 有两条路线直接法和迭代法。直接法的代表是高斯消去、LU分解、Cholesky分解。它们的特点是固定步骤、高精度对任意非奇异矩阵都能在有限步内给出理论精确解。但代价也摆在那里存储和计算量都随矩阵规模呈平方到三次方增长。一个 n×n 的稠密矩阵直接存下来就需要 n² 个浮点数。当 n 到一万时是 800MB到十万时就是 80GB。这还只是存储。Cholesky分解的计算量大概是 n³/3 次浮点运算n 10⁵ 时那就是约 3.3×10¹⁴ 次。哪怕你的机器有 1 TFLOP 的算力也要跑好几天。所以在机器学习里动不动就是几十万样本、几万特征、或者核矩阵稠密到没法存这时候直接法基本没戏。迭代法思路完全不同从一个初始猜测 x₀ 出发通过不断修正得到逐渐逼近的解序列。每一步的主要成本是矩阵-向量乘法matvec一次 matvec 只要 O(n²) 甚至更少。更妙的是如果矩阵是稀疏的或者你能用某种隐式方式定义 A比如神经网络里的 Hessian-vector product那连矩阵本身都不需要显式构造出来。这就是迭代法能处理超大规模问题的根本原因。1.2 CG不是“某种优化技巧”它是解方程和优化的桥梁共轭梯度法最早是用来解对称正定SPD线性方程组的但它同时也是一个优化算法用来最小化下面这个二次型f(x) ½ xᵀAx − bᵀx const求一下梯度得到 ∇f(x) Ax − b −r其中 r b − Ax 就是残差。所以解 Ax b 和目标函数 f(x) 取极值完全等价。这个视角很重要因为一旦你把它看作优化问题就能用几何直觉去理解CG每一步在干什么后面很多实现细节也就顺理成章了。在机器学习里这种“解方程 ⇄ 优化”的等价关系经常出现。岭回归的正规方程、高斯过程后验推断、最小二乘问题本质全是 Ax b。识别出这一层你就能把CG这个工具箱用在非常多看似无关的任务上。1.3 什么情况下你该用PCG而不是其他优化器需要澄清一点PCG 适合的是“局部强凸/二次型”的问题。如果你面对的是深度学习中那种非凸的大规模损失函数标准的 CG 不会直接拿来当主优化器但二阶优化里会出现 CG 子问题。它的最佳适用场景很明确矩阵 A 对称正定至少在实际求解的子问题里是正定的你只想要一个一定精度的近似解而不是机器精度的精确解矩阵规模太大直接分解内存炸了矩阵能以 matvec 形式高效计算但你不一定有空间存它。像线性回归、高斯过程、泊松方程求解、图上的拉普拉斯系统求解都是典型的高频使用场景。2. CG的数学原理从最速下降到共轭方向2.1 最速下降为什么会“锯齿”先回顾一下最容易想到的迭代方法最速下降法每一步都沿负梯度方向走。负梯度方向是局部下降最快的方向但问题在于局部最快不等于全局最快。如果目标函数的等高线是拉长了的椭圆每一步都垂直于当前等高线结果就是在椭圆的窄轴方向来回震荡走成之字形收敛速度非常慢。这个震荡的本质原因是你在某一维度上已经把误差减得很小了下一步一换方向又把它破坏了。优化里管这个叫“步之间的冗余”。解决思路就是让每一步的方向之间满足某种“正交性”使得前面方向上的贡献不会被后面破坏。2.2 A-内积与共轭方向为什么要用A来定义“正交”对于问题 Ax bCG 选择了一种非常聪明的正交方式。它不要求搜索方向在欧氏内积下正交而是在 A 内积下正交也就是满足pᵢᵀ A pⱼ 0, i ≠ j这种方向对就称为 A-共轭。你可能会问为什么要用 A 这么奇怪的“尺子”去量正交原因是 f(x) 的等高线是椭球形而椭球的形状完全由 A 决定。在 A-内积下坐标系的“圆形”恰好就是原空间中椭球面的“天然方向”。用 A-共轭方向去搜索相当于把椭圆坐标拉成圆之后每走一步都严格不破坏前面已经做好的优化。更直观地说如果 A 是一个对角矩阵A-共轭方向就是标准的坐标系方向。高维空间里 A 不是对角的但我们可以想象成先做了一次线性变换 y A^{1/2}x在这个新坐标系里目标函数变成了圆形的这时要找一个标准正交基来依次搜索它就对应回原坐标系的 A-共轭方向。把这个逻辑理解透了CG 其实也就一句话把椭圆坐标拉圆然后老老实实地沿标准正交基走一遍。2.3 标准CG迭代公式理论归理论实际计算当然不会真去做 A^{1/2} 这样的变换太贵了。CG 的高明之处在于用一套递推公式在原来坐标系里就能生成这些 A-共轭方向并且每步给出该方向上的最优步长。标准流程如下输入: SPD矩阵A, 右端项b, 初始猜测x0, 容差tol r0 b - A*x0 p0 r0 k 0 while ||r_k|| tol: alpha_k (r_k^T * r_k) / (p_k^T * A * p_k) x_{k1} x_k alpha_k * p_k r_{k1} r_k - alpha_k * A * p_k beta_k (r_{k1}^T * r_{k1}) / (r_k^T * r_k) p_{k1} r_{k1} beta_k * p_k k k 1这里每条公式都有清晰的几何含义。alpha 是当前方向上的最优步长计算公式来自“在该方向上一维精确搜索”的最优解。r 是残差也正好是梯度 ∇f(x) 的相反数。beta 这个系数用来生成新的共轭方向可以理解成把当前残差方向中已经被旧方向张成的分量剔除掉保证新方向与之前所有方向 A-共轭。值得强调的是CG 残差之间是两两正交的欧氏内积而搜索方向之间是 A-共轭的。这两组正交性保证了CG每一步都是在 n 维空间中沿着一个对偶的“完备方向系”行走所以理论上最多 n 步就能精确收敛。2.4 收敛速度条件数是真正的主导理论上的“n步收敛”只是天花板实际迭代多少步才停下取决于 A 的特征值分布。可以证明CG 第 k 步的误差满足一个上界近似为||x_k − x*||_A ≤ 2 * ( (√κ − 1) / (√κ 1) )^k * ||x0 − x*||_A这里的 κ λ_max / λ_min 就是矩阵 A 的条件数。如果 κ 1那一步就收敛。κ 100 时(√κ−1)/(√κ1) ≈ 0.818大概每 4 步误差下降一个量级κ 10⁶ 时这个比值约 0.998收敛慢得令人发指。所以CG的实际效率完全被条件数掐住了。这也是专门要加预条件器、发展出 PCG 的原因。浮点数运算本身很快但迭代步数一多谁都扛不住。机器学习里很多核矩阵条件数轻松到 10⁸ 甚至更高不加预条件根本没法用。3. PCG中预条件子的选型与实现3.1 预条件的本质换一把更合适的尺子PCG 的思路其实跟“坐标变换”这件事一脉相承。我们构造一个对称正定矩阵 M让它尽量接近 A同时 M 又很容易求逆。然后原本 Ax b 的问题被改写成M^{−1}Ax M^{−1}b如果 M 非常接近 A那 M^{−1}A 的条件数就会接近 1CG 收敛当然就快。但这个改写如果直接做会破坏对称性所以 PCG 实际是在 M^{-1}A 这个“非对称”矩阵上用 M-内积定义正交关系推导出一套保持原算法结构的递推式。你不需要在代码里显式计算 M^{-1}只需要能快速求解形如 Mz r 的线性系统即可。打个比方你要量一个细长健身房的对角线长度最笨办法是拿着米尺一步步走但健身房角落里有个现成的卷尺、或者你一眼能看出比例那先做个“预变换”就快得多。预条件子就是那个卷尺它不是答案本身却让答案变得唾手可得。3.2 预条件子选型对照表不同问题结构适合不同的预条件子。我整理了一个常用选择表格方便按场景快速定方向预条件子构造代价单次求解代价适用场景注意事项Jacobi / 对角O(n)O(n)对角占优系统、作为快速基线对强对角占优最有效条件数改善有限SSORO(nnz)约两次三角求解经典有限差分/有限元离散需要松弛因子一般取1.0附近不完全Cholesky IC(0)约等于一次matvec的几倍两次稀疏三角求解稀疏SPD矩阵工程中最常用填充元可能不稳定可用IC(τ)截断稀疏近似逆 SPAI较高一次稀疏矩阵向量乘并行计算、GPU友好构造成本可能超过收益多重网格/区域分解高极低理论上O(n)椭圆型PDE问题需要问题特定网格信息不是通用黑盒选预条件子有一个很反直觉的经验构造预条件子本身要花时间如果只用一次线性解那预条件的总成本可能反而比不预条件更贵。所以 PCG 划不划算要看你是不是在一个循环里反复求解多个类似的线性系统。比如高斯过程里在多个超参数组合下做交叉验证那预条件子构造一次、复用多次成本摊薄了收益就特别明显。3.3 我自己最常用的预条件策略在实际工作中我处理机器学习里的稠密但低秩加对角结构的系统时比如岭回归正规方程 (XᵀX λI)w Xᵀy我几乎总是先用对角预条件也就是取 M diag(A)。这个选择几乎不花成本有时就能把迭代次数降低个两三倍。如果系统来自稀疏图、或者大规模稀疏高斯过程我会试 IC(0)。这个实现并不复杂。对于 n100万左右的稀疏矩阵一次 IC(0) 分解的代价可以接受而它带来的收敛改善经常是数量级的。还有一类特殊场景A 是稠密核矩阵 K σ²I。对角预条件在核矩阵上效果很有限因为核矩阵的特征值往往按指数衰减对角元并不能反映整体谱分布。这时候我一般改用不完全Cholesky或者干脆利用已知的核函数先验做“Nyström近似预条件”就是把低秩的 K_approx 作为 M。这个思路适合核矩阵效果会好很多。3.4 不要被“迭代次数下降”迷惑这里一定要提醒一句PCG 的终极目标是减少总耗时不是单纯减少迭代次数。预条件子每次迭代都多一次 M z r 的求解如果这个求解比一次 A p 的 matvec 还贵那少掉的迭代次数可能都被预条件自身的成本吃回去了。我在一次大规模高斯过程实验中就犯过这个错误。当时数据点约有 12 万核矩阵没法显式构造我用了一种 RBF 核的快速 matvec 方法每次迭代很快。结果一个小伙伴建议的 IC 预条件虽然把迭代次数从 4600 次降到了 60 次但预条件子本身的构造和每次求解加起来的耗时比不预条件时还多了 40%。后来改用只做一次粗略的 Nyström 预条件构造成本低而且后续二十多次目标函数序列的求解都要复用同一套预条件总耗时反而下降了三倍多。所以永远要对着秒表说话不要盯着迭代次数曲线自我感动。4. 机器学习里的PCG实操从岭回归到高斯过程4.1 手写一个可用的PCG求解器不依赖魔改库我们用 NumPy 就可以写一个足够干净的 PCG 实现。关键点是 A 和 M 都以“可调用对象”传入这样就能支持隐式矩阵和自定义预条件子。import numpy as np def pcg(A, b, MNone, x0None, tol1e-6, max_iter1000, reportFalse): # A: callable, y A(x) # 如果是numpy矩阵则包一层matvec if not callable(A): matvec_A lambda x: A x else: matvec_A A if M is not None and not callable(M): # 如果M传的是矩阵那就直接矩阵乘法否则应该是callable解Mzr matvec_M lambda x: np.linalg.solve(M, x) else: matvec_M M # 可能为None n b.shape[0] if x0 is None: x np.zeros(n) else: x x0.copy() r b - matvec_A(x) z matvec_M(r) if matvec_M is not None else r.copy() p z.copy() rz_old np.dot(r, z) rz0 rz_old for it in range(max_iter): Ap matvec_A(p) alpha rz_old / np.dot(p, Ap) x alpha * p r - alpha * Ap if np.linalg.norm(r) tol * np.sqrt(rz0): break z matvec_M(r) if matvec_M is not None else r.copy() rz_new np.dot(r, z) beta rz_new / rz_old p z beta * p rz_old rz_new if report: return x, it 1, np.linalg.norm(r) return x这个代码里有几个处理很关键。r 的更新用的是r - alpha * Ap而不是每步重新算b − Ax可以省一次 matvec。但注意浮点累积误差会让这个残差慢慢漂移所以如果迭代次数非常多需要每隔比如 50 步重新计算一次真残差。M 是 None 时退化成普通 CG。需要明确的是M 传入的方式如果是矩阵则np.linalg.solve(M, r)是直接法只适合 M 比较稀疏或不太大的情况如果 M 是一个可以快速求解的隐式算子那应该传一个定义好的函数。4.2 场景一大规模岭回归别傻傻求逆岭回归的目标是求下面的 ww (XᵀX λI)^{−1} Xᵀy朴素做法是直接 Cholesky 分解这在特征维度几千左右还扛得住。但如果特征维度到几万、几十万或者 X 本身是稀疏的那直接分解就不划算了。我们可以把问题变成解线性方程组(XᵀX λI) w Xᵀy然后用 PCG。这里 A XᵀX λI 并不需要显式构造只要定义好 matvecdef ridge_matvec(w, X, lam): # X: (n_samples, n_features), 可能非常稀疏 return X.T (X w) lam * w理论上这一步需要两次稀疏矩阵向量乘一次 X w一次 X.T (Xw)并不贵。我试过一个 50 万样本、2 万维特征的稀疏文本分类问题直接求逆哪怕用稀疏 Cholesky 也会耗掉十几分钟而 PCG 配合最简单的对角预条件大概四十秒就收敛到相对残差 1e-6。对角线元素 A_ii 可以通过逐列求 X 的平方和得到不需要构造整矩阵。对于预条件取 M diag(XᵀX λI)。注意不要小看这个简单预条件它把特征尺度不均衡带来的条件数恶化基本压住了配合适当的小 λPCG 的迭代次数能控制在一两百以内。4.3 场景二高斯过程回归里的核线性系统高斯过程回归要解的线性系统是(K σ²I) α y其中 K 是 n×n 核矩阵。n 到一万以上直接对 K 做 Cholesky 需要存 n²/2 个元素大概 400MBfloat64 下 n10000 时约 400MB 以上n 到五万时光存储就超过 10GB基本告别大多数个人工作站。PCG 的登场条件是你必须给出一种高效的 K matvec。如果核函数是分段可并行的哪怕一次 matvec 是 O(n²) 的矩阵乘一次只需要 4GB 左右内存比直接存在矩阵里略好一些如果核函数是 RBF 这类标准核并且数据点分布在不太高的维度里还可以用树结构加速到 O(n log n)。总之只要 matvec 能算矩阵本身不存也没关系。实战场上我处理过一个 30 万点的时空高斯过程模型。显式 Cholesky 是不可能了我用了 RBF 核 快速多极子加速 matvec配合 Nyström 预条件。预条件 M 取 5000 个锚点上的低秩核矩阵求逆整体 PCG 收敛到 1e-5 大约花了 80 步每步 matvec 约 20 毫秒。整个后验推断不到两秒搞定。这个思路如果你要复现可以记住一个经验公式预条件子用的锚点数越多预条件效果越好但构造 M 和每次求解 Mz r 的花费也会越高。一般锚点数在几百到几千比较合适按收敛总时长调一调。4.4 场景三神经网络二阶优化中的CG子问题深度学习圈子里有一类优化方法是所谓“无Hessian优化”Hessian-free直接在 CG 子问题里用到 Hessian 向量积。大名鼎鼎的 Martens 2010 年论文用这个思路训练深度网络复现的关键点是CG 需要 Hessian 至少近似正定所以一般会用 Gauss-Newton 矩阵或阻尼正则化。如果你用 PyTorch 或 JAX 这类自动微分框架Hessian-vector product 很容易通过两次反向传播得到# PyTorch伪代码计算 H v g torch.autograd.grad(loss, params, create_graphTrue) flat_g torch.cat([p.view(-1) for p in g]) Hv torch.autograd.grad(flat_g, params, grad_outputsv)计算 Hessian 向量积只比一次反向传播贵一点但要完全避免显式构造 H 矩阵。CG 子问题出解之后再做一次线搜索更新参数。这套流程在凸问题区域收敛很快但到非凸地带容易遇到 Hessian 不正定。解决方法是加上阻尼项比如把 H 替换成 H λIλ 用一个自适应策略。这其实就是 PCG 里的预条件思想在用另一个维度上的体现——只不过这时预条件是“修正非正定性”而不单单是降条件数。5. PCG的数值细节与避坑指南5.1 残差更新 vs 真残差什么时候会漂移前面代码里用的是递归更新残差r ← r − αAp。这个方式的优点是省一次 matvec缺点是每次浮点运算都有一点舍入误差几十上百步之后 r 的每个分量可能都已经不是真实的 b − Ax。在条件数特别大的问题上误差会被放大得更厉害导致明明 r 已经很小了x 却离真解很远。我的习惯是如果矩阵条件数超过 1e8 或者迭代超过 200 步每隔 50 步强制重算一次 r b − Ax。如果发现重算后的残差比递归残差大了两个数量级以上那就是明显的漂移必须把这个真残差同步回 r、z、p 状态里。这个操作叫“重启”等价于丢掉已经积累的浮点噪声让迭代重新“校准”。如果问题有条件数爆炸的嫌疑还有一个更稳的选择直接用双精度浮点数而不是单精度。在 GPU 上很多人为了提速用 float16但 CG 在 float16 下极易发散。我个人建议PCG 至少用 float32最好 float64。5.2 容差到底设多少合适机器学习任务里解线性系统的精度要求和传统数值计算不太一样。有些阶段比如超参搜参你根本不需要 1e-10 的解反而追求快。典型经验值训练模型内部线性求解相对残差 1e-4 到 1e-6 足够了再高纯属浪费高斯过程预测1e-6 甚至 1e-5 都可以后验均值对噪声项很鲁棒和直接法结果做对比验证得设到 1e-10否则解析梯度和数值梯度差异会把误差归错因。判断收敛时我一般用相对残差||r_k|| / ||r_0|| tol。对规模差异很大的问题来说绝对残差没有意义——一个 b 的范数是 1e10 的问题和 b 范数是 1e-5 的问题同样的绝对残差含义完全不同。5.3 matvec 的优化优先级CG/PCG 每一步都要做一次 A p 和一次 M^{−1} r这两个操作的效率决定了整套算法的上限。所以如果你从默认实现转成自定义实现第一个要优化的就是 matvec。在稀疏矩阵上用 CSR/CSC 格式的稀疏矩阵乘法通常比稠密 NumPy 矩阵快得多。在深度学习框架里matvec 可以直接用算子而不必显式构建矩阵。多线程和 GPU 并行时注意 A p 这种操作能天然并行但 M z r 如果 M 是三角矩阵分并行会难一些。矩阵存储格式、融合算子、避免重复内存分配这些都是实际性能的主要来源。5.4 非正定矩阵的救火方法如果 A 不是对称正定CG 理论上不收敛或收敛极慢。实际中更常见的是你的 A 看似正定但由于浮点舍入或建模问题有一个接近零的特征值出现负值。这种时候有三个应急方案加一个对角线阻尼把 A 替换为 A δI这个 δ 可以动态调整你会发现 CG 突然又能收敛了改用最小残差类方法比如 GMRES、MINRES不需要正定性但每步开销更大如果你的 A 本身来源于某个优化问题的 Hessian 近似可以去检查是不是数据标准化出了问题很多“非正定”其实是数值尺度失衡造成的先做做特征缩放往往能解决。6. 常见问题速查表与排障思路症状可能原因处置办法迭代次数不降甚至越跑越慢矩阵非正定或条件数过大预条件没生效打印矩阵的特征值谱或对角元占比尝试换预条件子、加对角阻尼残差先降后涨发散浮点累积误差学习率等价物 α 过大每隔若干步重算真残差换 float64降低收敛容差要求某次迭代出现 NaN预条件子 M 奇异或接近奇异分母为零检查 M 分解时是否出现非正定方块IC 分解失败时加大截断阈值迭代正常但总时间比直接法还久预条件成本过高或 matvec 实现太慢用计时器分开统计构造预条件、每次 matvec、每次 M 求解的时间考虑更简单的预条件每次迭代都要新分配大数组内存碎片和分配开销大提前分配好 x, r, p, Ap, z 等向量在循环内复用换 GPU 后收敛行为变化浮点精度降为 float16非确定性原子操作强制 float64设置确定模式对比 CPU 上同样的迭代结果用深度学习框架做 Hessian-free 训练时 Hessian 向量积不对自动微分被错误使用了两次写一个小型数值梯度校验和有限差分对比验证 Hv 是否正确还有一个常被忽略的坑初值 x0 不是零向量时要确认 r0 的范数已经远小于 b 的范数。如果你初始化 w 0却在 PyTorch 里声明模型时默认权重不是零那 r0 可能本身就很小容差判断会提前终止。实际处理时我经常用零初始化这样 r0 b收敛判定逻辑最干净。7. 我对PCG的一次深刻整改经历最后聊一个印象很深的案例。之前做一个时空插值项目要解一个 80 万维的稀疏对称正定系统。最初版本是直接调用 scipy 的 CG 接口预条件子用了对角线对角占优的 Jacobi跑起来迭代了 5000 多次每次大概 3 毫秒总耗时十几秒人还能接受。但后来要在一个大的交叉验证循环里反复做耗时直接变成瓶颈。我把预条件换成 IC(0)迭代降到 300 次但构造 IC 要 2 秒而且每次求解 Mz r 也比对角慢一个数量级总体反而只快了 30%。这个结果让我意识到问题结构决定预条件策略不能光看迭代曲线。后来我分析矩阵的非零图模式发现它由三个子块构成边界耦合弱于是我换成了块对角预条件每个块做一次小规模的 Cholesky。结果迭代次数 200构造时间不到 0.5 秒每次 M 求解完美并行总耗时直接降到了原来的五分之一。这个项目给我的经验就一句话别迷信某个预条件子的名声把矩阵结构看清楚再选最简单、能契合结构的方式。很多问题里的矩阵都有自然的块结构或低秩结构预条件子如果能把这种结构吸收进去效果往往比通用 IC 好得多。这个经验在后来的高斯过程、图拉普拉斯系统包括深度学习二阶优化里反复应验算是 PCG 使用中我最有价值的一课。
阅读完成 · 觉得有帮助？
