企业官网建设流程全解析

做自动微分这块的人，多半遇到过这样的场景：想看看模型输出对输入的敏感程度，或者想给损失加一个梯度惩罚项，loss.backward()敲下去，拿到的却是和输入同形状的一坨梯度，跟脑子里那个“雅可比矩阵”怎么也对不上号。问题就出在 PyTorch 的反向传播引擎天生只给你一种乘积——向量雅可比积，而很多场景真正想要的是它的孪生兄弟，雅可比向量积（Jacobian-Vector Product，简称 JVP）。这两个东西长得像、名字像，方向却完全相反，搞混一次能让人调一晚上。这篇内容就把 PyTorch 里 JVP 这件事从头到尾拆一遍：它是什么、为什么反向传播给不了你、有哪几种写法、每种写法什么时候用、踩过哪些坑、数值上怎么验证对不对。

内容偏实战，适合已经能写loss.backward()但被“雅可比”“二阶导”“梯度惩罚”这些词卡住的人，也适合想搞清楚自动微分底层方向的读者。基础概念我会用生活化的类比铺一遍，代码都是可直接跑的完整片段，照着抄就能验证。

1. 从雅可比矩阵到雅可比向量积：为什么要盯住这个“乘积”

1.1 雅可比矩阵到底是什么，为什么几乎没人真的把它算出来

先把定义摆清楚。有一个函数 $f: \mathbb{R}^n \to \mathbb{R}^m$，输入是 $n$ 维向量，输出是 $m$ 维向量。它的雅可比矩阵 $J$ 是一个 $m \times n$ 的矩阵，第 $i$ 行第 $j$ 列的元素是 $\partial f_i / \partial x_j$，也就是“第 $i$ 个输出对第 $j$ 个输入的偏导”。单输出标量函数（比如分类的 loss）的雅可比就是一个 $1 \times n$ 的行向量，退化成了我们熟悉的梯度。

理论很干净，实际却很少有人真的把 $J$ 显式构造出来。原因是维度。举个例子：一个普通的全连接分类器，输入是 784 维的展平图像，输出是 10 类，那单样本的输入雅可比是 $10 \times 784$，不到八千个数，float32 存下来也就 30KB，随便算。但只要换成图像到图像的任务，比如超分、去噪、分割，输入和输出都是 $256 \times 256 \times 3$，那雅可比就是 $196608 \times 196608$，接近 390 亿个元素，光存一份 float32 就要 150GB 上下。这还只是一个样本。再叠上 batch，直接没法看。

所以真实工程里的做法从来不是“把雅可比算出来”，而是“把雅可比乘到某个向量上”。这一步乘法把 $m \times n$ 的矩阵压回一个 $m$ 维或 $n$ 维的向量，存储和计算都回到了可接受的范围。这也是整个自动微分体系的设计根基：没人想要矩阵，大家要的都是矩阵和向量的乘积。

1.2 JVP 与 VJP：同一个矩阵的两个方向

同样是“雅可比乘向量”，方向不同，名字不同，代价也完全不同。

雅可比向量积（JVP）：拿一个和输入同维的向量 $v \in \mathbb{R}^n$，算 $Jv \in \mathbb{R}^m$。换个说法，你沿着输入空间的某个方向 $v$ 轻轻推一下，看输出往哪个方向动、动多快，这本质上是方向导数。用二阶展开的语言说，$f(x + \epsilon v) \approx f(x) + \epsilon \cdot Jv$，$\epsilon$ 很小时这个近似相当准，这也是后面做数值校验的依据。

向量雅可比积（VJP）：拿一个和输出同维的向量 $u \in \mathbb{R}^m$，算 $J^\top u \in \mathbb{R}^n$。它的含义是，把输出空间上的一个方向“反投影”回输入空间，告诉你输入该怎么调才能让输出沿 $u$ 方向增长最快。

