NumPyro优化器扩展：支持带损失值输入的Optax优化器

在概率编程和变分推断中，学习率的选择对模型收敛至关重要。本文将介绍如何在NumPyro中实现对Optax优化器的扩展支持，特别是那些需要额外输入（如当前损失值）的高级优化器。

背景与挑战

在随机变分推断(SVI)中，学习率的选择直接影响模型的收敛性。传统方法需要手动调整学习率，这不仅耗时，而且在自动化流程（如交叉验证）中难以实现。Optax库提供了reduce_on_plateau等高级优化器，它们能根据损失值自动调整学习率，但这些优化器需要当前损失值作为额外输入参数。

解决方案设计

为了在NumPyro中支持这类优化器，我们需要扩展现有的优化器接口。核心思路是：

1. 创建一个新的优化器基类_NumPyroOptimValueArg，继承自NumPyro的标准优化器类
2. 修改更新方法，使其能够接收并传递损失值参数
3. 实现与Optax优化器的适配层

关键技术实现

优化器基类扩展

新的优化器基类需要重写三个关键方法：

class _NumPyroOptimValueArg(_NumPyroOptim):
    def update(self, g, state, value):
        # 扩展更新方法以接收损失值
        i, opt_state = state
        opt_state = self.update_fn(i, g, opt_state, value=value)
        return i + 1, opt_state
    
    def eval_and_update(self, fn, state, forward_mode_differentiation=False):
        # 计算值并更新，传递损失值
        params = self.get_params(state)
        (out, aux), grads = _value_and_grad(fn, x=params, ...)
        return (out, aux), self.update(grads, state, value=out)
    
    def eval_and_stable_update(self, fn, state, forward_mode_differentiation=False):
        # 稳定性检查版本
        params = self.get_params(state)
        (out, aux), grads = _value_and_grad(fn, x=params, ...)
        out, state = lax.cond(
            jnp.isfinite(out) & ...,
            lambda _: (out, self.update(grads, state, value=out)),
            lambda _: (jnp.nan, state),
            None)
        return (out, aux), state

Optax适配器

创建专门的适配器函数，将Optax优化器转换为NumPyro可用的形式：

def optax_to_numpyro_value_arg(transformation):
    def init_fn(params):
        opt_state = transformation.init(params)
        return params, opt_state
    
    def update_fn(step, grads, state, value):
        params, opt_state = state
        updates, opt_state = transformation.update(
            grads, opt_state, params, value=value)
        updated_params = optax.apply_updates(params, updates)
        return updated_params, opt_state
    
    def get_params_fn(state):
        params, _ = state
        return params
    
    return _NumPyroOptimValueArg(..., init_fn, update_fn, get_params_fn)

应用示例

在实际使用中，我们可以组合多种优化策略。例如，将Adam优化器与基于平台的学习率衰减结合：

optimizer = optax_to_numpyro_value_arg(optax.chain(
    optax.adam(0.01),
    reduce_on_plateau(
        cooldown=100, 
        accumulation_size=100, 
        patience=200
    )
))

svi = SVI(model, guide, optimizer, loss=Trace_ELBO())
svi_result = svi.run(random.PRNGKey(0), 2000, data)

这种组合优化器会在初始阶段使用Adam进行优化，当检测到损失值进入平台期时，自动降低学习率，从而获得更好的收敛效果。

技术优势

1. 自动化学习率调整：无需手动调整学习率，优化器根据损失变化自动调整
2. 兼容性：保持与现有NumPyro代码的兼容性
3. 灵活性：可以组合多种优化策略
4. 稳定性：内置数值稳定性检查

总结

通过扩展NumPyro的优化器接口，我们实现了对需要额外输入的Optax优化器的支持。这一改进特别适合变分推断场景，能够自动调整学习率，提高模型收敛的稳定性和效率。开发者现在可以更方便地使用高级优化策略，而无需担心底层实现细节。

这种设计模式也展示了如何将JAX生态中的优秀组件（如Optax）无缝集成到概率编程框架中，为更复杂的模型优化提供了可能。