4、自动微分深入：grad、vmap、pmap的协同、高阶梯度计算、Hessian矩阵的求法、自定义VJP规则

自动微分是JAX的灵魂。说实话，我当年从PyTorch转到JAX时，最不适应的就是这套东西。但用顺手之后，真香。

这一章我们聊点硬的。grad、vmap、pmap怎么配合？高阶梯度怎么算？Hessian矩阵怎么优雅地求？还有自定义VJP——这个我踩过不少坑。

4.1 grad、vmap、pmap的协同作战

先说说这三兄弟各自干啥：

• grad：求梯度。这个最常用。
• vmap：自动向量化。把标量函数变成批量处理。
• pmap：并行计算。把计算分布到多个设备上。

它们可以组合使用。我刚开始觉得这没啥，直到有一次做大规模参数估计，才体会到组合的威力。

4.1.1 grad + vmap 组合

假设你有一个函数，输入是单个样本，输出是loss。你想对一批样本分别求梯度。怎么做？

import jax
import jax.numpy as jnp

def loss_fn(params, x):
    return jnp.sum(params * x)

params = jnp.array([1.0, 2.0, 3.0])
batch_x = jnp.array([[1.0, 2.0, 3.0],
                     [4.0, 5.0, 6.0],
                     [7.0, 8.0, 9.0]])

# 对单个样本求梯度
grad_fn = jax.grad(loss_fn)
# 用vmap批量处理
batch_grad_fn = jax.vmap(grad_fn, in_axes=(None, 0))

grads = batch_grad_fn(params, batch_x)
print(grads)
# 输出：每个样本对应的梯度

这里in_axes=(None, 0)的意思是：第一个参数（params）不批量，第二个参数（x）沿着第0维批量。嗯，这个细节容易搞错，我刚开始就经常写反。

4.1.2 grad + pmap 组合

当你的模型太大，单卡装不下时，pmap就派上用场了。我做过一个NLP模型，参数量10亿+，单卡根本跑不动。

# 假设有4张GPU
def loss_fn(params, x):
    return jnp.sum(params * x)

# 先求梯度
grad_fn = jax.grad(loss_fn)
# 再并行
parallel_grad_fn = jax.pmap(grad_fn, axis_name='devices')

# 数据需要分片
params_sharded = jax.device_put_replicated(params, jax.devices())
batch_sharded = jax.device_put_sharded(
    jnp.split(batch_x, len(jax.devices())), 
    jax.devices()
)

grads = parallel_grad_fn(params_sharded, batch_sharded)

注意：pmap要求所有输入在设备间是分片的。我犯过一个低级错误——忘了把参数复制到所有设备上，结果梯度全乱了。

4.1.3 三兄弟一起上

最骚的操作是三个一起用。vmap处理批量，pmap处理多设备，grad求梯度。

# grad + vmap + pmap 三合一
def loss_fn(params, x):
    return jnp.sum(params * x)

grad_fn = jax.grad(loss_fn)
batch_grad_fn = jax.vmap(grad_fn, in_axes=(None, 0))
parallel_batch_grad_fn = jax.pmap(batch_grad_fn, axis_name='devices')

# 数据准备
params_sharded = jax.device_put_replicated(params, jax.devices())
# 每个设备处理一部分batch
batch_per_device = batch_x.reshape(len(jax.devices()), -1, batch_x.shape[-1])
batch_sharded = jax.device_put_sharded(batch_per_device, jax.devices())

result = parallel_batch_grad_fn(params_sharded, batch_sharded)

4.2 高阶梯度计算

高阶梯度，说白了就是梯度的梯度。这在一些优化算法和物理模拟中很常见。

4.2.1 二阶梯度

JAX的grad可以嵌套使用。你想求二阶导？那就grad两次。

def f(x):
    return x**3 + 2*x**2 + x

# 一阶导
df_dx = jax.grad(f)
print(df_dx(2.0))  # 3*4 + 4*2 + 1 = 21

# 二阶导
d2f_dx2 = jax.grad(jax.grad(f))
print(d2f_dx2(2.0))  # 6*2 + 4 = 16

# 三阶导
d3f_dx3 = jax.grad(jax.grad(jax.grad(f)))
print(d3f_dx3(2.0))  # 6

你想想看，这其实就是链式法则的递归应用。JAX在底层用了一套叫「线性化」的机制，每次grad调用都会构建一个计算图。

4.2.2 混合高阶梯度

有时候你需要对不同的参数求不同阶的梯度。比如在GAN训练中，生成器和判别器的更新频率不同。

def loss_fn(params_gen, params_disc, x):
    # 假设这是某个损失函数
    return jnp.sum(params_gen * x) + jnp.sum(params_disc * x**2)

# 对生成器参数求一阶导，对判别器参数求二阶导
grad_gen = jax.grad(loss_fn, argnums=0)
grad_disc = jax.grad(jax.grad(loss_fn, argnums=1), argnums=1)

params_gen = jnp.array([1.0, 2.0])
params_disc = jnp.array([3.0, 4.0])
x = jnp.array([5.0, 6.0])

print(grad_gen(params_gen, params_disc, x))
print(grad_disc(params_gen, params_disc, x))

