Optax项目中L-BFGS优化器使用问题深度解析

背景介绍

Optax是Google DeepMind开发的一个基于JAX的优化库，提供了多种优化算法的实现。L-BFGS（Limited-memory Broyden-Fletcher-Goldfarb-Shanno）是一种准牛顿优化方法，特别适合解决大规模非线性优化问题。

问题现象

用户在从Jaxopt迁移到Optax时，发现Optax实现的L-BFGS优化器与Jaxopt版本存在差异，主要表现为：

1. 在某些情况下损失函数会出现非单调下降的情况
2. 优化过程中可能出现损失值突然增加（"spikes"）
3. 在元学习场景下，由于使用jax.lax.while_loop实现，无法进行反向传播

技术分析

实现差异

Optax和Jaxopt的L-BFGS实现存在几个关键区别：

1. 初始步长选择：Optax默认在每次迭代中使用步长1作为初始猜测，而Jaxopt可能有不同的策略
2. 梯度处理：Optax在第一次迭代时会对梯度进行裁剪
3. 线搜索实现：两个库使用的默认线搜索算法可能不同

损失增加问题

当观察到损失函数非单调下降时，通常与线搜索配置有关：

1. 最大线搜索步数不足：max_linesearch_steps参数设置过小可能导致线搜索无法找到合适的步长
2. 梯度不匹配：如果在L-BFGS外部对梯度进行了处理（如裁剪），但value_fn没有相应调整，会导致优化方向不正确
3. 数值稳定性：某些情况下，梯度裁剪可能使搜索方向不再是下降方向

元学习场景限制

在元学习中使用L-BFGS作为内部优化器时，会遇到以下挑战：

1. 不可微分问题：由于使用jax.lax.while_loop实现，无法直接进行反向传播
2. 替代方案：
   • 使用隐函数定理和自定义JVP实现
   • 考虑使用支持可微分while循环的库如Optimistix

解决方案与最佳实践

1. 线搜索配置：

   • 适当增加max_linesearch_steps参数
   • 启用verbose模式监控线搜索过程
   • 确保传递给优化器的梯度与目标函数一致

2. 梯度处理：

   • 避免在L-BFGS外部单独处理梯度
   • 如需梯度裁剪，应在value_fn中保持一致处理

3. 元学习场景：

   • 考虑使用SGD等可微分优化器作为内部优化器
   • 如必须使用L-BFGS，可参考Jaxopt实现自定义可微分版本

示例代码修正

对于简单优化问题，正确的L-BFGS使用方式如下：

def fun(x):
    return jnp.sum(x**2)

x = jnp.array([1.0, 2.0, 3.0])
opt = optax.lbfgs(max_linesearch_steps=20)  # 增加线搜索步数
opt_state = opt.init(x)

for i in range(20):
    value, grad = jax.value_and_grad(fun)(x)
    updates, opt_state = opt.update(
        grad, opt_state, x, 
        grad=grad, value=value,
        value_fn=fun
    )
    x = optax.apply_updates(x, updates)

总结

Optax中的L-BFGS实现是一个高效的优化器，但在使用时需要注意其与Jaxopt的实现差异。特别是在线搜索配置和梯度处理方面需要格外小心。对于元学习等需要反向传播的场景，可能需要考虑替代方案或自定义实现。理解这些底层细节有助于更好地利用Optax进行优化任务。