非线性系列（三）—— 共轭梯度法在机器学习优化中的实战应用

1. 共轭梯度法从数学原理到机器学习优化第一次接触共轭梯度法(CG)是在研究生课程《数值分析》中当时只觉得这是个解线性方程组的数学工具。直到后来处理一个百万维度的推荐系统优化问题时我才真正体会到它的威力。相比常见的梯度下降法CG在机器学习优化中展现出三大独特优势收敛速度更快对于n维二次问题理论上最多n次迭代即可收敛内存效率高不需要存储完整的Hessian矩阵数值稳定性好特别适合病态矩阵问题举个实际例子在训练一个包含50万用户的推荐模型时使用普通梯度下降需要3000次迭代才能收敛而采用共轭梯度法仅需800次左右训练时间从6小时缩短到90分钟。这种优势在大规模稀疏数据场景如NLP的word2vec训练中更为明显。核心数学原理其实很直观想象你在山谷中寻找最低点。梯度下降就像蒙着眼沿最陡方向走难免之字形徘徊而共轭梯度法会记住之前的方向信息确保每个新方向都与之前方向共轭数学上表示为p_i^T A p_j0相当于消除了冗余搜索。2. 线性与非线性共轭梯度法实战对比2.1 线性CG对称正定问题的黄金标准上周帮同事调试一个金融风控模型时遇到了这样的场景需要求解形如Axb的线性系统其中A是用户行为特征的协方差矩阵天然对称正定。这时候线性CG就是首选工具。具体实现时要注意三个关键点# Python实现示例 def linear_cg(A, b, x0, tol1e-6, max_iterNone): if max_iter is None: max_iter len(b) x x0.copy() r b - A x p r.copy() rsold r.dot(r) for i in range(max_iter): Ap A p alpha rsold / p.dot(Ap) x alpha * p r - alpha * Ap rsnew r.dot(r) if np.sqrt(rsnew) tol: break p r (rsnew / rsold) * p rsold rsnew return x实际项目中我发现两个常见坑矩阵A不满足对称正定时需要预处理或改用其他算法条件数较大时建议使用Jacobi预处理器diag(A)^-12.2 非线性CG深度学习中的隐士高手处理神经网络训练这种非凸问题时Fletcher-Reeves和Polak-Ribiere两种变体最常用。去年优化一个图像生成模型时对比实验显示优化方法收敛步数最终损失内存占用Adam15000.021较高SGD with momentum22000.018低FR-CG8000.015最低PR-CG7500.014最低实现非线性CG时步长选择是关键。推荐使用强Wolfe条件线搜索from scipy.optimize import line_search def nonlinear_cg(f, grad, x0, methodPR, tol1e-6): x x0.copy() grad_fx grad(x) delta -grad_fx history [] while True: alpha line_search(f, grad, x, delta)[0] x_new x alpha * delta grad_fx_new grad(x_new) # 更新规则选择 if method FR: beta grad_fx_new.dot(grad_fx_new) / grad_fx.dot(grad_fx) elif method PR: beta max(0, grad_fx_new.dot(grad_fx_new - grad_fx) / grad_fx.dot(grad_fx)) delta -grad_fx_new beta * delta x, grad_fx x_new, grad_fx_new history.append(f(x)) if np.linalg.norm(grad_fx) tol: break return x, history3. 大规模机器学习中的工程实践3.1 分布式实现技巧当数据量超过单机内存时我通常采用如下方案矩阵分块将Hessian矩阵按行划分到不同worker异步通信各节点计算本地矩阵向量积后汇总混合精度使用float16存储float32计算在Spark环境中的核心代码结构def spark_cg(A_rdd, b, x0, n_workers4): # A_rdd: 分块存储的矩阵RDD x x0 r b - A_rdd.map(lambda block: block x).sum() p r rsold r.dot(r) for _ in range(max_iter): # 分布式矩阵向量乘 Ap A_rdd.map(lambda block: block p).treeAggregate( np.zeros_like(p), lambda x,y: xy, lambda x,y: xy) alpha rsold / p.dot(Ap) x alpha * p r - alpha * Ap rsnew r.dot(r) if np.sqrt(rsnew) tol: break p r (rsnew / rsold) * p rsold rsnew return x3.2 与深度学习框架集成在TensorFlow中自定义CG优化器的示例class ConjugateGradient(tf.keras.optimizers.Optimizer): def __init__(self, learning_rate0.01, nameCG, **kwargs): super().__init__(name, **kwargs) self._set_hyper(learning_rate, learning_rate) def _create_slots(self, var_list): for var in var_list: self.add_slot(var, prev_grad) self.add_slot(var, direction) def _resource_apply_dense(self, grad, var, apply_stateNone): lr self._get_hyper(learning_rate) prev_grad self.get_slot(var, prev_grad) direction self.get_slot(var, direction) # Fletcher-Reeves更新规则 beta tf.reduce_sum(grad**2) / (tf.reduce_sum(prev_grad**2) 1e-8) new_direction -grad beta * direction var_update var lr * new_direction var.assign(var_update) prev_grad.assign(grad) direction.assign(new_direction)4. 性能调优与避坑指南4.1 预处理技术对比不同预处理器的效果实测在CTR预测任务中预处理器类型迭代次数收敛时间内存开销无预处理12504.2h1xJacobi6802.1h1.2xSSOR(ω1.2)5201.8h1.5x不完全Cholesky3501.2h2.3x推荐一个实用的预处理选择策略当矩阵对角占优时用Jacobi足够对于稀疏矩阵尝试SSOR能接受较高内存开销时用不完全Cholesky4.2 常见问题排查最近调试一个推荐系统时遇到的典型问题迭代过程中残差突然增大。经过分析发现是数据类型溢出导致的解决方法包括使用64位浮点数对输入数据做标准化添加重启机制每k次迭代重置方向另一个坑是线搜索不收敛我的经验是调整Wolfe条件参数通常c11e-4, c20.9限制最大步长如alpha_max1.0当线搜索失败时改用固定学习率