使用Jax/Jaxopt/SciPy Minimize时优化变量x未更新，求问题原因

问题描述 #

需要最小化一个目标函数，优化参数x为(n,m)维度的NumPy数组，目标函数通过调用外部Java API包装的calculate_normX计算向量范数，核心代码如下：

# 初始值
normX0 = calculate_normX(x_start)

def objective(x) -> float:
    """目标函数"""
    x = x.reshape((n,m))
    normX = calculate_normX(x)
    return -(float(normX) / float(normX0))

分别尝试两种优化方案后，均出现优化过程中x无变化的问题：

• Jaxopt NonlinearCG实现：

solver = NonlinearCG(fun=objective, maxiter=5, verbose=True)
res = solver.run(x.flatten())

• SciPy L-BFGS-B结合Jax自动微分实现：

objective_jac = jax.jacrev(objective)
minimize(objective, jac=objective_jac, x0=x.flatten(), method='L-BFGS-B', options={'maxiter': 2})

更换初始随机值或其他求解器后，问题仍然存在。

核心原因 #

问题根源在于calculate_normX这个外部Java API包装函数：

• Jax的自动微分机制（包括jax.jacrev）仅支持追踪Jax原生操作，非Jax实现的外部函数会被视为常数函数——即对输入x的导数恒为0。
• 优化器依赖梯度信息更新参数，当梯度全为0时，优化器会判定当前点已是极值点，不会对x执行任何更新操作。

解决方案 #

1. 替换为Jax原生范数计算函数 #

如果calculate_normX计算的是标准向量范数（如L2、L1范数），直接使用Jax内置的jax.numpy.linalg.norm替代，确保函数完全可微分：

import jax.numpy as jnp

def objective(x) -> float:
    x = x.reshape((n,m))
    normX = jnp.linalg.norm(x)  # 替换为Jax原生范数计算
    return -(float(normX) / float(normX0))

2. 手动实现梯度（保留外部API） #

若无法替换外部API，需手动推导并实现目标函数的梯度，再传给优化器。以L2范数为例，目标函数的梯度为-x / (normX0 * normX)，实现代码如下：

def objective_jac(x):
    x_reshaped = x.reshape((n,m))
    normX = calculate_normX(x_reshaped)
    grad = -x_reshaped.flatten() / (normX0 * normX)
    return grad.astype(float)

# SciPy优化时使用手动梯度
minimize(objective, jac=objective_jac, x0=x.flatten(), method='L-BFGS-B', options={'maxiter': 2})

3. 用Jax自定义JVP包装外部函数 #

若必须保留外部API，可通过jax.custom_jvp为其添加可微分支持，以L2范数为例：

from jax import custom_jvp, jnp

# 包装外部API
@custom_jvp
def calculate_normX_jax(x):
    return float(calculate_normX(x))

# 自定义JVP（Jacobian-Vector Product）
@calculate_normX_jax.defjvp
def calculate_normX_jax_jvp(primals, tangents):
    x, = primals
    x_dot, = tangents
    normX = calculate_normX(x)
    primal_out = float(normX)
    tangent_out = float(jnp.dot(x.flatten(), x_dot.flatten()) / normX)  # L2范数的JVP逻辑
    return primal_out, tangent_out

# 更新目标函数
def objective(x) -> float:
    x = x.reshape((n,m))
    normX = calculate_normX_jax(x)
    return -(normX / normX0)

此时Jax的自动微分机制可正确识别梯度，优化器能正常更新参数。

内容的提问来源于stack exchange，提问作者agentsmith