这两个操作在自动微分里有明确的对应关系：反向模式自动微分天然产出 VJP，前向模式自动微分天然产出 JVP。反向模式就是 PyTorch 的backward()，它从输出往输入走，每过一层就把上游传来的向量乘上该层雅可比的转置，一路下来正好是 $J^\top u$。前向模式反过来，从输入往输出走，带着一个“切向量”一路传播，得到的就是 $Jv$。

代价上有个很好记的口诀：一次 VJP 的开销大约是单次前向传播的常数倍（通常 2 到 3 倍），一次 JVP 也差不多。但如果要拼出完整雅可比，反向模式得跑 $m$ 次，前向模式得跑 $n$ 次。所以选择哪种模式，完全取决于是输出维度大还是输入维度大。

1.3 PyTorch 默认走的是哪条路，这一点必须先掰清楚

loss.backward()之所以能工作，是因为 loss 是标量，$m = 1$。此时雅可比退化成 $1 \times n$ 的行向量，VJP 里的 $u$ 取标量 1，算出来就是完整梯度。换句话说，反向传播是 VJP 在 $m=1$ 时的特例。

这个认知非常关键，因为很多人第一次接触torch.autograd.grad时会写出这样的代码，并且误以为自己在算 JVP：

import torch def f(x): return torch.stack([x[0] * x[1], x[1] ** 2, torch.sin(x[0])]) x = torch.tensor([1.5, -0.5], requires_grad=True) y = f(x) u = torch.tensor([1.0, 0.0, 0.0]) out = torch.autograd.grad(y, x, grad_outputs=u, retain_graph=True)[0] print(out) # tensor([-0.5000, 1.5000])

这段代码里，grad_outputs=u的形状必须和输出 $y$ 一致，它扮演的是 VJP 里的 $u$。算出来的结果是 $J^\top u$，也就是雅可比第一行转置，形状是输入维度 2。它不是 $Jv$。判断依据特别简单：grad_outputs的形状永远跟着输出走，结果的形状永远跟着输入走，这就是 VJP 的指纹。反过来，如果某个接口要求你传入的向量和输入同形状、返回值和输出同形状，那才是 JVP。

搞混这一点，轻则梯度方向反了导致训练不收敛，重则在二阶导场景里算出完全错误的惩罚项，而且数值上还不报错，非常隐蔽。

2. PyTorch 里做雅可比向量积的四种写法与选型

2.1 写法一：forward_ad 双数模式，最贴近“前向模式”原意

先讲最底层、最能说明原理的一种。PyTorch 在torch.autograd.forward_ad里提供了“双数”机制（dual number），思想很朴素：把每个参与计算的数扩展成“值 + 切向量”两部分，值按普通规则算，切向量按求导规则同步传播。这和复数运算的写法几乎同构，学过一点数值方法的人会觉得很亲切。

import torch from torch.autograd.forward_ad import dual_level, make_dual, unpack_dual def f(x): return torch.stack([x[0] * x[1], x[1] ** 2, torch.sin(x[0])]) x = torch.tensor([1.5, -0.5]) v = torch.tensor([1.0, 2.0]) with dual_level(): dual_x = make_dual(x, v) # 把 x 包装成 dual 张量，切向量是 v dual_y = f(dual_x) # 正常调用函数，切向量自动传播 y, jvp = unpack_dual(dual_y) # 拆出原始输出和 JVP print(y) # tensor([-0.7500, 0.2500, 0.9975]) print(jvp) # tensor([ 2.5000, -2.0000, 0.0707])

这段代码有几个地方必须留意。第一，dual_level()是一个上下文管理器，dual 张量只能在它的作用域内参与运算，出了with块再拿它去做任何计算都会直接报错，很多人把unpack_dual写在块外面（这样是对的），但把f(dual_x)写在块外面（这样会炸）。第二，make_dual的两个参数形状必须完全一致。第三，被包装的x不需要requires_grad=True，前向模式和反向模式是两套独立机制，混着设反而容易让人困惑。

