大多数用过 PyTorch 的人对loss.backward()这套反向传播都很熟但一旦被问到雅可比向量积怎么算很多人第一反应是打开torch.autograd.functional.jacobian把整个雅可比矩阵求出来再手动做一次矩阵乘法。这种做法在小模型上没毛病一旦输出维度上千显存和耗时立刻教做人。雅可比向量积Jacobian-Vector Product简称 JVP存在的意义就是你根本不需要显式构造那个 m×n 的雅可比矩阵也能拿到它和一个向量的乘积代价只有一次前向传播的量级。这篇内容围绕 JVP 在 PyTorch 里的几种实现路径展开重点拆torch.func.jvp的用法、torch.autograd.functional.jvp的性能陷阱、以及踩过的坑。适合已经会写autograd.grad、想进一步搞清楚前向模式自动微分、并且需要在二阶优化、雅可比正则、灵敏度分析这些场景里真正用起来的人。1. 把 JVP 从矩阵里拆出来谁乘在左边谁乘在右边1.1 一个 3×2 的例子先把概念钉死先不谈自动微分回到线性代数本身。假设有一个函数 f: R² → R³在点 x 处它的雅可比矩阵 J 是一个 3 行 2 列的矩阵。JVP 就是拿一个 2 维向量 v从右边乘上去得到 J·v结果是一个 3 维向量。注意 v 的维度必须和输入维度一致结果的维度必须和输出维度一致。反过来VJPVector-Jacobian Product是拿一个 3 维向量 u 从左边乘uᵀ·J结果是一个 2 维向量维度跟着输入走。这两个操作的形状规则完全不同这是后面所有坑的根源。拿一个具体例子import torch def f(x): return torch.stack([x[0] * x[1], x[1] ** 2, torch.sin(x[0])]) x torch.tensor([2.0, 3.0], dtypetorch.float64)手工求偏导J 应该是J [[x1, x0 ], [0, 2*x1 ], [cos(x0), 0 ]]代入 x [2, 3]得到[[3, 2], [0, 6], [cos(2), 0]]其中 cos(2) ≈ -0.4161。现在取 v [1, 0]J·v 就是 J 的第一列即[3, 0, -0.4161]。取 v [0, 1]得到第二列[2, 6, 0]。这两个结果后面用来验证代码是否正确。1.2 为什么 PyTorch 默认只给你 VJPloss.backward()背后跑的是反向模式自动微分它天然产出的是 VJP不是 JVP。原因很实际深度学习里绝大多数函数是从一个巨大的参数向量映射到一个标量损失也就是 f: R^N → R¹。这种情况下输出维度 m 1只需要一次反向传播就能拿到完整的梯度向量代价极低。而如果换成 JVP 的思路你得做 N 次前向传播才能拿全梯度完全不可行。所以框架的设计重心一直压在 VJP 上前向模式自动微分在很长一段时间里属于有但不太好用的状态。这直接导致一个现象当你遇到 f: R^n → R^m 且 n 很小、m 很大的场景时比如一个几十维的控制参数映射到上千维的仿真输出用反向模式硬算会非常亏而这恰恰是 JVP 的主场。用一个表把两种模式的成本关系摆清楚这张表值得贴在显示器边上目标操作需要调用几次适合的形状PyTorch 对应接口完整雅可比逐列前向n 次前向n 小 m 大torch.func.jacfwd完整雅可比逐行反向m 次反向m 小 n 大torch.func.jacrev单次 JVPJ·v1 次前向量级输入小torch.func.jvp单次 VJPuᵀ·J1 次反向量级输出小torch.func.vjp/autograd.grad1.3 前向模式到底在算什么dual number 的直觉前向模式自动微分的实现方式叫对偶数dual number。把每个标量 a 扩展成 a b·ε 的形式规定 ε ≠ 0 但 ε² 0。任何函数作用上去之后实部就是函数值ε 的系数就是方向导数。举个手算的例子x 2 1·ε那么 x² 4 4ε ε² 4 4εε 的系数 4 正好是 2x 在 x2 处的导数乘上方向 1。这个机制的好处是一次前向传播就能同时算出函数值和沿着某个方向 v 的方向导数而且天然支持任意深的网络不需要保存中间激活。代价是每个标量都要额外存一份 ε 系数显存占用大约翻倍但和前向传播同一量级。PyTorch 里这套机制的底层入口是torch.autograd.forward_adtorch.func.jvp是它的高层封装。理解这一层很关键因为后面会看到torch.autograd.functional.jvp走的根本不是这条路。2. PyTorch 里算出 JVP 的四条路各自的真实代价2.1 torch.func.jvp目前最该用的那条PyTorch 2.0 之后torch.func命名空间下提供了一套函数式变换 APIjvp是其中之一。它的签名是jvp(func, primals, tangents)primals是函数位置参数的元组tangents是切线方向的元组两者结构必须一致。返回值是一个二元组(output, jvp_out)output是函数原本的输出jvp_out是 J·v 的结果。import torch from torch.func import jvp def f(x): return torch.stack([x[0] * x[1], x[1] ** 2, torch.sin(x[0])]) x torch.tensor([2.0, 3.0], dtypetorch.float64) v torch.tensor([1.0, 0.0], dtypetorch.float64) out, jvp_out jvp(f, (x,), (v,)) print(out) # tensor([ 6.0000, 9.0000, 0.9093], dtypetorch.float64) print(jvp_out) # tensor([ 3.0000, 0.0000, -0.4161], dtypetorch.float64)跑出来和手算的[3, 0, -0.4161]完全对上。注意两个细节第一x本身不需要requires_gradTrue前向模式不依赖计算图第二f内部不允许有原地操作in-place比如x[0] 1这种写法会直接报错因为 dual tensor 的 ε 分量不允许被就地覆盖。2.2 torch.autograd.functional.jvp名字一样实现完全不同很多人第一次搜到的是torch.autograd.functional.jvp。它确实能用签名也几乎一样from torch.autograd.functional import jvp out, jvp_out jvp(f, (x,), (v,), create_graphFalse, strictFalse)但这里有个非常容易被忽略的事实这个函数内部是用反向模式实现的。它的做法是先构造出 VJP 的形式再对切向量 v 求一次导本质上是 reverse-over-reverse也就是两次反向传播拼出来的 JVP。这意味着什么意味着当你的函数是输入维度小、输出维度大时torch.autograd.functional.jvp反而会比真正的torch.func.jvp慢一大截因为它要沿着输出维度反向走一遍。我在一个 50 维输入、2000 维输出的仿真函数上做过对比torch.func.jvp耗时大约是functional.jvp的三分之一差距相当明显。提示torch.autograd.functional.jvp的create_graphTrue是它为数不多的优势场景——需要 JVP 结果本身再参与求导比如构造二阶以上的混合导数时它比函数式 API 更省事。日常只要一阶 JVP优先torch.func.jvp。2.3 手搓 double-backward 与有限差分第三条路是老派做法算两次autograd.grad。这不是真正的 JVP而是构造 VJP 再套一层def manual_jvp(f, x, v): x x.detach().requires_grad_(True) y f(x) # 先拿一个哑元的 VJP把 v 塞进 grad_outputs g torch.autograd.grad(y, x, grad_outputstorch.zeros_like(y), create_graphTrue) ...说实话这条路现在没什么推荐价值了代码绕、易错、还容易在create_graph上翻车除非你在维护一份必须兼容老版本 PyTorch 的代码库否则直接跳过。第四条路是有限差分(f(x h*v) - f(x)) / h。它的缺点写在教科书里——截断误差和舍入误差互相拉扯h 太小结果被噪声吃掉h 太大非线性误差明显。不过它有个独特价值作为数值校验的金标准。你写完任何 JVP 实现拿有限差分对一遍能立刻发现是不是方向搞反了、是不是 v 的维度用错了。我习惯用h 1e-6配 float64 做对照误差在 1e-6 量级就算通过。2.4 四条路的横向对比路径是否真前向模式相对耗时需要 requires_grad支持二阶推荐场景torch.func.jvp是1x基线否配合create_graph式组合首选尤其 n 小时的批量 JVPtorch.autograd.functional.jvp否反向两次2~4x否strictFalse原生支持需要高阶混合导手搓autograd.grad否3x 以上是麻烦老版本兼容有限差分近似2x否不支持数值校验3. torch.func.jvp 的实战用法从单方向到整块雅可比3.1 参数结构不匹配是最常见的报错来源jvp(func, primals, tangents)这三者的结构约束很严格。primals和tangents必须是元组长度等于func的位置参数个数每个位置的张量形状必须逐一对上。如果你的函数是def f(a, b)那就得写jvp(f, (a, b), (va, vb))漏一个或者包错层级都会报形状错误。def g(a, b): return (a * b).sum() b.pow(2).sum() a torch.randn(3, dtypetorch.float64) b torch.randn(4, dtypetorch.float64) va torch.ones(3, dtypetorch.float64) vb torch.zeros(4, dtypetorch.float64) out, jvp_out jvp(g, (a, b), (va, vb)) # 输出是标量对 a 方向的偏导之和这里vb传全零意思是只要 a 方向的导数。这是一个非常实用的技巧当你只关心某一个输入的分量时其余位置传零张量即可前向模式会自动跳过它们对结果的贡献严格说是它们的 ε 分量为零不影响计算但也不省时间只是省得你写新函数。如果函数的输出不是一个张量而是一个元组jvp的返回也会是同结构的元组一一对应。这一点比torch.autograd.grad友好得多后者在输出为元组时经常让人抓狂。3.2 用 vmap 把逐方向的 JVP 拼成完整雅可比单次 JVP 只给一列。想拿完整雅可比把标准基向量逐个喂进去就行但 Python 循环太慢正确做法是vmapfrom torch.func import jvp, vmap def full_jacobian_via_jvp(f, x): n x.numel() basis torch.eye(n, dtypex.dtype, devicex.device) # 每次调用返回 J e_i即雅可比的第 i 列 cols vmap(lambda v: jvp(f, (x,), (v,))[1])(basis) return cols.T # 转置成标准的 (m, n) 形状 J full_jacobian_via_jvp(f, x) print(J) # tensor([[ 3.0000, 2.0000], # [ 0.0000, 6.0000], # [-0.4161, 0.0000]], dtypetorch.float64)结果和手算完全一致。这里有个容易搞反的地方vmap出来后cols的形状是(n, m)第 i 行是第 i 列所以必须转置。我第一次写的时候忘了转后面所有矩阵乘法全错排查了半天才发现是形状问题。实际上官方已经把这个封装好了torch.func.jacfwd干的就是这件事from torch.func import jacfwd J jacfwd(f)(x) # 直接得到 (m, n)jacfwd内部就是 vmap jvp性能基本一致能直接用就用它少写一堆胶水代码。3.3 jacfwd 和 jacrev按形状选方向别选反这一对函数经常被混用但它们的成本差异可以到一个数量级。规则很简单jacfwd前向模式需要 n 次前向量级的计算适合n m输入维度小于输出维度jacrev反向模式需要 m 次反向量级的计算适合m n输出维度小于输入维度import time x_small torch.randn(10, dtypetorch.float64) x_large torch.randn(2000, dtypetorch.float64) def f_wide(x): # 10 - 2000输入小输出大 return torch.sin(x).repeat(200).reshape(-1)[:2000] * 1.5 def f_tall(x): # 2000 - 10输入大输出小 return torch.sin(x).reshape(200, 10).sum(0)对f_widejacfwd只要 10 次前向jacrev要 2000 次反向差距悬殊。对f_tall正好相反。跑一下就能直观感受到我这里实测f_wide上jacfwd比jacrev快两个数量级还多f_tall上则反过来。注意不要无脑用jacrev也不要无脑用jacfwd。判断依据只有一条——比较输入元素总数和输出元素总数谁小就用谁做方向。这是前向/反向模式选择的唯一标准和函数长什么样没关系。还有一个隐蔽的点jacfwd在输入维度超过几千时因为要一次性 vmap 出几千个方向的 dual tensor显存会爆。这时候该用chunk_size分批或者干脆切回jacrev。3.4 手工构造 dual tensor 的底层写法有时候你需要绕开torch.func的封装直接操作 dual level比如需要在自定义 autograd Function 内部嵌一段前向模式逻辑。底层写法是这样from torch.autograd.forward_ad import make_dual, unpack_dual, dual_level x torch.tensor([2.0, 3.0], dtypetorch.float64) v torch.tensor([1.0, 0.0], dtypetorch.float64) with dual_level(): x_dual make_dual(x, v) y_dual f(x_dual) y, jvp_out unpack_dual(y_dual) print(jvp_out) # tensor([ 3.0000, 0.0000, -0.4161])dual_level()是必须的上下文管理器出了这个 with 块之后 dual tensor 就失效了再拿去做运算是未定义行为。另外注意只有支持前向模式的算子才能在里面用像某些第三方扩展的自定义算子如果没注册 forward AD 规则会抛RuntimeError。遇到这种要么退回torch.autograd.functional.jvp要么给算子补一份jvp规则。4. 踩过的坑让 JVP 结果对不上的七个原因4.1 原地操作和 BatchNorm 的 training 状态torch.func.jvp对原地操作零容忍。x 1、x.relu_()、x.masked_fill_()这类都会直接报错因为在给 dual tensor 覆写 ε 分量时前向模式无法确定正确的语义。解决办法是全局搜一遍下划线结尾的算子改成非原地版本。更隐蔽的是 BatchNorm。model.train()状态下 BN 会更新 running_mean/running_var这属于副作用side effect函数式变换默认禁止副作用。要么切model.eval()要么用torch.func提供的functional_call把 buffer 显式传进去。我遇到过一次 JVP 结果每次调用都不一样的诡异情况查了两小时才发现是 BN 在偷偷更新统计量。4.2 create_graph 开着然后拿结果去 backward这是最高频的翻车点。假设你这么写_, jv jvp(f, (x,), (v,)) jv.sum().backward() # RuntimeError 或者结果全 None因为jvp默认不构造可微图jv是叶子张量backward 拿不到梯度。想对 JVP 的结果再求导得用torch.func的变换组合方式from torch.func import jvp, grad # 求 d/dx [J(x) v]即二阶信息 def hvp_scalar(x, v): _, jv jvp(grad(lambda t: f(t).sum()), (x,), (v,)) return jv.sum()也就是说把jvp和高阶变换组合起来而不是指望 autograd 的图自动接上。这是函数式 API 和传统 autograd 在思维方式上最大的差别。4.3 随机性、dropout 和数值精度模型里有 dropout 时jvp的两次内部调用一次算输出、一次算切线如果各自采样一份 mask结果就会莫名其妙。torch.func系列变换提供了randomness参数可以设成differentiate或error后者会在检测到随机算子时直接报错逼你显式处理。我的习惯是调试阶段一律设error把问题提前暴露出来。精度方面JVP 涉及大量累加float32 在方向导数上经常只留 4 到 5 位有效数字。做数值校验时务必切到 float64否则有限差分和 JVP 的差异会被舍入噪声盖住你会误以为代码写错了其实是精度不够。4.4 一个完整的排查链路我把自己遇到的一次典型问题完整记录下来这个排查思路可以复用。现象一个 8 维输入、512 维输出的函数torch.func.jvp的结果和有限差分的差异稳定在 3% 左右不是随机噪声而是系统性偏差。第一步把输入换成 float64 重跑差异降到 0.1%说明有大半是精度问题但仍有残差。第二步把函数拆到最简只留一层线性变换此时 JVP 和手算完全一致说明前向模式本身没问题。第三步逐层加回发现加上某个归一化层后偏差出现。第四步检查那个层发现它在内部对输入做了clamp(min0)而测试点 x 的分量有负数clamp 在 0 处不可导有限差分跨过了这个折点于是两边不一致。第五步把测试点挪到所有分量都为正的区域差异消失。结论是JVP 本身没错是有限差分在不可导点附近失效。这个教训很值钱——做数值校验时先确认测试点远离 ReLU、abs、clamp 这类折点否则校验工具本身就在说谎。5. JVP 真正常用的四个场景5.1 Hessian-vector product二阶优化的基石海森矩阵是 n×n参数量上亿时根本存不下但你往往只需要 H·v 这个乘积。而 H·v 恰好等于梯度的 JVPimport torch from torch.func import jvp, grad def loss_fn(x): return (x ** 3).sum() x torch.tensor([1.0, 2.0, 3.0], dtypetorch.float64) v torch.tensor([1.0, 1.0, 1.0], dtypetorch.float64) _, hv jvp(grad(loss_fn), (x,), (v,)) print(hv) # tensor([ 6., 12., 18.])正好是 6x 逐元素乘 v手算验证loss Σx³梯度是 3x²海森是对角阵 diag(6x)H·v 6x·v [6, 12, 18]。完全对上。这条式子在整个二阶优化领域是基础设施。共轭梯度法解 H·p -g 需要反复算 H·v信赖域方法、自然梯度、K-FAC 的近似版本全都建立在这个乘积上。用传统写法torch.autograd.grad(g, x, grad_outputsv)也能算但那个要求 g 必须由create_graphTrue得到显存开销大jvp(grad(f), ...)走前向模式显存友好得多。5.2 雅可比正则与输入梯度惩罚对抗训练和鲁棒性研究里常加一项雅可比正则惩罚模型对输入变化的敏感度形式是 ||∂f/∂x||²。直接算完整雅可比代价太高用 Hutchinson 估计def jacobian_frobenius_sq(model, x, num_samples4): 估计 ||J||_F^2x 是单个样本 def f_single(xi): return model(xi.unsqueeze(0)).squeeze(0) total 0.0 for _ in range(num_samples): v torch.randint(0, 2, x.shape, dtypex.dtype, devicex.device) * 2 - 1 _, jv jvp(f_single, (x,), (v,)) total total (jv ** 2).sum() return total / num_samples原理是 E[||J·v||²] ||J||_F²只要 v 的分量是零均值、单位方差的随机向量Rademacher 分布 ±1 就够。四次采样的估计就足够稳定成本只有四次前向。这个技巧在很多需要控制 Lipschitz 常数的工作里都是标配。5.3 灵敏度分析与参数扫描工程仿真里经常有这种需求一个模型有 30 个可调参数输出是 5000 个测点的响应曲线想知道参数动一点点输出怎么变。这就是标准的 JVP 场景输入小输出大反向模式完全使不上劲。def simulate(params): # params: (30,) 返回: (5000,) ... params torch.randn(30, dtypetorch.float64) direction torch.zeros(30, dtypetorch.float64) direction[3] 1.0 # 只看第 4 个参数的单独影响 _, sensitivity jvp(simulate, (params,), (direction,))这里的sensitivity就是第 4 个参数单独变动一个单位时5000 个测点各自的响应。用 VJP 做同样的事得对 5000 个输出各反向一次成本差两个数量级。神经 ODE 的伴随方法里也大量依赖这类操作只不过那边用的是 VJP因为目标通常是标量损失。方向选择永远跟着形状走这一条不会变。5.4 用 trace 估计代替完整雅可比有时候你只想要雅可比的迹trace比如在流模型里估算对数行列式的一阶近似。迹可以用 Hutchinson 直接估def trace_hutchinson(f, x, num_samples8): n x.numel() total 0.0 for _ in range(num_samples): v torch.randint(0, 2, (n,), dtypex.dtype, devicex.device) * 2 - 1 _, jv jvp(f, (x,), (v,)) total total (v * jv).sum() return total / num_samples因为 E[vᵀJv] tr(J)。对一个 100 维的对角函数八次采样就能把迹估到 5% 以内的相对误差而完整雅可比要 100 次前向。这个差距在参数维度上千时会变得不可接受。6. 性能与工程取舍什么时候不该用 JVP6.1 批量化vmap 一次比循环一百次强torch.func.jvp真正的威力在于它能被vmap包裹把一百个方向变成一次核函数调用。手写 for 循环的话每个方向都要重新走一遍算子调度Python 开销直接吃掉所有收益。实测在一个中等规模的网络上vmap 版本比 Python 循环快 20 倍以上。如果不想用vmapPyTorch 还提供了另一条路torch.autograd.grad的is_grads_batchedTrue参数把多个 VJP 的 cotangent 堆成一个 batch 一次算完。y f(x) u_batched torch.eye(y.numel(), dtypey.dtype).reshape(y.numel(), *y.shape) grads torch.autograd.grad(y, x, grad_outputsu_batched, is_grads_batchedTrue)[0] # grads 形状 (y.numel(), *x.shape)每一片是一行雅可比这条路的优点是底层实现更贴近 autograd 原语某些老算子上比 vmap 兼容性更好。缺点是它算的是 VJP 批量方向选择依然受形状限制。6.2 显存前向模式不是免费的前向模式的显存开销常常被低估。每个参与运算的张量都要额外存一份 ε 分量理论上翻倍再叠加vmap的批维度实际可能是三到四倍。当输入维度 n 达到几千、中间激活又很大的时候jacfwd直接 OOM 是常事。几条降显存的手段按有效性排序用chunk_size参数分批 vmap把 n 个方向拆成若干组一组一组算再拼起来在torch.no_grad()外层包一层只让需要求导的部分进计算图换jacrev如果你的形状允许的话反向模式的显存模型完全不同有时候反而更省降低精度到 bf16前向模式对精度的容忍度比反向模式高一些因为没有梯度累积的链式放大但要先验证数值稳定性。6.3 一条选择决策线把这些经验压缩成一条可以照着走的判断线如果你的目标是一个方向上的方向导数且函数的输入维度不大用torch.func.jvp一次完事。如果目标是完整雅可比先看形状输入元素少于输出元素走jacfwd反之走jacrev。如果目标是海森矩阵乘向量用jvp(grad(f), ...)组合别去碰functional.hvp那个内部是反向两次慢。如果目标是雅可比范数或迹这类标量统计量别算完整雅可比用 Hutchinson 加随机方向估计采样四到八次就够。我个人的体会是torch.func这套 API 刚上手时最大的障碍不是语法而是思维方式的切换——它不依赖全局的requires_grad状态也不依赖.backward()一切都是显式的函数变换。一旦适应了这种写法代码反而比传统的 autograd 更清晰因为对什么求导、沿着什么方向全都写在函数签名里不用去猜计算图长什么样。真正需要警惕的只有两点一是别在需要前向模式的场景里误用torch.autograd.functional.jvp二是做数值校验时务必用 float64 并且避开折点这两条踩过一次就够记一辈子了。