别再死记硬背公式了！用Python从零实现共轭梯度法（CG），直观理解每一步

别再死记硬背公式了用Python从零实现共轭梯度法CG直观理解每一步数值优化算法的学习常常陷入两个极端要么被复杂的数学推导吓退要么沦为调用库函数的黑箱操作。共轭梯度法Conjugate Gradient, CG作为求解大型稀疏线性系统的经典算法教科书上充斥着正交化、Krylov子空间等抽象概念。本文将带您用Python从零实现CG算法在代码调试和可视化中感受数学背后的直觉。1. 准备工作构建一个病态矩阵实验室理解CG算法的优势首先需要创造一个能凸显其特性的实验环境。我们故意构造一个条件数很大的病态矩阵观察不同优化方法的表现差异。import numpy as np from matplotlib import pyplot as plt def generate_ill_conditioned_matrix(n, kappa1e6): 生成指定条件数的对称正定矩阵 U np.random.randn(n, n) Q, _ np.linalg.qr(U) # 随机正交矩阵 s np.logspace(0, np.log10(kappa), n) # 对数分布的特征值 return Q np.diag(s) Q.T n 50 A generate_ill_conditioned_matrix(n) # 条件数约1e6的病态矩阵 b np.random.randn(n) # 随机生成右端项 x_true np.linalg.solve(A, b) # 真实解仅用于验证这个矩阵生成技巧的关键在于通过QR分解获得随机正交基让特征值呈对数分布人为制造巨大条件数保证矩阵对称正定CG算法的基本要求提示条件数condition number是矩阵最大与最小特征值的比值数值越大表示矩阵越病态常规迭代法收敛越慢。2. 热身最速下降法的锯齿困境作为对比基准我们先实现最速下降法Steepest Descent观察其在病态系统中的表现def steepest_descent(A, b, max_iter1000, tol1e-6): x np.zeros_like(b) residuals [] for _ in range(max_iter): r b - A x alpha np.dot(r, r) / np.dot(r, A r) x x alpha * r residuals.append(np.linalg.norm(r)) if residuals[-1] tol: break return x, residuals x_sd, res_sd steepest_descent(A, b) plt.semilogy(res_sd, labelSteepest Descent)运行后会看到残差范数虽然总体下降但呈现明显的锯齿模式。这是因为最速下降法相邻搜索方向正交在病态情况下反复震荡。这种现象在二维情况下可以直观可视化# 二维情况演示 A_2d np.array([[3, 2], [2, 6]]) b_2d np.array([2, -8]) def plot_contour(): x np.linspace(-3, 3, 100) y np.linspace(-2, 2, 100) X, Y np.meshgrid(x, y) Z 0.5*(A_2d[0,0]*X**2 2*A_2d[0,1]*X*Y A_2d[1,1]*Y**2) - b_2d[0]*X - b_2d[1]*Y plt.contour(X, Y, Z, levels20) plot_contour() x_sd_2d, _ steepest_descent(A_2d, b_2d, max_iter10) plt.plot(x_sd_2d[:,0], x_sd_2d[:,1], o-, labelSD Path) plt.legend()从等高线图上可以清晰看到最速下降法的路径呈90°转折这正是收敛慢的根本原因。3. 共轭梯度法的实现与魔法CG算法的核心思想是构造一组A-共轭的搜索方向确保每个方向上的优化一步到位。让我们拆解实现def conjugate_gradient(A, b, max_iterNone, tol1e-6): if max_iter is None: max_iter len(b) x np.zeros_like(b) r b - A x p r.copy() residuals [np.linalg.norm(r)] for k in range(max_iter): Ap A p alpha np.dot(r, r) / np.dot(p, Ap) x x alpha * p r_new r - alpha * Ap residuals.append(np.linalg.norm(r_new)) if residuals[-1] tol: break beta np.dot(r_new, r_new) / np.dot(r, r) p r_new beta * p r r_new return x, residuals x_cg, res_cg conjugate_gradient(A, b) plt.semilogy(res_sd, labelSteepest Descent) plt.semilogy(res_cg, labelConjugate Gradient) plt.legend()关键改进点在于搜索方向p通过beta系数继承前序方向信息确保新方向与之前所有方向A-共轭p_i^T A p_j 0 (i≠j)理论上n步内收敛实际受浮点误差影响在二维示例中CG的路径直接指向最小值plot_contour() x_cg_2d, _ conjugate_gradient(A_2d, b_2d, max_iter2) plt.plot(x_cg_2d[:,0], x_cg_2d[:,1], s-, labelCG Path) plt.legend()神奇的是CG在两步内就到达了精确解这正是共轭方向的威力——每个方向上的优化互不干扰。4. 进阶预条件技术实战对于极端病态系统即使CG也可能收敛缓慢。预条件技术通过变换系统改善条件数。以Jacobi预条件为例def jacobi_preconditioner(A): return np.diag(1/np.diag(A)) def pcg(A, b, M, max_iterNone, tol1e-6): if max_iter is None: max_iter len(b) x np.zeros_like(b) r b - A x z np.linalg.solve(M, r) p z.copy() residuals [np.linalg.norm(r)] for k in range(max_iter): Ap A p alpha np.dot(r, z) / np.dot(p, Ap) x x alpha * p r_new r - alpha * Ap residuals.append(np.linalg.norm(r_new)) if residuals[-1] tol: break z_new np.linalg.solve(M, r_new) beta np.dot(z_new, r_new) / np.dot(z, r) p z_new beta * p r, z r_new, z_new return x, residuals M jacobi_preconditioner(A) x_pcg, res_pcg pcg(A, b, M) plt.semilogy(res_sd, labelSD) plt.semilogy(res_cg, labelCG) plt.semilogy(res_pcg, labelPCG (Jacobi)) plt.legend()预条件CG的收敛曲线更加平滑快速。不同预条件子效果对比预条件子类型构造复杂度适用场景JacobiO(1)对角占优矩阵SSORO(n)一般SPD矩阵不完全CholeskyO(nnz)稀疏矩阵几何多重网格问题相关结构网格问题实际项目中好的预条件子往往能带来数十倍的加速。在SciPy中可以方便地使用预处理过的CGfrom scipy.sparse.linalg import cg, LinearOperator def preconditioned_solver(A, b): M_inv np.diag(1/np.diag(A)) # Jacobi预条件逆 M_operator LinearOperator(A.shape, matveclambda x: M_inv x) x, info cg(A, b, MM_operator) return x5. 调试技巧与数值验证实现CG时常见的坑包括矩阵不对称导致的算法失效浮点误差累积破坏共轭性停止准则设置不当验证算法正确性的方法def verify_cg_implementation(A, b): # 与权威解对比 x_ref np.linalg.solve(A, b) x_cg, _ conjugate_gradient(A, b) print(fReference vs CG error: {np.linalg.norm(x_ref - x_cg):.2e}) # 检查残差正交性 r b - A x_cg print(fFinal residual norm: {np.linalg.norm(r):.2e}) # 验证A-共轭性 _, history conjugate_gradient(A, b, max_iter3) p0 history[directions][0] p1 history[directions][1] print(fA-conjugacy check: {p0.T A p1:.2e})在实现过程中建议记录以下调试信息# 增强版CG记录器 def debug_cg(A, b, max_iterNone): history {x: [], residuals: [], directions: []} # ...实现与之前类似但记录更多信息 return x, history可视化工具能极大提升理解效率。比如绘制搜索方向的夹角history debug_cg(A_2d, b_2d, max_iter2) p0, p1 history[directions] angle np.degrees(np.arccos(p0 A_2d p1 / (np.linalg.norm(A_2dp0)*np.linalg.norm(p1)))) print(fAngle between search directions: {angle:.2f}°)6. 从线性到非线性CG的扩展应用虽然我们以线性系统为例但CG的思想可推广到非线性优化。以Rosenbrock函数为例def rosenbrock(x): return 100*(x[1]-x[0]**2)**2 (1-x[0])**2 def nonlinear_cg(f, x0, max_iter100, methodFR): x x0.copy() grad gradient(f, x) p -grad history [x.copy()] for _ in range(max_iter): alpha line_search(f, x, p) x x alpha * p grad_new gradient(f, x) if method FR: beta np.dot(grad_new, grad_new) / np.dot(grad, grad) elif method PR: beta np.dot(grad_new, grad_new - grad) / np.dot(grad, grad) p -grad_new beta * p grad grad_new history.append(x.copy()) return x, np.array(history) def gradient(f, x, eps1e-6): n len(x) grad np.zeros(n) for i in range(n): h np.zeros(n) h[i] eps grad[i] (f(x h) - f(x - h)) / (2*eps) return grad不同beta计算方式的对比方法公式特点Fletcher-Reeves (FR)β ∇f_{k1}^T ∇f_{k1} / ∇f_k^T ∇f_k保证下降性可能慢收敛Polak-Ribière (PR)β ∇f_{k1}^T (∇f_{k1}-∇f_k) / ∇f_k^T ∇f_k更高效但需要重启机制在机器学习中非线性CG常作为梯度下降的替代方案。例如在逻辑回归中def logistic_loss(w, X, y): z X w return np.mean(np.log1p(np.exp(-y * z))) def logistic_cg(X, y, max_iter100): w0 np.zeros(X.shape[1]) w, _ nonlinear_cg(lambda w: logistic_loss(w, X, y), w0, max_iter) return w相比SGDCG在中小规模数据集上通常需要更少迭代次数尤其适合参数数量适中的情形。