这种写法的好处是语义最直白，你能亲眼看到切向量是怎么一层层传播的，调试时打印中间层的 dual 张量非常直观。缺点是啰嗦，而且不是所有算子都完整支持前向模式（PyTorch 这些年补得很勤，但边缘算子仍有缺失），一旦碰上不支持的算子会直接抛异常而不是静默降级。

2.2 写法二：torch.autograd.functional.jvp，一行搞定的高层封装

不想手搓双数的话，torch.autograd.functional模块里有现成的：

import torch from torch.autograd.functional import jvp def f(x): return torch.stack([x[0] * x[1], x[1] ** 2, torch.sin(x[0])]) x = torch.tensor([1.5, -0.5]) v = torch.tensor([1.0, 2.0]) y, jv = jvp(f, x, v) print(y) # tensor([-0.7500, 0.2500, 0.9975]) print(jv) # tensor([ 2.5000, -2.0000, 0.0707])

返回的是一个二元组：原始输出和 JVP，顺序别记反。inputs可以是单个张量，也可以是张量元组；如果函数有多个输入，v也要给成同结构的元组。

这个接口有两个需要注意的地方。一是它要求函数是“纯”的——不能修改外部可变状态、不能有依赖全局变量的副作用，否则内部构造计算图时会得到和预期不符的结果。二是关于它的实现路径，我记得早期版本里它并不是真正走前向模式，而是借了双重反向的路径把结果拼出来，所以性能上未必比手写 VJP 循环快多少。这一点随着版本演进而变化，我的建议是别把性能赌在实现细节上，需要压性能就自己测一遍。

同一个模块里还有vjp、jacobian、hessian、hvp。其中jacobian有个vectorize=True参数，会在内部用批量化技巧提速，算中小规模雅可比时比逐个循环快几倍，值得记住。

2.3 写法三：torch.func.jvp，目前最推荐的通用解

torch.func是把原来的 functorch 合并进主仓之后的产物，设计目标很明确：让函数变换（transform）可以自由组合。它的jvp写法是这样的：

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([1.5, -0.5]) v = torch.tensor([1.0, 2.0]) y, jv = jvp(f, (x,), (v,)) print(y) # tensor([-0.7500, 0.2500, 0.9975]) print(jv) # tensor([ 2.5000, -2.0000, 0.0707])

注意参数结构变了：primals和tangents都必须是元组，哪怕只有一个输入也要写成(x,)。这是新手最常见的报错来源，忘了加逗号就会得到一个整数而不是元组，然后报一个看起来毫不相干的类型错误。

torch.func系的最大价值是可组合。它和vmap、grad、jacrev、jacfwd能自由嵌套，这一点在需要“对一批样本分别求雅可比”或者“对一组方向批量求 JVP”时体现得淋漓尽致。代价是对函数纯度要求更严格：不能有 in-place 操作，不能有随机性（除非显式处理随机数生成器），不能依赖 Python 侧的副作用。写惯了普通 PyTorch 代码的人第一次迁移过去，十有八九会被 in-place 报错拦住，x += 1要改成x = x + 1，relu(inplace=True)要改成inplace=False。

2.4 写法四：用 vmap 批量铺开，一次拿到整块雅可比

单方向 JVP 只是起点。真实需求往往是“我要看输入空间所有基方向上的响应”，也就是完整雅可比。这时把jvp和vmap组合起来用：

import torch from torch.func import jvp, vmap def f(x): return torch.stack([x[0] * x[1], x[1] ** 2, torch.sin(x[0])]) x = torch.tensor([1.5, -0.5]) basis = torch.eye(2) # 两个基向量 e1, e2 def one_jvp(v): return jvp(f, (x,), (v,))[1] cols = vmap(one_jvp)(basis) print(cols) # tensor([[-0.5000, 0.0000, 0.0707], # [ 1.5000, -1.0000, 0.0000]])

