JAX 高阶偏导数计算技巧与方法

本文介绍在 JAX 中避免全阶导数计算、直接获取特定混合偏导数（如三阶导 ∂³f/∂x₀∂x₂∂x₄）的两种高效方法：基于 JVP 的链式方向导数法和解构参数的 argnums 分离法，显著提升高维向量函数高阶微分的计算效率。

在科学计算与可微编程中，常需对向量输入函数（如 f(x) = ∏ᵢ xᵢ）计算特定索引组合的高阶偏导数，例如三阶导数 ∂³f/∂x₀∂x₂∂x₄。若直接使用 jax.jacobian(jax.jacobian(jax.grad(f)))，JAX 会完整构建维度为 (n, n, n) 的三阶张量（n=len(x)），造成大量冗余计算与内存开销——尤其当 n=5 时仍需生成 125 个元素，而目标仅是其中 1 个切片。

幸运的是，JAX 提供了更轻量、更精准的替代方案。核心思想是：将“对某坐标求偏导”转化为“沿对应单位向量的方向导数（JVP）”，从而跳过无关维度的梯度传播。

✅ 方法一：链式 JVP（推荐 —— 保持单参数签名）

利用 jax.jvp（Jacobian-vector product）逐层施加方向向量，每次仅激活一个坐标轴：

import jax
import jax.numpy as jnp

def f(x):
    return jnp.prod(x)

def deriv(f, x, v):
    """计算 f 在 x 处沿方向 v 的一阶方向导数"""
    return jax.jvp(f, (x,), (v,))[1]

def one_hot(i, size):
    """生成 size 维 one-hot 向量，第 i 位为 1"""
    return jnp.zeros(size).at[i].set(1.0)

# 构造目标三阶导数：∂³f/∂x₀∂x₂∂x₄
size = 5
df_x0     = lambda x: deriv(f, x, one_hot(0, size))      # ∂f/∂x₀
df2_x0_x2 = lambda x: deriv(df_x0, x, one_hot(2, size))  # ∂²f/∂x₀∂x₂
df3_x0_x2_x4 = lambda x: deriv(df2_x0_x2, x, one_hot(4, size))  # ∂³f/∂x₀∂x₂∂x₄

# 批量计算（支持 vmap）
x_batch = jnp.array([
    [ 1.,  2.,  3.,  4.,  5.],
    [ 6.,  7.,  8.,  9., 10.],
    [11., 12., 13., 14., 15.],
    [16., 17., 18., 19., 20.],
    [21., 22., 23., 24., 25.],
    [26., 27., 28., 29., 30.]
])

result = jax.vmap(df3_x0_x2_x4)(x_batch)
print(result)  # [  8.  63. 168. 323. 528. 783.]

✅ 优势：完全保留 f(x) 的单参数签名；计算复杂度从 O(n³) 降至 O(1)（固定 3 次 JVP）；内存占用恒定。
⚠️ 注意：确保 one_hot 向量与 x 的 shape[1] 对齐；所有中间函数必须可被 jax.jit 或 vmap 转换（本例满足）。

✅ 方法二：argnums 解构（适用于坐标天然分离场景）

若允许将向量输入显式展开为独立标量参数，可借助 argnums 精确指定每次微分作用的参数位置：

def f_explicit(x0, x1, x2, x3, x4):
    return x0 * x1 * x2 * x3 * x4

# 逐层指定求导变量：先对 x4，再对 x2，最后对 x0
df_dx4     = jax.grad(f_explicit, argnums=4)
df2_dx4_dx2 = jax.jacobian(df_dx4, argnums=2)
df3_dx4_dx2_dx0 = jax.jacobian(df2_dx4_dx2, argnums=0)

# 将批量数据转置后解包为独立参数
x_T = x_batch.T  # shape: (5, 6)
result_v2 = df3_dx4_dx2_dx0(*x_T)  # 自动广播
print(result_v2)  # 同样输出 [  8.  63. 168. 323. 528. 783.]

✅ 优势：语义清晰，JAX 编译器可极致优化；无需手动构造方向向量。
⚠️ 限制：要求输入维度固定且已知；函数签名改变，不适用于动态长度向量。

总结与选型建议

? 底层原理：JAX 的 jvp 实质是前向模式自动微分，天然适合“单方向、多层链式”求导；而 jacobian 默认使用反向模式（vjp），适合“多输出、单输入”，但全张量计算代价高昂。二者结合，方得高效之钥。

无论选择哪种方式，都应配合 jax.jit 进一步加速（尤其在循环或多次调用时）：

df3_jitted = jax.jit(df3_x0_x2_x4)
result_fast = jax.vmap(df3_jitted)(x_batch)