高阶导数

主线位置：变换主线的「grad / vjp / jvp / hvp」一步，把变换叠起来求二阶导数；最后一节回到 TRG 主线的「分段求导」之前，看整图求二阶导数为什么行不通。

版本：JAX 0.12.0.1（CPU，x64）、tenferro-rs 0.8.0（CPU，native，单线程）

flowchart TB
  g["grad 的图"] --> fr["jvp 正向套反向"]
  g --> rr["grad 反向套反向"]
  fr --> t0["对角切向量为 0"]
  t0 --> ok["和差分一致"]
  rr --> big["对角余切很大"]
  big --> bad["抵消不净 偏离"]

tenferro 的导数还是 TracedTensor，而且沿用原叶子的输入键，所以对导数再求导就是把变换叠起来，和 JAX 一样（机制见 traced AD 的「高阶导数」）。本页先在变换主线的 f(x) = Σ sin(x)·x² 上对照 hvp 和 Hessian，再看经过 SVD 时两种叠法的差别，最后看 TRG 整图求二阶导数的代价。

实验代码：tenferro 一侧是 bin higher、trg_second（release，CpuEngine::serial()），JAX 一侧是 higher.py。M2 Max、macOS 27、rustc 1.99.0；负载均值 4–7。

hvp

fwd_rev = jax.jit(lambda x, v: jax.jvp(jax.grad(f), (x,), (v,))[1])
rev_rev = jax.jit(lambda x, v: jax.grad(lambda x: jnp.vdot(jax.grad(f)(x), v))(x))

let g = ad.grad(&f(&x)?, &x)?;
let hv = ad.jvp(&g, &x, &vc)?;                            // vc 带数据
let v = TracedTensor::input_concrete_shape(DType::F64, &[N])?;
let hr = ad.grad(&(&g * &v)?.reduce_sum(None)?, &x)?;      // v 是占位符

# JAX
jvp(grad)       max|.-exact|=1.332e-15 jit first 18749 us steady 12.4 us
grad(<grad,v>)  max|.-exact|=1.776e-15 jit first 20525 us steady 12.4 us
# tenferro
jvp(grad)       max|.-exact|=1.332e-15 transform 167.4 us inputs=3 ops=21 steady 23.7 us
  x index Some(0)
  run [x, seed, 2v]: Err: binding rank mismatch for symbolic-shape placeholder: expected rank 1, got rank 0
  run [x, 2v, seed]: max|.-2*exact|=2.665e-15
grad(<grad,v>)  max|.-exact|=1.776e-15 transform 162.1 us inputs=4 ops=23 steady 25.0 us

两种叠法两边都对，误差都在 1e-15 量级（hvp 的绝对值最大约 5.1），连误差的数值都一样。

• 前向套反向（jvp(grad)）是 JAX 文档推荐的写法，tenferro 上游的 hvp.rs 测试也主要用它。tenferro 里切向量建图时要带数据，程序的 3 个输入是 x、切向量和 grad 的种子；和 上一页 一样，run 时给全 3 个输入就能换 v。顺序要试出来：按 [x, 种子, 2v] 传，被秩检查拦下；按 [x, 2v, 种子] 传，得到 2 倍的 hvp。如果种子和切向量形状相同，传反了不会报错。
• 反向套反向（grad(⟨grad, v⟩)）里 v 只是函数里的一个张量，可以是占位符，按 compile_with_input_specs 的顺序传，4 个输入里另外两个是两层 grad 的种子。
• tenferro 的 hvp 稳态约 24 µs，是 JAX 的两倍，程序有 21、23 个算子。一阶的 grad 是 14 µs 对 11 µs，阶数越高差得越多；原因没有细查，tenferro 怎么把逐元素运算分区执行见 编译 的「分段与逐元素区」。

Hessian：没有 jax.hessian

hess = jax.hessian(f)(x4)                 # x4 = [0.1, 0.7, -1.3, 2.0]

for j in 0..4 {
  let mut e = vec![0.0_f64; 4];
  e[j] = 1.0;
  let col = ad.jvp(&g4, &x4, &TracedTensor::from_vec_col_major(vec![4], e)?)?;
  // 编译、准备、执行，得到第 j 列
}

# JAX
jax.hessian n=4 max|.-exact|=0.000e+00
# tenferro
4 columns: transform+compile+prepare+run 988.8 us total, max|H - exact|=2.220e-16

jax.hessian 是 jacfwd(jacrev(f))，靠 vmap 把所有基向量一次推过去。tenferro 没有 vmap、jacfwd、hessian，只能一列一列地做 jvp(grad)。上面每列都重新变换一次；既然切向量是可以在 run 时替换的输入，也可以只变换、编译一次，再用 n 个基向量各跑一次。n 大时两种写法都是 n 次执行，JAX 是一次批量执行。

