PyTorch LBFGS：突破传统优化范式，以闭包之力驾驭非凸地形

1. 为什么LBFGS在PyTorch中如此特别当你第一次在PyTorch中使用LBFGS优化器时可能会被它奇怪的用法搞得一头雾水。其他优化器如SGD、Adam都是三步走计算损失、反向传播、参数更新。但LBFGS却要求你把整个计算过程封装在一个叫闭包(closure)的函数里再把这个函数传给step()方法。这到底是怎么回事其实这种设计背后隐藏着LBFGS算法的核心思想。LBFGS全称是Limited-memory Broyden-Fletcher-Goldfarb-Shanno是一种拟牛顿法。与普通的一阶优化器不同它需要近似计算Hessian矩阵二阶导数信息来指导优化方向。而闭包机制正是为了支持这种复杂的多步计算过程。想象你在爬山普通优化器就像只看着脚下的路一步步走而LBFGS则会先观察周围地形规划出一条更优的路径。闭包函数就是让它能够反复观察当前点的地形特征。2. 深入理解LBFGS的闭包机制2.1 闭包函数的作用原理闭包在LBFGS中扮演着关键角色。当你调用optimizer.step(closure)时LBFGS会在内部多次执行这个闭包函数。每次执行都会清空梯度zero_grad计算当前参数下的损失值反向传播计算梯度返回损失值这种设计允许LBFGS在单次参数更新中执行多次前向-反向计算这是因为它需要进行线搜索(line search)来确定最优步长。普通优化器通常使用固定学习率而LBFGS会动态调整。def closure(): optimizer.zero_grad() loss model(inputs, targets) loss.backward() return loss2.2 与传统优化器的对比让我们用经典的Rosenbrock函数来对比LBFGS和SGD的表现。这个函数被称为香蕉函数因为它的等高线呈香蕉形状有一个狭长的谷底通向全局最小值(1,1)。def rosenbrock(x): return (1 - x[0])**2 100 * (x[1] - x[0]**2)**2使用SGD时参数更新是直线向下的x torch.tensor([-1.0, 2.0], requires_gradTrue) opt torch.optim.SGD([x], lr1e-5) for i in range(1000): opt.zero_grad() loss rosenbrock(x) loss.backward() opt.step()而LBFGS则能更好地沿着谷底前进x torch.tensor([-1.0, 2.0], requires_gradTrue) opt torch.optim.LBFGS([x], line_search_fnstrong_wolfe) def closure(): opt.zero_grad() loss rosenbrock(x) loss.backward() return loss for i in range(100): opt.step(closure)实测下来LBFGS通常能在几十次迭代内收敛而SGD可能需要上千次而且容易在谷底两侧振荡。3. LBFGS的核心参数解析3.1 history_size记忆窗口的大小LBFGS通过存储最近的梯度变化来近似Hessian矩阵。history_size决定了存储多少步的历史信息。较大的值能提供更准确的二阶近似但会消耗更多内存。# 存储最近20步的梯度信息 optim.LBFGS(params, history_size20)经验表明对于大多数问题10-20的history_size已经足够。更大的值带来的收益会递减。3.2 max_iter每次更新的最大迭代次数这个参数控制每次调用step()时内部线搜索的最大迭代次数。注意这不是整个训练的总迭代次数。# 每次step最多尝试5次线搜索 optim.LBFGS(params, max_iter5)设置太大可能导致单次更新耗时过长太小可能找不到好的步长。通常4-10是比较合理的范围。3.3 line_search_fn线搜索策略PyTorch提供了两种线搜索算法None基本的Armijo条件回溯线搜索strong_wolfe满足强Wolfe条件的线搜索# 使用强Wolfe条件线搜索 optim.LBFGS(params, line_search_fnstrong_wolfe)强Wolfe搜索通常能带来更稳定的收敛但计算量稍大。对于简单问题使用None可能就够了。4. 实战在神经网络中使用LBFGS4.1 全连接网络的训练虽然LBFGS常用于小规模优化问题但在某些神经网络场景下也能发挥作用。比如训练小型全连接网络model nn.Sequential( nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10) ) optimizer optim.LBFGS(model.parameters(), lr1, max_iter10) def closure(): optimizer.zero_grad() output model(inputs) loss criterion(output, targets) loss.backward() return loss for epoch in range(10): optimizer.step(closure)注意这里的学习率lr意义不同因为实际步长由线搜索决定。通常可以设为1让线搜索来调整。4.2 与Adam优化器的对比在图像分类任务上对比LBFGS和Adam优化器训练时间最终准确率内存占用Adam快高低LBFGS慢相当高LBFGS的优势在于通常能找到更优的解超参数较少主要调整history_size不需要学习率调度但缺点也很明显每个step耗时更长内存占用随history_size线性增长大批量数据下表现不佳5. LBFGS的适用场景与注意事项5.1 最适合使用LBFGS的情况根据我的经验LBFGS在以下场景表现优异小规模参数优化参数数量1e5需要高精度解的科学计算问题损失函数比较平滑且计算不昂贵批量训练batch_size全部数据一个典型的例子是物理模拟中的参数拟合比如用神经网络拟合分子势能面。5.2 需要避免的情况LBFGS不太适合大规模深度学习如ResNet、Transformer随机梯度下降场景mini-batch训练损失函数有大量局部极小值需要频繁评估的在线学习我曾尝试在BERT微调中使用LBFGS结果内存直接爆掉。对于这种大模型还是Adam/AdamW更合适。5.3 常见问题排查如果LBFGS不收敛可以检查闭包函数是否正确实现了zero_gradforwardbackwardline_search_fn是否设为strong_wolfehistory_size是否足够大尝试增加到20初始参数是否合理有时需要先用SGD预热一个实用的技巧是先运行少量SGD迭代再用其输出初始化LBFGS。这能帮助避开一些糟糕的初始区域。