这里有一个非常容易搞反的地方：对基向量 $e_1, e_2$ 分别做 JVP，得到的是 $Je_1$ 和 $Je_2$，而 $Je_j$ 恰好是雅可比的第 $j$ 列。所以堆起来的形状是 $n \times m$，是雅可比的转置，不是雅可比本身。想要标准的 $m \times n$ 形式，得再transpose(0, 1)一下。我见过不止一次有人拿这个结果去做行方向的归一化，维度对不上却因为方形矩阵碰巧能跑通，最后数值全错。

2.5 四种写法的选型对照

选型上我的一般原则是：临时验证用functional，正式代码用torch.func，需要理解机制时手搓forward_ad。如果模型里有大量 in-place 操作又不想改代码，functional系列的容错性反而更好一些。

3. 实操：从零跑通一个 JVP 并做数值校验

3.1 先手算一遍，把标准答案钉死

写代码之前，先拿纸笔把答案算出来，后面所有实现都拿它对照。用前面那个函数：

$f_0 = x_0 x_1$，$f_1 = x_1^2$，$f_2 = \sin(x_0)$，在 $x = (1.5, -0.5)$ 处。

偏导逐个求：$\partial f_0/\partial x_0 = x_1 = -0.5$，$\partial f_0/\partial x_1 = x_0 = 1.5$；$\partial f_1/\partial x_0 = 0$，$\partial f_1/\partial x_1 = 2x_1 = -1.0$；$\partial f_2/\partial x_0 = \cos(1.5) \approx 0.070737$，$\partial f_2/\partial x_1 = 0$。

于是雅可比是：

$$J = \begin{bmatrix} -0.5 & 1.5 \ 0 & -1.0 \ 0.070737 & 0 \end{bmatrix}$$

取方向 $v = (1, 2)$，$Jv = (-0.5 \times 1 + 1.5 \times 2,\ 0 \times 1 + (-1.0) \times 2,\ 0.070737 \times 1 + 0 \times 2) = (2.5, -2.0, 0.070737)$。

再取 $u = (1, 0, 0)$，$J^\top u = (-0.5,\ 1.5)$。这两个答案就是后面所有代码的验收标准，一个数字都不能差。

3.2 三种实现并排跑，对答案

把前面的代码整合成一份可执行脚本，一次跑完三种写法：

import torch from torch.autograd.forward_ad import dual_level, make_dual, unpack_dual from torch.autograd.functional import jvp as func_jvp from torch.func import jvp as tf_jvp def f(x): return torch.stack([x[0] * x[1], x[1] ** 2, torch.sin(x[0])]) x = torch.tensor([1.5, -0.5]) v = torch.tensor([1.0, 2.0]) # 写法 A：双数模式 with dual_level(): y_a, jv_a = unpack_dual(f(make_dual(x, v))) # 写法 B：functional 封装 y_b, jv_b = func_jvp(f, x, v) # 写法 C：torch.func y_c, jv_c = tf_jvp(f, (x,), (v,)) print("双数 :", jv_a) print("functional:", jv_b) print("torch.func:", jv_c) print("理论值 : tensor([ 2.5000, -2.0000, 0.0707])")

三种写法输出完全一致。实际跑的时候建议用torch.allclose断言而不是肉眼看，float32 下末位数字的差异肉眼很难分辨。如果换成 float64（torch.tensor(..., dtype=torch.float64)），精度能压到 $10^{-15}$ 量级，后续做校验也更从容。

3.3 用有限差分做独立校验，别信自己写的解析解

上一节的“理论值”是我手算的，手算也可能错。真正独立的验证手段是中心差分：

def central_diff(f, x, v, eps=1e-3): return (f(x + eps * v) - f(x - eps * v)) / (2 * eps) print(central_diff(f, x, v)) # tensor([ 2.5000, -2.0000, 0.0707])