经过 SVD：jvp(grad) 对，grad(grad) 偏

两种叠法在数学上是同一个二阶导数，经过 SVD 时数值上却不一样。例子沿用试验阶段的 5×5 情形：a0、w 是固定种子的正态随机矩阵（奇异值 2.84、2.57、1.84、0.49、0.33，没有简并），

l(t) = Σ ((U[:, :2] · √s[:2]) Vᴴ[:2]) ⊙ w，A = a0 + t·w，在 t = 0 求 l''。

def mk(eps, masked):
  @jax.custom_jvp
  def svd(a):
    return jnp.linalg.svd(a, full_matrices=False)

  @svd.defjvp
  def rule(p, t):
    # 和 JAX 内置规则相同，只是 F = d / (d² + eps²)，d = (s_j - s_i)(s_j + s_i)
    ...
    fb = d / (d * d + eps * eps)
    f_ = jnp.where(jnp.eye(s.shape[0]) > 0, 0.0, fb) if masked else fb
    ...

let t = TracedTensor::input_concrete_shape(DType::F64, &[])?;
let l = svd_loss(&t, a0, w, eps)?;   // a.svd() 或 svd_with_options(SvdOptions::default().derivative_eps(e))
let g = ad.grad(&l, &t)?;
let gj = ad.jvp(&g, &t, &one)?;      // jvp(grad)
let gg = ad.grad(&g, &t)?;           // grad(grad)

tenferro 的 SVD 规则 linearize_svd 用的就是这个展宽的 F，默认 ε = 1e-12，对角线不屏蔽（规则本身见 线性代数 的「SVD 的反向规则」）。JAX 一侧用 custom_jvp 写了同样的规则，屏蔽和不屏蔽对角线各一份：

对照值：JAX 内置 jnp.linalg.svd 的 grad(grad) 和 jvp(grad) 都是 −11.0069745057，梯度的中心差分是 −11.0069745042（JAX）和 −11.0069745040（tenferro）。ε = 1e-3 时连正确的二阶导数也被展宽改了一点（各方的 jvp(grad) 都是 −11.0069704925），那一行只看三列彼此差多少。

# tenferro，默认 ε
default eps  grad=5.924521608687 grad(grad)=-10.3648209072 jvp(grad)=-11.0069745057 FD(grad)=-11.0069745040
             transform ms: grad 1.57, jvp(grad) 1.72, grad(grad) 1.22; ops: grad 82, jvp(grad) 205, grad(grad) 211

• jvp(grad) 在两边、所有 ε 下都和差分一致，ε ≤ 1e-6 时前 10 位和 JAX 内置规则相同。
• 不屏蔽对角线时，grad(grad) 两边都错，错的方向不一样（JAX −11.49，tenferro −10.36，真值 −11.01），偏差 4% 和 6%。屏蔽之后 JAX 的 grad(grad) 就对了。
• ε 越小越糟；ε = 1e-3 时三列彼此只差 5e-9。

原因在对角线。F[i,i] = 0/(0 + ε²) = 0，值没有问题，但 F 对 d 的斜率在 d = 0 处是 1/ε²，ε = 1e-12 时是 1e24。

• 前向模式里，这个斜率乘的是 d_ii 的切向量 (ṡ_i − ṡ_i)(s_i + s_i)，正好是 0，乘出来还是 0。所以 jvp(grad) 不受影响。
• 反向模式里，它乘的是 F[i,i] 收到的余切，这个余切不为 0，d_ii 得到一个 1/ε² 量级的余切；它再分成一正一负两份流回同一个 s_i，数学上抵消，但两份是在不同的求和里先和别的项加在一起、各自舍入之后才相减的。屏蔽对角线（where）让这个余切在源头就是 0。

ε 从 1e-12 放大到 1e-6，偏差从 0.48 和 0.64 降到 4.0e-3 和 6.8e-3，这符合「巨大的余切抵消不净」。但偏差并不随 1/ε² 无限变大（见动脑筋），更像是正常项在大数旁边被舍掉了；这一点是推断，没有逐项打印余切核实。

tenferro 的 SvdOptions 只有 derivative_eps，没有屏蔽对角线的开关。所以在 tenferro 里对经过 SVD（以及用同一种展宽的 eigh，这一点没有测）的函数求二阶导数，用 jvp(grad)，别用 grad(grad)。JAX 里也一样，除非自己写的规则屏蔽了对角线。JAX 内置 SVD 规则这里没问题，但它遇到简并奇异值会给出 NaN，那是 简并奇异值 的话题。

二阶 TRG：变换本身的代价

TRG 主线要比热 c = β² ∂² lnZ/N / ∂β²。JAX 对整个 TRG 写 jit(jvp(grad)) 就行（自己的 SVD 规则屏蔽对角线），χ = 16、β = 0.3：