4.3 Hessian矩阵的求法

Hessian矩阵是二阶偏导数的矩阵。在优化理论中，它描述了函数的曲率。说白了，就是告诉你梯度变化得有多快。

4.3.1 用grad求Hessian

最直接的方法：对grad再求一次grad。但注意，grad返回的是标量函数的梯度向量，所以第二次grad得到的是雅可比矩阵，也就是Hessian。

def f(x):
    return x[0]**2 + 3*x[0]*x[1] + x[1]**2

# 方法1：grad嵌套
hessian_f = jax.jacfwd(jax.grad(f))
x = jnp.array([1.0, 2.0])
print(hessian_f(x))
# 输出：[[2, 3], [3, 2]]

这里用了jacfwd而不是grad。为什么？因为grad要求输出是标量，而grad(f)的输出是向量，所以第二次求导要用jacfwd或jacrev。

4.3.2 用jacfwd和jacrev组合

JAX提供了更灵活的方式：

# 方法2：jacfwd + jacrev
hessian_f = jax.jacfwd(jax.jacrev(f))
print(hessian_f(x))

# 方法3：直接使用hessian函数
from jax import hessian
hessian_f = hessian(f)
print(hessian_f(x))

三种方法结果一样。我个人习惯用hessian函数，代码更简洁。但如果你需要控制计算顺序（比如正向模式还是反向模式），那就用jacfwd+jacrev组合。

4.3.3 Hessian向量积（HVP）

实际应用中，你很少需要完整的Hessian矩阵。更常见的是Hessian向量积（HVP）。比如在共轭梯度法中。

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

def hvp(f, x, v):
    return jax.jvp(jax.grad(f), (x,), (v,))[1]

x = jnp.array([1.0, 2.0, 3.0])
v = jnp.array([0.1, 0.2, 0.3])
print(hvp(f, x, v))
# 输出：array([0.2, 0.4, 0.6])

HVP的计算复杂度是O(n)，而完整Hessian是O(n²)。对于大模型，这差别就是天壤之别。

4.4 自定义VJP规则

VJP（Vector-Jacobian Product）是反向模式自动微分的核心。说白了，就是反向传播时算的那个东西。

有时候JAX内置的自动微分不够用。比如你有一个自定义操作，或者你想手动控制梯度计算。这时候就需要自定义VJP。

4.4.1 基本用法

from jax import custom_vjp

@custom_vjp
def my_op(x):
    return x**2

# 定义前向传播
def f_fwd(x):
    return my_op(x), x  # 保存中间结果

# 定义反向传播
def f_bwd(res, g):
    x = res
    # 手动计算梯度：d(x^2)/dx = 2x
    return (2 * x * g,)

my_op.defvjp(f_fwd, f_bwd)

# 测试
x = jnp.array(3.0)
grad_fn = jax.grad(my_op)
print(grad_fn(x))  # 输出：6.0

这里的关键是：前向传播要返回一个「残差」（res），反向传播要用到这个残差。我刚开始写的时候，经常忘记保存中间结果，导致反向传播算不出来。

4.4.2 复杂一点的例子

假设你有一个操作，前向传播很简单，但反向传播需要特殊处理。比如一个带约束的归一化操作。

@custom_vjp
def constrained_norm(x, threshold):
    norm = jnp.linalg.norm(x)
    return jnp.where(norm > threshold, x * threshold / norm, x)

def norm_fwd(x, threshold):
    norm = jnp.linalg.norm(x)
    out = jnp.where(norm > threshold, x * threshold / norm, x)
    return out, (x, norm, threshold)

def norm_bwd(res, g):
    x, norm, threshold = res
    # 手动推导梯度
    # 如果norm <= threshold，梯度就是单位矩阵
    # 如果norm > threshold，梯度需要特殊处理
    def grad_when_clipped(x, norm, threshold, g):
        # 这里省略复杂的推导过程
        return g * threshold / norm
    
    grad_x = jnp.where(norm > threshold, 
                       grad_when_clipped(x, norm, threshold, g), 
                       g)
    grad_threshold = jnp.array(0.0)  # threshold不需要梯度
    return (grad_x, grad_threshold)

constrained_norm.defvjp(norm_fwd, norm_bwd)

4.4.3 自定义VJP的常见坑

• 残差太大：保存的中间结果太多，导致内存爆炸。只保存必要的。
• 梯度形状不匹配：反向传播返回的梯度形状必须和输入形状一致。
• 非可微操作：如果操作中有不可微的部分（比如argmax），需要手动处理。

我曾经在一个项目中自定义了一个排序操作的VJP。排序本身不可微，但我们可以用「软排序」的梯度来近似。这个坑我爬了整整两天。

知识体系总览

下面这张图总结了本章的核心逻辑：

这张图展示了自动微分的四个核心方向。它们不是孤立的，实际项目中经常需要组合使用。比如我在做物理模拟时，既要用到高阶梯度，又要用vmap批量处理，还得自定义VJP处理边界条件。

好了，这一章的内容就到这里。记住：自动微分是工具，不是目的。理解背后的原理，才能用好它。