经过 SVD 的 grad(grad) 静默偏了约 6%

主线位置：变换主线的「grad / vjp / jvp / hvp」一步：把两个变换叠起来求二阶导数，中间经过 SVD。

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

flowchart TB
  l["l of t 经过 SVD"] --> g["grad"]
  g --> gg["再 grad"]
  g --> jg["再 jvp"]
  gg --> bad["-10.3648 偏 5.8%"]
  jg --> good["-11.0070 对"]
  fd["grad 的差分"] --> good

实验代码：JAX 脚本一个，tenferro 程序 p12_svd_grad_grad（AdContext 带线代规则，默认 SvdOptions，CpuEngine::serial()，release）。a0、w 是固定种子的 5×5 正态随机矩阵，两边从同一个文件读入。平台 M2 Max、macOS 27、rustc 1.99.0。

现象

l(t) = Σ ((U[:, :2] · √s[:2]) Vᴴ[:2]) ⊙ w，A = a0 + t·w，在 t = 0 求 l''。JAX 里习惯直接 grad(grad) 或 jax.hessian：

def loss(t):
  u, s, vh = jnp.linalg.svd(a0j + t * wj, full_matrices=False)
  return jnp.sum((u[:, :2] * jnp.sqrt(s[:2])) @ vh[:2] * wj)

g = jax.grad(loss)
jax.grad(g)(0.0); jax.jvp(g, (0.0,), (1.0,))[1]; jax.hessian(loss)(0.0)

singular values: [2.842414953926728, 2.5719761392455838, 1.8351127914139929, 0.4928116168814265, 0.33471429328282876]
grad=5.924521608687 grad(grad)=-11.0069745057 jvp(grad)=-11.0069745057 hessian=-11.0069745057 FD(grad)=-11.0069745043

奇异值互不相同，三种写法和梯度的中心差分（h = 1e-5）一致。照搬到 tenferro：

let t = TracedTensor::input_concrete_shape(DType::F64, &[])?;
let l = svd_loss(&t, &rows[0], &rows[1])?;     // a.svd()，默认选项
let g = ad.grad(&l, &t)?;
let gg = ad.grad(&g, &t)?;                     // grad(grad)
let gj = ad.jvp(&g, &t, &one)?;                // jvp(grad)

grad=5.924521608687 grad(grad)=-10.3648209072 jvp(grad)=-11.0069745057 FD(grad)=-11.0069745040
grad(grad) rel. error vs FD = -5.834e-2, jvp(grad) rel. error vs FD = 1.486e-10

一阶导数和 JAX 打印出的小数点后 12 位相同，grad(grad) 却偏了 5.8%，没有任何报错或警告；jvp(grad) 和 JAX 打印出的小数点后 10 位相同。

出处

tenferro 的 SVD 求导规则 linearize_svd 用洛伦兹展宽的 F = d/(d² + ε²)（默认 ε = 1e-12）处理接近简并的奇异值，对角线不屏蔽。F[i,i] 的值是 0，但它对 d 的斜率在 0 处是 1/ε²。前向模式里这个斜率乘的切向量正好是 0；反向套反向时，它乘的是一个不为 0 的余切，产生两份 1/ε² 量级、数学上相互抵消的贡献，舍入之后抵消不净。JAX 内置规则不展宽，所以没有这个问题；用 custom_jvp 写同样的展宽规则时 JAX 也会偏（−11.49），屏蔽对角线就对了。完整的对照表、ε 的影响和推导见高阶导数的「经过 SVD」一节，规则本身见 tenferro 卷的线性代数。

SvdOptions 只有 derivative_eps，没有屏蔽对角线的开关；把 ε 调大只能减小偏差（高阶导数：ε = 1e-6 时 −11.0138）；ε 大到 1e-3 时，连正确的二阶导数也被展宽改动。

修法

二阶导数用「先反向、再前向」的 jvp(grad)：

let one = TracedTensor::from_vec_col_major(vec![], vec![1.0_f64])?;
let gj = ad.jvp(&g, &t, &one)?;

输出见上：−11.0069745057，和差分的相对差 1.5e-10（差分本身的截断误差量级），和 JAX 的三种写法在小数点后 10 位上相同。多维输入时，Hessian 按列用 jvp(grad) 拼（一列一个基向量切向量），hvp 也用 jvp(grad)，写法见高阶导数。

怎么确认没中招

• 二阶导数的程序里出现了 SVD（或用了同样展宽的分解）时，搜一下有没有 ad.grad(&ad.grad(..)..) 这种叠法。
• 用梯度的中心差分核对一次：(g(t+h) − g(t−h)) / 2h，h = 1e-5 时相对差应在 1e-9 以下；上面 grad(grad) 的相对差是 6e-2。
• 只有一阶导数的程序不受影响：上面的 grad 和 JAX 在小数点后 12 位上相同。

小结

• JAX 内置 SVD 上 grad(grad)、jvp(grad)、jax.hessian 都对；tenferro 默认 SVD 规则上 grad(grad) 静默偏 5.8%。
• 原因是展宽后的 F 在对角线上的斜率 1/ε²，只在反向套反向时乘到非零余切。
• 经过 SVD 的二阶导数用 jvp(grad)，并用梯度的差分核对一次。