Optax项目中LBFGS优化器与线搜索在自定义类中的应用实践

背景介绍

在机器学习模型训练过程中，优化算法的选择对模型性能有着重要影响。Optax作为JAX生态中的优化库，提供了多种优化算法实现。其中L-BFGS算法因其优秀的收敛特性，特别适合中小规模问题的优化。本文将重点探讨如何在自定义神经网络模型中使用Optax的LBFGS优化器，并配合线搜索功能实现更高效的参数优化。

核心问题分析

当开发者尝试在自定义神经网络类中使用LBFGS优化器时，常会遇到以下技术难点：

1. 参数处理复杂性：自定义类通常包含可训练参数和静态参数，需要正确处理
2. 线搜索接口适配：线搜索需要特定的值函数接口，与常规训练循环不同
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)

2. 参数分区处理

使用Equinox的partition和combine方法分离可训练参数和模型结构：

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

3. 线搜索适配

创建适配线搜索的lambda函数，确保在每次评估时都能正确组合模型参数和结构：

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

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. 参数分区的重要性：确保在优化过程中只更新可训练参数，保持模型结构不变
2. 线搜索机制：LBFGS的线搜索需要纯值函数来评估不同步长下的损失值
3. JIT编译兼容：使用eqx.filter_jit确保整个步骤可以被JAX正确编译优化
4. 梯度处理：明确区分可训练参数的梯度和模型的其他部分

实际应用建议

1. 对于中小规模问题，LBFGS+线搜索通常能获得更好的收敛性
2. 监控线搜索过程中的函数评估次数，避免不必要的计算开销
3. 考虑结合学习率调度器来动态调整初始步长
4. 对于大规模问题，可能需要改用随机优化方法或有限内存LBFGS变种

总结