PyTorch autograd：标量函数多输入下的Jacobian与JVP高效计算

我有一个接收5个输入值并返回标量的函数，映射形式为f:R^5 -> R。因此，它的Jacobian矩阵J维度为(1x5)，也可表示为梯度g形式的行向量。

我可以通过torch.autograd.functional.jacobian轻松计算单个输入x的Jacobian：

J = torch.autograd.functional.jacobian(func=f, inputs=(x[0], x[1], x[2], x[3], x[4]))

我的问题如下：

1. 当其中一个参数存在一系列取值时，多次计算Jacobian（梯度）的最优高效方式是什么？
2. 我不需要每个取值对应的单独梯度，而是希望将JacobianJ与向量v相乘，得到同维度的向量y=Jv。能否使用jvp、vjp或vmap等函数提升代码性能？
3. 我在其他地方看到过Jax的相关提及，它在这类问题上表现是否更出色？

示例代码（当前使用循环/列表推导，希望优化） #

import torch

def f(x0, x1, x2, x3, x4):
        return x0 ** 2 + x1 ** 3 + x2 ** 4 + x3 ** 5 + x4 ** 6

if __name__ == '__main__':
    a = torch.tensor(1.0)
    b = torch.tensor(1.0)
    c = torch.tensor(1.0)
    d = torch.tensor(1.0)
    e = torch.tensor(1.0)

    g = torch.autograd.functional.jacobian(func=f, inputs=(a, b, c, d, e))
    print(g)  
    """ 输出: (tensor(2.), tensor(3.), tensor(4.), tensor(5.), tensor(6.)) """

    x0_values = torch.arange(1.0, 10.0, 1.0)

    g_list = []
    for x0 in x0_values:
        g = torch.autograd.functional.jacobian(func=f, inputs=(x0, b, c, d, e))
        g_list.append(torch.hstack(g))

    J = torch.vstack(g_list)
    print(J)

    """
    输出:
    tensor([[2., 3., 4., 5., 6.],
            [4., 3., 4., 5., 6.],
            [6., 3., 4., 5., 6.],
            [8., 3., 4., 5., 6.],
            [10., 3., 4., 5., 6.],
            [12., 3., 4., 5., 6.],
            [14., 3., 4., 5., 6.],
            [16., 3., 4., 5., 6.],
            [18., 3., 4., 5., 6.]])
    """

问题1：批量计算梯度的高效方式 #

循环单输入会重复构建计算图，带来额外开销，最优方案是向量化输入+PyTorch批量自动微分：

• 修改函数兼容批量输入，利用广播机制处理固定参数；
• 使用torch.autograd.grad替代逐次调用jacobian，针对标量输出的梯度计算更直接，且支持批量处理。

优化代码示例：

import torch

def f(x0, x1, x2, x3, x4):
    return x0 ** 2 + x1 ** 3 + x2 ** 4 + x3 ** 5 + x4 ** 6

if __name__ == '__main__':
    b = torch.tensor(1.0, requires_grad=True)
    c = torch.tensor(1.0, requires_grad=True)
    d = torch.tensor(1.0, requires_grad=True)
    e = torch.tensor(1.0, requires_grad=True)
    x0_values = torch.arange(1.0, 10.0, 1.0, requires_grad=True)

    # 一次性计算批量输出
    outputs = f(x0_values, b, c, d, e)
    # 批量计算所有参数的梯度
    grads = torch.autograd.grad(outputs, [x0_values, b, c, d, e], grad_outputs=torch.ones_like(outputs))
    
    # 拼接成批量Jacobian矩阵
    J = torch.hstack([
        grads[0].reshape(-1,1),
        grads[1].repeat(len(x0_values),1),
        grads[2].repeat(len(x0_values),1),
        grads[3].repeat(len(x0_values),1),
        grads[4].repeat(len(x0_values),1)
    ])
    print(J)

问题2：用jvp/vjp/vmap优化Jv计算 #

若仅需y=Jv，无需显式计算Jacobian矩阵，**JVP（Jacobian-vector product）**是更高效的选择——它避免存储完整Jacobian，计算复杂度更低。

方法1：vmap + jvp批量计算（PyTorch 2.0+支持） #

用vmap将单输入JVP逻辑映射到整个批量，无需手动循环：

import torch
from torch import vmap

def f(x):
    # 整合输入为张量，适配vmap
    return x[0]**2 + x[1]**3 + x[2]**4 + x[3]**5 + x[4]**6

v = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])  # 目标向量v

if __name__ == '__main__':
    # 构造批量输入：x0为取值序列，其余参数固定为1.0
    x_batch = torch.stack([
        torch.arange(1.0,10.0,1.0),
        torch.ones(9),
        torch.ones(9),
        torch.ones(9),
        torch.ones(9)
    ], dim=1)
    
    # 定义单输入JVP逻辑
    def jvp_single(x):
        return torch.autograd.functional.jvp(f, (x,), (v,))[1]
    
    # 批量计算Jv
    y_batch = vmap(jvp_single)(x_batch)
    print(y_batch)

方法2：手动推导（函数形式已知时） #

对于示例函数，Jv可直接用导数公式计算：Jv = 2x0*v0 + 3x1²*v1 +4x2³*v2 +5x3⁴*v3 +6x4⁵*v4，这种方式比自动微分更快，但仅适用于函数形式明确的场景。

问题3：Jax在这类问题上的优势 #

Jax在批量微分、向量化计算上确实有显著优势：

1. 原生vmap支持：vmap是Jax核心特性，对批量计算的适配更自然，可与自动微分、JIT编译无缝结合；
2. JIT编译：jax.jit能将批量微分、JVP等逻辑编译为高效机器码，大规模计算时性能远超PyTorch动态图；
3. 函数式模型：纯函数式设计避免了PyTorch中张量状态管理的开销，更适合无状态的批量微分计算。

如果代码已基于PyTorch构建，迁移Jax需要适配整个栈；小规模计算时PyTorch优化足够，大规模场景下Jax的性能优势会很明显。

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