这个式子来自二阶展开：$f(x + \epsilon v) - f(x - \epsilon v) \approx 2\epsilon Jv + O(\epsilon^3)$，误差是三阶小量，所以中心差分比前向差分的精度高一阶。$\epsilon$ 的取值很有讲究：float32 下取 $10^{-3}$ 附近比较合适，太小会被舍入误差淹没（float32 的机器精度约 $1.2 \times 10^{-7}$），太大则截断误差抬头。float64 下可以取到 $10^{-6}$，精度能到 $10^{-10}$，这时如果自动微分的结果和你差出 $10^{-5}$ 以上，基本可以确定代码写错了。

注意：做差分校验时一定要先把张量转成 float64，并且不要让它们参与梯度计算（requires_grad=False）。混着requires_grad=True做差分，会不知不觉把计算图撑大，显存悄悄涨上去。

3.4 逐样本统计：vmap 组合拳的正确姿势

真实场景里更常见的是“一个 batch 里每个样本都要算雅可比的某个统计量”，比如雅可比范数。这时vmap是刚需：

import torch from torch.func import jacfwd, vmap def model(xi): # 单个样本：输入 4 维，输出 3 维 return torch.tanh(xi @ torch.ones(4, 3)) def jac_norm(xi): J = jacfwd(model)(xi) # 形状 (3, 4) return (J ** 2).sum() batch = torch.randn(8, 4) norms = vmap(jac_norm)(batch) # 形状 (8,) print(norms.shape)

这里的jacfwd内部就是沿着输入维度做 $n$ 次 JVP，再拼成矩阵，所以 $n$ 小的时候非常划算。反过来，如果输出维度远小于输入维度，jacrev更合适，它走的是 $m$ 次 VJP 的路径。

有一个坑必须提：vmap(jac_norm)要求jac_norm内部的所有操作都支持批量化语义，而jacfwd本身又会在内部插入vmap，形成嵌套。嵌套层数一多，报错信息就会非常难读，通常只报一句“某个算子在 vmap 下不支持”。排查办法是把嵌套拆开，先确认model在单样本下能跑，再确认jacfwd(model)能跑，最后才加vmap，一层一层定位。

4. 雅可比向量积真正派上用场的地方

4.1 二阶优化：海森矩阵向量积是 JVP 的经典战场

一阶梯度不够用时，人们会想到海森矩阵。但海森是 $n \times n$，百万参数就是万亿元素，不可能显式构造。真正被用的是海森矩阵向量积（HVP）：给定方向 $v$，算 $Hv$。而 $Hv$ 本身就是“梯度的雅可比”乘 $v$，天然是个 JVP 问题。

PyTorch 里的标准写法是两次反向传播：

import torch def loss_fn(x): return (x ** 3).sum() x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True) v = torch.tensor([1.0, 0.0, -1.0]) g = torch.autograd.grad(loss_fn(x), x, create_graph=True)[0] # 一阶导 3x^2 hvp = torch.autograd.grad(g, x, grad_outputs=v)[0] # 再对 v 求 VJP print(g) # tensor([ 3., 12., 27.], grad_fn=<...>) print(hvp) # tensor([ 6., 0., -18.])

手算对一下：$H = \mathrm{diag}(6x) = \mathrm{diag}(6, 12, 18)$，$Hv = (6 \times 1,\ 12 \times 0,\ 18 \times (-1)) = (6, 0, -18)$，对上了。

这里的关键是create_graph=True。它告诉 autograd“这次求导的结果还要参与后续求导，别把图拆了”。少了这个参数，第二次grad会报“元素不需要梯度”，因为第一次求导已经把计算图释放了。另一个容易被忽略的点是retain_graph：autograd.grad的retain_graph默认值等于create_graph，所以开了create_graph一般不用再显式指定，但如果你在同一张图上要做第三次求导，还是得加retain_graph=True，否则第二次就把图用掉了。

torch.autograd.functional里还有一个现成的hvp，一行出结果，适合快速验证。但要注意它的返回结构和jvp一样是二元组，别只接一个变量。

