Optax项目中使用LBFGS优化器结合线搜索的实践指南

背景介绍

在机器学习模型训练过程中，二阶优化算法因其收敛速度快的特点受到广泛关注。L-BFGS作为经典的拟牛顿法优化算法，在深度学习领域也有重要应用。Google DeepMind开发的Optax库提供了LBFGS优化器的实现，并支持线搜索功能，能够进一步提升优化效果。

核心问题

当在Equinox框架下构建自定义模型时，如何正确使用Optax的LBFGS优化器并启用线搜索功能，需要特别注意几个关键点：

1. 线搜索需要单独定义纯损失函数（value_fn），而非同时计算值和梯度的函数
2. 在Equinox框架下需要正确处理模型的可微分和不可微分部分
3. 参数更新时需要保持模型结构的完整性

解决方案详解

1. 定义纯损失函数

线搜索算法需要一个仅返回损失值的函数，因此我们需要单独定义：

def loss_fn(model, ts, ys_true):
    y0 = jnp.array([0.0])
    y_pred = model(ts, y0)
    return jnp.mean((y_pred - ys_true) ** 2)

这与常见的value_and_grad函数不同，后者会同时返回损失值和梯度。

2. 模型参数处理

在Equinox框架下，我们需要区分模型的可微分和不可微分部分：

model_params, model_struct = eqx.partition(model, eqx.is_array)

这种分割确保了在优化过程中只更新可训练参数，同时保持模型结构不变。

3. 创建线搜索兼容的损失函数

为了在线搜索中使用，我们需要创建一个闭包函数，将当前模型结构和输入数据绑定：

loss_fn_ = lambda model_params: loss_fn(
    eqx.combine(model_params, model_struct), ti, yi)

这个lambda函数将模型参数与固定结构重新组合，确保每次线搜索评估时模型结构保持一致。

4. 优化步骤实现

完整的训练步骤实现如下：

@eqx.filter_jit
def make_step(ti, yi, model, opt_state):
    # 计算损失和梯度
    loss, grads = eqx.filter_value_and_grad(loss_fn)(model, ti, yi)
    
    # 准备优化器输入
    grads = eqx.filter(grads, eqx.is_array)
    opt_state = eqx.filter(opt_state, eqx.is_array)
    model_params, model_struct = eqx.partition(model, eqx.is_array)
    
    # 创建线搜索兼容的损失函数
    loss_fn_ = lambda model_params: loss_fn(
        eqx.combine(model_params, model_struct), ti, yi)
    
    # 执行优化步骤
    updates, opt_state = optim.update(
        grads, opt_state, model_params, 
        value=loss, grad=grads, value_fn=loss_fn_)
    
    # 更新模型
    model = eqx.apply_updates(model, updates)
    return loss, model, opt_state

关键注意事项

1. 函数类型匹配：确保传递给线搜索的是纯损失函数，而非value_and_grad函数
2. 模型完整性：在参数更新前后保持模型结构的完整性
3. 性能考量：线搜索会增加每次迭代的计算量，但通常能减少总迭代次数
4. 学习率设置：即使启用了线搜索，初始学习率的设置仍然会影响优化效果

实际应用建议

对于科学计算和物理信息神经网络(PINN)等场景，LBFGS+线搜索的组合往往能取得比一阶优化器更好的效果。建议：

1. 先使用Adam等一阶优化器进行预训练
2. 切换到LBFGS进行精细优化
3. 监控线搜索的接受率，调整初始学习率

通过合理使用Optax提供的LBFGS实现，可以在保持代码简洁的同时，获得接近二阶优化算法的性能表现。