十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

共轭梯度法与PCG:机器学习大规模线性方程组求解实践

共轭梯度法与PCG:机器学习大规模线性方程组求解实践 直接聊结论如果你的机器学习模型里出现了一个大规模对称正定线性方程组兜底方案基本就是它——共轭梯度法Conjugate Gradient加了预条件处理之后就叫PCGPreconditioned Conjugate Gradient。这名字听着像是数值代数教材里的老古董但实际上从岭回归、高斯过程、图神经网络到有限元仿真里的物理信息神经网络凡是需要解Ax b的地方几乎都会碰到它。这篇文章我不打算给你抄一遍维基百科而是从机器学习实际会踩到的场景出发把CG为什么快、PCG到底在预处理什么、以及你写代码时最容易翻车的细节一次讲透。很多人第一次接触CG是在数值计算课上学完只觉得“哦这是个迭代法”。但真正到机器学习里用的时候比如你有一万个样本、一万个特征X^T X这个矩阵是10000×10000直接求逆要算(X^T X)^{-1}内存和时间都扛不住。这时候你才意识到我需要一个不用显式求逆、只靠矩阵乘法就能迭代求解的方法。共轭梯度法正好是这一族方法里对对称正定矩阵最友好的一个。它不需要你手动调学习率不需要你存庞大的Hessian矩阵也不需要你赌初始点选得好不好——只要矩阵是SPD对称正定理论上最多n步收敛配合预条件往往几十步就能达到机器精度。这篇文章适合这么几类人正在准备机器学习期末复习、想把“PCG”作为面试或考试亮点的同学做实验时发现np.linalg.inv或np.linalg.solve在大矩阵上慢到怀疑人生的研究生以及在化工、电化学、物理仿真里搭机器学习模型需要把求解器嵌进训练循环的工程师。下面我会从“为什么需要它”讲起然后是数学原理、完整的推导流程、可抄作业的Python实现最后是我实际调试中遇到的各种坑。1. 为什么机器学习里绕不开PCG1.1 从线性方程组说起机器学习里到处都是“解方程”的影子。最简单的线性回归最优权重是w (X^T X)^{-1} X^T y正规方程里那个(X^T X)^{-1}就是在解一个对称正定线性方程组。岭回归加一个λI照样是。高斯过程回归要算核矩阵的逆和行列式本质也是解方程组。甚至深度学习里某些优化方法比如L-BFGS在计算搜索方向时也需要解近似的线性系统。但要命的点在于这些方程组的规模会随着数据量飞速膨胀。你看下面这个对比就明白了方法内存复杂度时间复杂度10000×10000矩阵的直观感受直接求逆(X^T X)^{-1}O(n²)O(n³)内存约800MB计算量10^12量级小服务器直接卡死Cholesky分解O(n²)O(n³/6)比求逆省一点但依旧扛不住大规模共轭梯度CGO(n²)只存矩阵迭代次数×O(n²)迭代几十次每次一次矩阵乘法轻松很多PCG预条件CGO(n²)矩阵预条件子通常远少于CG的迭代次数收敛更快尤其病态矩阵注意上面这个表格里CG的“O(n²)”是矩阵向量乘法的复杂度。如果你的矩阵是稀疏的比如某些图模型或有限元离散出来的矩阵那么每次矩阵乘法只需要O(非零元个数)直接法的复杂度反而会高到没法用。这就解释了为什么PCG在实际工程中几乎是“大规模对称正定线性系统”的标准答案。1.2 直接法和迭代法的取舍逻辑很多人的第一反应是“我直接调numpy.linalg.solve不就行了吗”对几千维的矩阵确实行。可一旦到了几万维、几十万维直接法就有两个致命伤。第一填充fill-in问题。直接法无论是LU分解还是Cholesky分解过程中都会产生大量新的非零元。一个稀疏矩阵分解完之后可能变成稠密矩阵内存直接爆掉。第二没法利用矩阵的结构。如果A是稀疏的、或者有某种特殊结构比如低秩扰动、Toeplitz结构直接法完全无视这些愣头青一样地做分解。迭代法就不一样它只关心“你给我一个向量我能返回A乘以这个向量”至于A内部长什么样它不关心。这个特性在机器学习里特别重要——有时候你甚至不需要显式构造出A只需要一个能算A v的黑盒函数就够了。比如在核方法里核矩阵本来就是隐式定义的你根本不会真的把它存下来这时候PCG几乎是唯一可行的路径。我自己的体会是遇到对称正定线性系统先别急着求逆先想想能不能用PCG。这不是情怀而是当你试过一次100000维的岭回归之后就再也回不去了。2. 共轭梯度法的数学原理2.1 从最速下降法讲起要理解CG最好的入口是最速下降法。假设你要求解Ax b其中A对称正定。定义一个二次型f(x) 0.5 * x^T A x - b^T x这个函数的梯度是Ax - b令梯度为零就得到原方程。最速下降法做的事情很简单每次沿负梯度方向走一步步长通过一维搜索精确求出也就是在方向上找让f(x)最小的α。这个方法一定收敛但它有个著名的毛病——zigzag现象。如果矩阵的条件数很大比如一来一回两个方向上的“坡度”差别悬殊最速下降法会在两个方向之间来回震荡收敛慢得让人崩溃。形象一点说就像你在一个狭长的山谷里下山每次都垂直于等高线走结果一直在谷壁之间弹来弹去。2.2 共轭方向的核心思想共轭梯度法的聪明之处在于它不只是沿着梯度方向走而是构造一组A-共轭方向。什么叫A-共轭两个非零向量p_i和p_j如果满足p_i^T A p_j 0 (i ≠ j)就称它们关于A共轭。理解这个概念的钥匙在于x^T A x c在几何上是一个高维椭球椭球的轴方向就是A的特征向量方向。如果你能沿着这些“主轴方向”依次搜索每次都精确优化一个方向那么最多n步就能走到椭球中心——也就是最优解。而共轭方向恰好等价于在A定义的内积空间中正交所以在这些方向上依次做精确线搜索不会有任何一步破坏之前的方向上已经取得的优化成果。你可能会问我直接对A做特征分解然后沿特征向量方向走不就行了吗理论上行但特征分解的代价是O(n³)比直接求逆还贵毫无意义。CG的厉害之处在于它不需要事先知道共轭方向而是在迭代过程中用梯度信息自动生成共轭方向。每一轮的搜索方向是当前残差与上一个方向的线性组合组合系数由A-共轭条件唯一确定。这样既能保证方向两两共轭又不需要存下所有历史方向——每一轮只要记住上一个方向就行内存开销极小。2.3 为什么还需要预条件PCGCG的收敛速度理论上由矩阵的条件数κ(A)决定。条件数越大收敛越慢。具体来说CG的误差上界和(√κ - 1)/(√κ 1)的k次方有关。如果κ 1000这个比值大约是0.94意味着每步只减少6%的误差你想想得迭代多少轮才能收敛。预条件的思想特别直白找一个和A接近又容易求逆的矩阵M把原方程转换成条件数更小的等价方程。你希望M^{-1}A的特征值尽可能聚拢到1附近。最极端的情况是M A那一步就收敛了但问题是要先求A的逆——这等于没解决。所以预条件子就是在“像A”和“容易求逆”之间做平衡。PCG里预条件器和普通CG的唯一区别是标准的CG作用在等价系统M^{-1}Ax M^{-1}b上但为了数值稳定性和避免显式构造M^{-1}A我们把它隐含地嵌进算法流程里。你只需要额外提供一个算子给定向量r返回z M^{-1}r。常见的预条件子包括对角预条件Jacobi、不完全Cholesky、块对角预条件甚至在某些深度学习场景里会直接用另一个优化器的对角近似来充当预条件子。3. 逐步推导的PCG迭代流程3.1 算法伪代码先给你一个可以直接照抄的伪代码框架。假设我们需要求解Ax b给定初始猜测x_0预条件子Mr_0 b - A x_0 z_0 M^{-1} r_0 p_0 z_0 for k 0, 1, 2, ...: alpha_k (r_k^T z_k) / (p_k^T A p_k) x_{k1} x_k alpha_k p_k r_{k1} r_k - alpha_k A p_k if ||r_{k1}|| tol: break z_{k1} M^{-1} r_{k1} beta_k (r_{k1}^T z_{k1}) / (r_k^T z_k) p_{k1} z_{k1} beta_k p_k如果你把M取成单位矩阵那这就是标准CG如果M是对角矩阵那就是对角预条件CG也叫Jacobi预条件。这段代码里每一步都有几何意义我们逐一拆开看。3.2 每一行算子的实际含义残差与搜索方向。残差r_k b - A x_k它同时也是二次型梯度f(x_k)的负梯度方向。搜索方向p_k是当前残差的“预条件修正版”加上上一个方向的修正。CG之所以不需要存储所有历史方向是因为可以证明在精确算术下当前的搜索方向已经天然与之前所有方向A-共轭。这个性质源于每一步线搜索都到达了当前子空间里的极小点——这是Krylov子空间方法最漂亮的地方它在一个不断扩张的子空间里做最优近似而且扩张的方式完全由残差自动决定。步长alpha的推导。沿着方向p_k走多远作精确线搜索令f(x_k alpha p_k)对alpha的导数为零推导出来正好是上面那个算式。分母p_k^T A p_k是一个标量每次迭代只需要做一次矩阵向量乘法和两次内积整个算法最重的计算就是A p_k这一下。这也是为什么我说PCG的复杂度由“一次矩阵向量乘法”主导。beta与共轭性的保障。beta的经典形式在推导时有两种等价写法一个是Fletcher-Reeves的(r_{k1}^T r_{k1}) / (r_k^T r_k)一个是Polak-Ribière的r_{k1}^T (r_{k1} - r_k) / (r_k^T r_k)。在有预条件时正规写法是上面伪代码里带z的那个版本。需要注意的是当矩阵是对称正定、且用精确算术时这些公式是等价的但在浮点环境下Polak-Ribière对非精确搜索更鲁棒。实践里我倾向于用带预条件的标准公式但如果你发现残差不降反升换一种beta公式往往有奇效。3.3 终止条件与误差估计到底迭代多少轮才算收敛不能只看残差绝对大小因为如果b本身量纲很大残差绝对值意义不大。更通用的判据是相对残差||r_k|| / ||r_0|| tol这里tol通常取1e-6到1e-10之间。注意CG在迭代过程中残差范数未必单调递减——这是数值误差导致的正常现象尤其是当矩阵条件数很大时残差可能出现小幅震荡。所以如果只是用PCG做机器学习中间步骤我一般建议设一个宽松点的容差比如1e-6没必要追求机器精度否则多出来的迭代次数纯粹是浪费算力。而在物理仿真那种需要高精度的场合再把容差降到1e-12也不迟。4. 在机器学习中的典型落地场景4.1 岭回归与弹性网岭回归的目标是最小化||y - Xw||² λ||w||²其正规方程是(X^T X λI)w X^T y。这里的系数矩阵A X^T X λI天然对称正定完美符合CG的要求。问题在于如果X的维度是几万甚至几十万X^T X根本没法显式构造。但你看PCG的算法流程里面所有对A的操作都是“输入一个向量v输出A v”。而A v X^T (X v) λ v这一步只需要先算X v再算X^T (X v)根本不需要把X^T X这个矩阵组装出来。如果你是做稀疏特征的点击率预估特征维度动辄上亿这个技巧就是唯一可行的方案。这也是我特别喜欢PCG的原因它强迫你想清楚A的本质而不是把它当成一个黑盒矩阵。4.2 高斯过程回归高斯过程回归要预测新点x*的均值公式是μ* k_*^T (K σ²I)^{-1} y。这里的K是核矩阵N个训练点就是N×N。核矩阵不但对称正定而且通常病态——因为核函数在相近的点之间高度相关特征值衰减极快。你如果直接np.linalg.solve不仅慢而且数值上也不稳。更麻烦的是如果核函数是像RBF这种核矩阵本身就是稠密的没有稀疏性可以利用。这时候PCG的优势在于核矩阵的乘法可以按块处理甚至配合FGT快速高斯变换类方法把矩阵乘法近似加速到O(N log N)。不过一般来说高斯过程回归更常用的是Cholesky分解加低秩近似只有在N特别大或者你手头只有评估核函数能力时PCG才真正登场。但即便这样PCG依然是替代直接法的可用方案尤其是当核矩阵的逆没法稳定计算时PCG配合对角加秩一预条件往往能收敛得又快又稳。4.3 深度学习里的隐含线性求解很多人以为深度学习跟线性求解器八竿子打不着。但L-BFGS优化器内部要解一个关于搜索方向的线性系统Implicit Deep Learning这类模型更是直接把不动点迭代/线性求解嵌进了网络前向传播。再有就是图神经网络里常见的归一化拉普拉斯方程以及某些Transformer变体里对注意力矩阵的谱归一化——这些都会涉及大规模稀疏线性系统。我见过的实际案例里用PCG加速Implicit Layer训练能把单步前向的耗时降低一个数量级。在这个场景下PCG不再是一个“数值代数作业”而是模型架构的一部分性能直接决定训练速度。4.4 理工交叉场景热词里“机器学习 化工”、“电化学 机器学习”看着跨度大本质上都是同一类问题把物理规律和实验数据拟合在一起。比如电化学阻抗谱的等效电路拟合、化工过程传质方程的偏微分反问题往往要离散化后用有限元或有限体积法求解大型稀疏线性系统然后嵌入到优化目标里反复调用。在这些场景中PCG不但是“可选方案”还是默认的求解器因为离散化后的刚度矩阵就是典型的对称正定稀疏矩阵。你在实验室搭建机器学习服务器时装一套带PCG功能的科学计算库比如SciPy或者PETSc比手动去解那些上千维的稠密方程组要高效得多。5. 写一个能直接用的PCG实现5.1 Python代码numpy版我直接给你一段能跑通的numpy实现适合中等规模稠密矩阵。它不追求最高性能但逻辑非常清晰你可以照着改造成稀疏版本或者嵌入到自己的模型里。import numpy as np def pcg(A, b, MNone, x0None, tol1e-6, max_iterNone): Preconditioned Conjugate Gradient for symmetric positive definite A. Parameters ---------- A : callable or 2D array 如果A是可调用对象预期 A(v) 返回 A v 如果是2D数组内部会转成 lambda v: A v b : 1D array M : callable or 2D array, optional 预条件子M(v) 返回 M^{-1} v。默认None即单位阵 x0 : 1D array, optional 初始猜测 tol : float 相对残差容差 max_iter : int, optional 最大迭代次数默认 len(b) if not callable(A): A_mat A A lambda v: A_mat v if M is None: M lambda v: v elif not callable(M): M_mat M M lambda v: M_mat v n len(b) if x0 is None: x np.zeros_like(b, dtypefloat) else: x x0.copy() r b - A(x) z M(r) p z.copy() rz np.dot(r, z) if max_iter is None: max_iter n r0_norm np.linalg.norm(r) if r0_norm 0: return x, 0 for k in range(max_iter): Ap A(p) pAp np.dot(p, Ap) if pAp 1e-30: # 理论上 SPD 矩阵不会出现该情况数值异常时提前退出 break alpha rz / pAp x x alpha * p r r - alpha * Ap r_norm np.linalg.norm(r) if r_norm / r0_norm tol: return x, k 1 z M(r) rz_new np.dot(r, z) beta rz_new / rz p z beta * p rz rz_new return x, max_iter这段代码的核心是只用一次A v和一次M(r)整个迭代基本没有多余的重量级运算。用的时候正着来n 10000 A np.random.randn(n, n) A A A.T n * np.eye(n) # 构造对称正定矩阵 b np.random.randn(n) x, iters pcg(A, b, tol1e-8) print(fPCG迭代次数: {iters}) print(f残差: {np.linalg.norm(A x - b) / np.linalg.norm(b)})如果矩阵太大没法一次性读进内存就把A传成函数def matvec(v): return X.T (X v) lam * v x, iters pcg(matvec, X.T y, tol1e-8)5.2 代码里的几个关键细节pAp等于零的处理。理论上SPD矩阵对任意非零p都有p^T A p 0但浮点误差下可能出现极小甚至负值。一旦pAp接近零说明已经收敛或者数值已经乱套直接break是最稳妥的选择。预条件子的实现。代码里把M定义成“返回M^{-1}v”的算子这一步要特别小心。很多人把M理解成“存一个矩阵M”然后算np.linalg.solve(M, v)。这也没错但如果你能利用M的特殊结构比如对角、三角、稀疏就千万别做显式求逆。比如Jacobi预条件就是M np.diag(A)那么M(r) r / np.diag(A)一步搞定。强制用浮点。注意代码里x np.zeros_like(b, dtypefloat)如果你传入的b恰好是整型后面加减法会直接出int溢出。这种小坑我在调试时没少踩。5.3 预条件怎么选从Jacobi到不完全Cholesky选预条件子是PCG的灵魂。如果选了个烂预条件子可能比不预处理还慢。反过来说一个好的预条件子能让你从上百次迭代直接降到十几次。我按实用度排个序对角预条件JacobiM diag(A)。实现零成本效果中等。对于对角占优的矩阵效果不错但对病态严重的矩阵帮助有限。不完全CholeskyIC对A做近似的Cholesky分解但丢弃掉填充元。效果通常比Jacobi好很多尤其对有限元类的稀疏矩阵。SciPy里scipy.sparse.linalg.spilu就提供ILU分解你可以取上三角作为预条件子。块对角预条件把矩阵按某种物理含义分块对每个块分别求逆。在多物理场耦合问题里很常用。特定问题预条件如果你知道A的结构比如A B^T B λI那么可以构造更精巧的预条件子。比如用M (B^T B的对角 λI)之类的近似。我强烈建议你在动手之前先打印一下矩阵的条件数或者运行一次不预处理的CG看看收敛曲线。如果迭代次数超过几百次还没收敛基本可以断定需要更好的预条件子。6. 常见问题与排查技巧实录6.1 收敛慢先查条件数和预条件如果你发现PCG迭代了几百次都不收敛有两个地方是首要怀疑对象。第一矩阵真的是对称正定吗浮点运算中X^T X λI一般没问题但如果你做了一个原地更新导致矩阵不再对称CG会产生完全不稳定的行为。你可以在迭代前快速检查A是否对称np.linalg.norm(A - A.T)如果大于1e-8先修正再谈收敛。第二预条件子是否真的“预”了有时候你随手选了DTO对角预条件而矩阵本来就是强对角占优那预条件几乎没提升。再给一个实用技巧打收敛曲线。代码里把每次迭代的相对残差存下来画成log图。如果曲线平滑下降说明一切正常如果曲线先降后升那要么是数值误差累积要么是预条件子在某个方向放大了误差。6.2 数值不稳定残差震荡的真相即使矩阵是SPDCG在浮点环境下依旧会失去“理论上的共轭性”。典型表现是迭代次数超过一定阈值后残差不降反升。这不是算法错了而是舍入误差逐渐累积破坏了搜索方向的共轭性。解决办法有几个方向一是提高预条件子的质量让条件数变小自然就减少了舍入误差的累积二是重新开始restart也就是每隔几十步把搜索方向重置为当前残差方向三是改用双精度甚至更高精度。实践里如果你的PCG需要超过500次迭代才能收敛建议认真考虑预条件子而不是闷头加大迭代上限。调试时设一个硬性最大迭代次数。我一般设成min(n, 500)防止程序无限循环卡死。同时在每个迭代里定期打印残差肉眼观察趋势。6.3 内存不足别组装矩阵前文说过PCG最大的好处是“矩阵向量乘法”可以黑盒化。但初学时很容易犯一个错——为了用PCG先把A完整构造出来结果内存爆了。反过来想一想如果A是X^T X你可以用两次矩阵乘法替代如果A是核矩阵你可以按minibatch算块对角近似。这是PCG工程化和NumPy教学代码之间最核心的差别。实验室搭服务器跑大规模实验时我建议用scipy.sparse.linalg.cg它内置了PCG而且对稀疏矩阵有优化。你的职责是提供A的线性算子和预条件子而不是把A组装成稠密矩阵往内存里塞。6.4 常见问题速查表现象可能原因排查/解法迭代次数超过n不收敛矩阵不是SPD舍入误差累积检查对称性换预条件子重启残差先快速下降后停滞预条件子效果一般容差太严换IC预条件放宽tol程序内存溢出显式构造了稠密A改成黑盒矩阵向量乘法残差偶尔上升浮点精度问题beta公式选择不当换Polak-Ribière提高精度PCG结果和np.linalg.solve相差大容差没调到足够小对比时用tol1e-107. 我在实际项目里怎么用它说这么多我给一个具体的使用场景。之前我在做电化学阻抗谱的等效电路参数拟合模型里有一个几千维的参数协方差矩阵需要反复求逆。刚开始直接np.linalg.inv一次拟合要跑几个小时。后来换成PCG预条件子用对角加低秩近似一次矩阵向量乘法只涉及稀疏存取和少量密集运算最终把单次迭代时间压缩到原来的二十分之一。说实话那次调研给我的最大收获不是PCG本身而是“先想清楚A的结构再决定用什么算法”这个习惯。如果你也想在项目里真正落实PCG我建议从一维回归问题开始。拿波士顿房价数据集虽然现在不推荐用了但做练习没毛病把特征标准化之后构造X^T X λI先不用预条件跑一遍CG再上Jacobi预条件对比迭代次数。能清晰看到预条件带来的差别之后再往高斯过程、图模型那些更大规模的问题上迁移。根据个人经验PCG最值得记的一句话是它不关心A长什么样只关心你能不能给它一个快速的矩阵向量乘法。能把这句话用明白机器学习里一大部分线性代数瓶颈就解开了。
返回列表