steps= 1 jvp(grad)=2.440159969659 FD(grad)=2.440159971018 jit first 376 ms steady 0.09 ms
steps= 2 jvp(grad)=3.861353061955 FD(grad)=3.861353060830 jit first 445 ms steady 0.26 ms
steps=20 jvp(grad)=3.180981990549 FD(grad)=3.180981980461 jit first 4182 ms steady 780.48 ms

tenferro 把同样的 TRG 建成一张图（β 是标量占位符），再叠两层变换：

let beta = TracedTensor::input_concrete_shape(DType::F64, &[])?;
let lnz = trg(&beta, &ks, None)?;
let g = ad.grad(&lnz, &beta)?;
let h = ad.jvp(&g, &beta, &one)?;     // 和 ad.grad(&g, &beta) 各计时

tenferro 0.8.0 traced TRG chi=16 steps=1 ks=[2]
ad.grad              4.48 ms
ad.jvp(grad)       212.41 ms
ad.grad(grad)      207.63 ms
grad       ops 266
jvp(grad)  ops 715
grad(grad) ops 737
beta=0.30 grad=1.667309214024 FD(grad)=2.440159971118 jvp(grad)=2.440159969659 grad(grad)=2.440159969659
beta=0.40 grad=1.843337108813 FD(grad)=1.204216607620 jvp(grad)=1.204216606546 grad(grad)=1.204216606546

tenferro 0.8.0 traced TRG chi=16 steps=2 ks=[2, 4]
ad.grad              9.36 ms
exit=142 elapsed=46s

（FD(grad) 是梯度的中心差分；2 步那次用 alarm 限时 45 s，退出码 142 是被 SIGALRM 终止。）

• 1 步时，tenferro 的二阶导数和 JAX 前 12 位相同，两种叠法在这里都对。但光是变换就要 0.2 s，导数程序 715 个算子。
• 2 步时，ad.grad 只要 9 ms，ad.jvp(grad) 45 s 内没有做完。试验阶段在同一台机器上不限时跑过，2 步的 ad.jvp(grad) 和 ad.grad(grad) 各约 240 s，而且那次 grad(grad) 在五个 β 上偏了 2%–4%，jvp(grad) 和差分一致到 1e-9。
• JAX 第一次调用（追踪、变换、编译）1 步 0.38 s、2 步 0.45 s、20 步 4.2 s，增长平缓。

从 1 步到 2 步，tenferro 的二阶变换至少慢了两百倍，一阶的 ad.grad 只慢了一倍。试验阶段把这种随深度急剧增长的耗时追到多输出算子（SVD）的输出被反复使用的地方，复现和测量是 求导的代价 的主题。实用的写法是按步分段、手工串反向，再对 AD 得到的能量做差分求比热，见 分段 和 动手：Ising 模型的 TRG。

动脑筋答案

# JAX
eps=1e-09 plain      grad=5.924521608687 grad(grad)=-11.4884476856 jvp(grad)=-11.0069745057 FD(grad)=-11.0069745042
# tenferro
eps=1e-9     grad=5.924521608687 grad(grad)=-10.3648209072 jvp(grad)=-11.0069745057 FD(grad)=-11.0069745040

• 两边都和 ε = 1e-12 时一模一样，10 位都不变：JAX −11.4884476856，tenferro −10.3648209072。
• 斜率小了一百万倍，偏差却一点没少，所以偏差不是「1/ε² 乘一个舍入误差」那样的噪声。一个说得通的解释（推断，没有逐项核实）：ε 小到一定程度后，那一项大到把同一个和式里的正常项全部舍掉，结果停在一个固定的错误值上；ε = 1e-6 时只舍掉一部分，ε = 1e-3 时几乎不舍。
• 两边停在不同的值上，反向程序的算子排列不同，被舍掉的项也就不同。

小结

• hvp 两边写法一样：jvp(grad) 或 grad(⟨grad, v⟩)，误差 1e-15 量级；tenferro 稳态约 24 µs，JAX 约 12 µs。
• tenferro 没有 vmap、jacfwd、hessian；Hessian 按列做 jvp(grad)，切向量可以在 run 时替换，变换编译一次就够。
• 经过 SVD 时，tenferro 的展宽 F 不屏蔽对角线，grad(grad) 在默认 ε 下偏 6%（JAX 不屏蔽的同样规则偏 4%），jvp(grad) 对。二阶导数用前向套反向。
• ε ≤ 1e-9 时偏差停在固定值，不随 1/ε² 增大；放大 ε 能减小偏差（1e-6 时约 5e-3），但也会改动真正的导数，不如换成 jvp(grad)。
• 整图 TRG 的二阶变换：1 步 0.2 s，2 步超过 45 s（试验阶段约 240 s）；JAX 20 步首次调用 4.2 s。tenferro 上要分段。