4.2 梯度惩罚：把“梯度范数接近 1”写成可优化的损失

梯度惩罚是一类非常典型的 JVP 使用场景。核心想法是：让判别器或评分网络对输入的梯度范数保持在 1 附近，从而限制函数的变化剧烈程度，稳定训练。实现上需要先对插值样本求梯度，再对梯度求范数并纳入损失：

def gradient_penalty(critic, x_real, x_fake, alpha=None): if alpha is None: alpha = torch.rand(x_real.size(0), *([1] * (x_real.dim() - 1)), device=x_real.device) x_hat = (alpha * x_real + (1 - alpha) * x_fake).requires_grad_(True) score = critic(x_hat) grads = torch.autograd.grad( outputs=score.sum(), inputs=x_hat, create_graph=True, retain_graph=True, )[0] norm = grads.flatten(1).norm(2, dim=1) return ((norm - 1.0) ** 2).mean()

这里score.sum()里的.sum()就是在构造 VJP 的种子向量（全 1 向量）。因为输出是逐样本的标量，全 1 向量恰好等价于“对每个样本分别求一次梯度再拼起来”，数学上是等价的，只是省下了一次循环。create_graph=True必须有，否则惩罚项本身没有梯度，参数更新时这一项等于不存在——这个错误极其常见，因为代码能跑、损失也在下降，只是惩罚完全没生效。

实际做的时候还有两个经验点。一是alpha的采样维度要小心，样本形状是(B, C, H, W)时要生成(B, 1, 1, 1)而不是(B,)，否则广播出来的结果是错的，而且不报错。二是插值样本一定要requires_grad_(True)，很多人忘了这一步，grad会报“输入不需要梯度”。

4.3 输入敏感度分析：JVP 是天然的工具

想回答“输入的某个方向被扰动时，输出变化多剧烈”，这就是标准的 JVP。比如一个回归模型，输入是几百维的特征，你想知道沿“用户年龄”这个方向扰动 1 个单位，输出会变多少。用 JVP 一次就能拿到，不需要构造整个雅可比：

import torch from torch.func import jvp def model(x): return torch.stack([x.sum(), (x ** 2).sum()]) x = torch.randn(5, dtype=torch.float64) direction = torch.zeros(5, dtype=torch.float64) direction[0] = 1.0 # 只沿第 0 维扰动 y, dy = jvp(model, (x,), (direction,)) print(dy) # 两个输出分量的方向导数

结果向量告诉你：沿这个方向走，每个输出分量各自的变化率。这就是“局部敏感度”。如果模型是几十层深、上百万参数，显式雅可比不可能构造，但 JVP 的开销只比前向传播多一点，完全可接受。

顺带提一句，有些场景会在损失里加“输入雅可比范数”的正则项，用来提升模型对输入扰动的鲁棒性。这本质上是把 JVP 的模长当作惩罚项，属于同一类思路。写法上还是jacfwd加vmap，需要注意内存占用会随 batch 和输入维度平方级增长，batch 开大之前先小规模测一下显存。

4.4 前向模式什么时候真的比反向模式划算

很多人有个误解，觉得“反向传播是主流，前向模式没用”。这只在有大量参数、标量损失的情况下成立。一旦输出维度远大于输入维度，局面立刻反转。

举个具体例子：某个物理仿真场景，输入是 3 维的初始条件，输出是 1000 步的时间序列，每步 10 个观测量，总共 10000 维输出。这时雅可比是 $10000 \times 3$。用反向模式拼完整雅可比要跑 10000 次；用前向模式只需要 3 次。差了三个数量级。

反过来，标准分类任务输入是百万像素、输出是一个 loss，输入维度远大于输出维度，反向模式一次搞定，前向模式要跑上百万次。判断标准就一条：数一数输入维度和输出维度，哪边小就用哪种模式的方向。JVP 对应前向模式，适合输出维度大的情况；VJP 对应反向模式，适合输入维度大的情况。

