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

1. 共轭梯度法：从数学原理到机器学习优化

第一次接触共轭梯度法(CG)是在研究生课程《数值分析》中，当时只觉得这是个解线性方程组的数学工具。直到后来处理一个百万维度的推荐系统优化问题时，我才真正体会到它的威力。相比常见的梯度下降法，CG在机器学习优化中展现出三大独特优势：

1. 收敛速度更快：对于n维二次问题，理论上最多n次迭代即可收敛
2. 内存效率高：不需要存储完整的Hessian矩阵
3. 数值稳定性好：特别适合病态矩阵问题

举个实际例子，在训练一个包含50万用户的推荐模型时，使用普通梯度下降需要3000+次迭代才能收敛，而采用共轭梯度法仅需800次左右，训练时间从6小时缩短到90分钟。这种优势在大规模稀疏数据场景（如NLP的word2vec训练）中更为明显。

核心数学原理其实很直观：想象你在山谷中寻找最低点。梯度下降就像蒙着眼沿最陡方向走，难免"之字形"徘徊；而共轭梯度法会记住之前的方向信息，确保每个新方向都与之前方向"共轭"（数学上表示为p_i^T A p_j=0），相当于消除了冗余搜索。

2. 线性与非线性共轭梯度法实战对比

2.1 线性CG：对称正定问题的黄金标准

上周帮同事调试一个金融风控模型时，遇到了这样的场景：需要求解形如Ax=b的线性系统，其中A是用户行为特征的协方差矩阵（天然对称正定）。这时候线性CG就是首选工具。具体实现时要注意三个关键点：

# Python实现示例 def linear_cg(A, b, x0, tol=1e-6, max_iter=None): 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

实际项目中我发现两个常见坑：

1. 矩阵A不满足对称正定时，需要预处理或改用其他算法
2. 条件数较大时，建议使用Jacobi预处理器（diag(A)^-1）

2.2 非线性CG：深度学习中的隐士高手

处理神经网络训练这种非凸问题时，Fletcher-Reeves和Polak-Ribiere两种变体最常用。去年优化一个图像生成模型时，对比实验显示：

实现非线性CG时，步长选择是关键。推荐使用强Wolfe条件线搜索：

from scipy.optimize import line_search def nonlinear_cg(f, grad, x0, method='PR', tol=1e-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, history

3. 大规模机器学习中的工程实践

3.1 分布式实现技巧

当数据量超过单机内存时，我通常采用如下方案：

1. 矩阵分块：将Hessian矩阵按行划分到不同worker
2. 异步通信：各节点计算本地矩阵向量积后汇总
3. 混合精度：使用float16存储，float32计算

在Spark环境中的核心代码结构：

def spark_cg(A_rdd, b, x0, n_workers=4): # 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: x+y, lambda x,y: x+y) 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

3.2 与深度学习框架集成

在TensorFlow中自定义CG优化器的示例：

class ConjugateGradient(tf.keras.optimizers.Optimizer): def __init__(self, learning_rate=0.01, name="CG", **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_state=None): 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预测任务中）：

推荐一个实用的预处理选择策略：

1. 当矩阵对角占优时，用Jacobi足够
2. 对于稀疏矩阵，尝试SSOR
3. 能接受较高内存开销时，用不完全Cholesky

4.2 常见问题排查

最近调试一个推荐系统时遇到的典型问题：迭代过程中残差突然增大。经过分析发现是数据类型溢出导致的，解决方法包括：

1. 使用64位浮点数
2. 对输入数据做标准化
3. 添加重启机制（每k次迭代重置方向）

另一个坑是线搜索不收敛，我的经验是：

• 调整Wolfe条件参数（通常c1=1e-4, c2=0.9）
• 限制最大步长（如alpha_max=1.0）
• 当线搜索失败时，改用固定学习率