5. 常见坑与排查清单

5.1 形状对不上：三类高频报错

第一类是grad_outputs的形状给错。用torch.autograd.grad(y, x, grad_outputs=u)时，u必须和y形状完全一致，不是和x一致。想要广播是行不通的，autograd 要求精确匹配。

第二类是torch.func.jvp的参数不是元组。jvp(f, x, v)会报一个非常不直观的错误，因为第一个位置参数被当成了函数，第二个被当成 primals，传入的张量会被当成元组去迭代。看到“object is not iterable”这类报错，先检查是不是漏了括号和逗号。

第三类是把 vmap 出来的列方向当成行方向。前面说了，对基向量做 JVP 得到的是雅可比的列，堆起来是转置。方阵还能糊弄过去，非方阵直接形状错误，但如果你的代码里用了reshape或者view，可能把它强行掰成能跑的形状，然后静默算错。

5.2 非纯函数、in-place 与随机性

torch.func系列对函数纯度的要求，是新手最大的门槛。三条硬规矩：

• 不要有 in-place 操作。x.add_(1)、x.relu_()、tensor[0] = 1统统不行，改成等价的非原地写法。
• 不要在函数内部改 Python 容器（列表、字典）。results.append(...)这种副作用在变换下行为未定义。
• 随机性要显式处理。函数里调torch.randn在vmap下会得到所有批次共享同一份随机数的怪异结果，正确做法是把随机数当作输入传进来。

排查办法很土但有效：先把函数写成最朴素的形式，确认能在vmap下单样本跑通，再逐步加复杂度。报错信息通常只说“某个算子不支持”，逐行二分注释是最快的定位手段。

5.3 二阶导的链路陷阱

想做 HVP 或者二阶正则，create_graph=True是必须的。但光有它还不行，还有几种情况会翻车：

• 使用了不支持二阶导的算子（比如某些自定义的autograd.Function，它的backward里用了不可微的操作），表现为二阶导结果为 None 或零。
• 中间对张量做了.detach()或.item()，链路断在这里，后面所有二阶导都会变成零。这种情况最隐蔽，因为一阶导完全正常。
• 混合精度下做二阶导，float16 的动态范围撑不住，容易出 inf 或 nan。我的建议是二阶相关计算统一在 float32 甚至 float64 下做，混合精度只用于前向和反向的一阶部分。

验证办法：拿一个小规模可手算的例子（比如前面那个 $x^3$ 的例子），把 HVP 和手算结果逐个元素对一遍，别偷懒。

5.4 性能与显存：什么时候该停下来

JVP 本身开销不大，问题通常出在“批量铺开”这一步。用vmap算 $n$ 个方向的 JVP，显存占用大致是单次 JVP 的 $n$ 倍（图模式会做一部分优化，但量级不变）。输入维度上千时，一次算完整雅可比就可能直接 OOM。

几个实用对策：一是确认真的需要完整雅可比，很多统计量（范数、迹、对角线）都有更省的估计方法；二是分块计算，把基向量切成几组，每组算完立刻做归约再释放；三是用 float32 而不是 float64 做大头计算，只在最终校验时切 float64。

另外，torch.compile和torch.func的组合目前还有一些边界问题，如果开启编译后 JVP 结果异常，先关掉编译验证一遍，再决定要不要保留。

5.5 排查速查表

最后分享一个我自己踩过的坑：早期写梯度惩罚时，create_graph忘了开，训练跑了两周，loss 曲线漂亮得很，直到做消融实验才发现在去掉惩罚项后结果几乎没变——那一项从头到尾就没生效过。从那以后我养成了一个习惯，凡是涉及二阶导或者 JVP 的代码，都先拿一个能手算的小例子跑三遍：一次看数值、一次对差分、一次看梯度是否真的流回去了。这三遍加起来不到十分钟，能省下两周。