Rust in Everything
基于 tenferro-linalg 0.8.0

简并奇异值

主线位置：TRG 主线的「分段求导得内能」，反向穿过每一步的两次 SVD；本页看奇异值成对相等时，各家的 SVD 求导规则给出什么。

版本：JAX 0.12.0.1（CPU，x64）、PyTorch 2.15.0（CPU，torch.set_num_threads(1)）、tenferro-rs 0.8.0（CPU，native，单线程）

flowchart TB
  s["s_i 等于 s_j"] --> d["平方差为 0"]
  d --> jx["JAX torch 1/d"]
  d --> tf["tenferro 展宽"]
  jx --> inf["inf"]
  inf --> nanv["乘 0 得 NaN"]
  tf --> zero["F 为 0"]
  zero --> ok["梯度有限"]

SVD 的反向规则里有一个奇异值两两耦合的系数矩阵 F。教科书写法是 F[i,j] = 1/(s_j² − s_i²)，两个奇异值相等时除以 0。JAX 内置规则用的是这个写法；PyTorch 的反向是 C++ 写的，这里没有读源码，但下面的实验里它的表现和 1/d 一致（逐位相等时 NaN，差一个 ulp 时有限）；tenferro 默认把它展宽成 (s_j² − s_i²)/((s_j² − s_i²)² + ε²)，ε = 1e-12，相等时恰好是 0。展宽的公式和它在 tenferro 里的实测见 线性代数 的「SVD 的反向规则」，这里只做对照。

实验代码：trg_onset.py、degenerate.py、fma_probe.py、fma_trg.py（JAX、PyTorch），tenferro 一侧是 bin degenerate 和 trg_staged（release，CpuEngine::serial()）。M2 Max、macOS 27、rustc 1.99.0。

现象：TRG 第 5 步起梯度是 NaN

TRG 主线，χ = 16，截断数按秩取（k = 2, 4, 4, 16, …），β = 0.4，对 lnZ/N 求导，步数从 1 加到 8。JAX 一侧三种 SVD：内置 jnp.linalg.svd；custom_jvp 写的和 tenferro 相同的展宽规则（下文）；PyTorch 的 torch.linalg.svd 加 backward。

jax 0.12.0.1 torch 2.15.0 chi=16 rank schedule beta=0.4
steps=1 ks=[2] grad builtin=1.843337108813 lorentz=1.843337108813 torch=1.843337108813
steps=2 ks=[2, 4] grad builtin=1.602167517441 lorentz=1.602167517441 torch=1.602167517441
steps=3 ks=[2, 4, 4] grad builtin=1.499230514558 lorentz=1.499230514558 torch=1.499230514558
steps=4 ks=[2, 4, 4, 16] grad builtin=1.379116482259 lorentz=1.379116482259 torch=1.379116482259
steps=5 ks=[2, 4, 4, 16, 16] grad builtin=nan lorentz=1.295053391807 torch=nan
steps=6 ks=[2, 4, 4, 16, 16, 16] grad builtin=nan lorentz=1.222308598298 torch=nan
steps=7 ks=[2, 4, 4, 16, 16, 16, 16] grad builtin=nan lorentz=1.167575613423 torch=nan
steps=8 ks=[2, 4, 4, 16, 16, 16, 16, 16] grad builtin=nan lorentz=1.131237714886 torch=nan

tenferro 不需要任何改动，用的就是默认的 a.svd()（分段写法见 分段）：

tenferro 0.8.0 staged chi=16 steps=5 schedule=rank width=1 ks=[2, 4, 4, 16, 16]
  u=-1.295053391806669 exact=-1.106079203745793 err=-1.890e-1 | grad=1.295053391806669 fd=1.295053391542345 |grad-fd|=2.643e-10
tenferro 0.8.0 staged chi=16 steps=8 schedule=rank width=1 ks=[2, 4, 4, 16, 16, 16, 16, 16]
  u=-1.131237714886462 exact=-1.106079203745793 err=-2.516e-2 | grad=1.131237714886462 fd=1.131237716411881 |grad-fd|=1.525e-9

（这里 u 和 Onsager 精确解差得多，是因为只做了 5 步、8 步，格子还小；20 步的对照见 动手：Ising 模型的 TRG。）

  • 前 4 步三家逐位相同（打印的 12 位），JAX 内置规则和展宽规则在这里没有区别。
  • 第 5 步起 JAX 内置和 PyTorch 都是 nan，之后每一步都是。
  • tenferro 的梯度和 JAX 展宽规则在打印的 12 位上相同，和中心差分差 1e-10 到 1e-9。20 步全程也是这样，见 动手。

为什么是第 5 步

打印每一步被分解的矩阵 ma 的奇异值，看有没有恰好相等的两个：

step 0: n=4 min s=7.624e-18 exactly equal pairs=0 min|s_j^2-s_i^2| off-diag=4.075e-34 1/that=2.454e+33
step 1: n=4 min s=6.037e-02 exactly equal pairs=0 min|s_j^2-s_i^2| off-diag=6.772e-03 1/that=1.477e+02
step 2: n=16 min s=1.632e-19 exactly equal pairs=0 min|s_j^2-s_i^2| off-diag=6.665e-38 1/that=1.500e+37
step 3: n=16 min s=4.563e-05 exactly equal pairs=0 min|s_j^2-s_i^2| off-diag=3.474e-09 1/that=2.878e+08
step 4: n=256 min s=4.261e-17 exactly equal pairs=25878 min|s_j^2-s_i^2| off-diag=0.000e+00 1/that=inf
step 5: n=256 min s=1.970e-17 exactly equal pairs=1 min|s_j^2-s_i^2| off-diag=0.000e+00 1/that=inf

第 5 步（下标 4）的矩阵是 256×256，只留前 16 个奇异值，其余 240 个在舍入意义上是 0。LAPACK 把其中 228 个算成了同一个数 7.124e-17，两两配对就是 25878 对恰好相等的奇异值，s_j² − s_i² 恰好是 0。前几步也有 1e-34、1e-38 这样小的差，1/d 到 1e33、1e37，大，但是有限，没有出 NaN。

JAX 0.12.0.1 的 _svd_jvp_rule（jax/_src/lax/linalg.py）这样算 F：

s_diffs = (s_dim + _T(s_dim)) * (s_dim - _T(s_dim))
s_diffs_zeros = lax._eye(s.dtype, (s.shape[-1], s.shape[-1]))  # jnp.ones((), dtype=A.dtype) * (s_diffs == 0.)  # is 1. where s_diffs is 0. and is 0. everywhere else
s_diffs_zeros = lax.expand_dims(s_diffs_zeros, range(s_diffs.ndim - 2))
F = 1 / (s_diffs + s_diffs_zeros) - s_diffs_zeros

注释里写的是「s_diffs 为 0 的地方是 1」，代码里换成了单位矩阵，只挡住了对角线。非对角线上两个奇异值相等时，F 是 inf。这些奇异向量都被截断丢掉了，它们收到的余切是 0，可 inf · 0 是 NaN，经过 U @ (F ∘ …) 的矩阵乘扩散到整个梯度，再一路传回 β。PyTorch 在这条链上也是 NaN。tenferro 的 F 在 d = 0 时是 0/(0 + ε²) = 0，这些耦合直接消失。

最小复现：4×4，奇异值 (3, 2, 2, 1)

两个矩阵：对角的 diag(3, 2, 2, 1)，和用两个 Householder 反射转过的同一组奇异值。三个损失都和简并块里怎么选基无关，在这一点可导：Σ wᵢsᵢ（重复的两个权重相同），最大奇异向量的投影 P₁ = u₁u₁ᵀ 的加权和（3 和 2 之间有间隙），前三个奇异向量张成的投影 P₃ 的加权和（整个简并块在里面）。

== diag(3,2,2,1): LAPACK s = [3.0, 2.0, 2.0, 1.0], s1 - s2 = 0.000e+00
  loss sum w*s
    jax builtin       : nan= 0 max|g|=1.000e+00 max|g-fd|=1.398e-10
    jax builtin jit   : nan= 0 max|g|=1.000e+00 max|g-fd|=1.398e-10
    torch             : nan= 0 max|g|=1.000e+00 max|g-fd|=1.398e-10
    jax lorentz jit   : nan= 0 max|g|=1.000e+00 max|g-fd|=1.398e-10
  loss P1
    jax builtin       : nan=16 max|g|=nan max|g-fd|=nan
    jax builtin jit   : nan=16 max|g|=nan max|g-fd|=nan
    torch             : nan=16 max|g|=nan max|g-fd|=nan
    jax lorentz jit   : nan= 0 max|g|=6.419e-01 max|g-fd|=7.816e-11
  loss P3
    jax builtin       : nan=16 max|g|=nan max|g-fd|=nan
    jax builtin jit   : nan=16 max|g|=nan max|g-fd|=nan
    torch             : nan=16 max|g|=nan max|g-fd|=nan
    jax lorentz jit   : nan= 0 max|g|=5.493e-01 max|g-fd|=6.610e-11
== rotated (3,2,2,1): LAPACK s = [2.9999999999999996, 2.000000000000001, 2.000000000000001, 1.0000000000000002], s1 - s2 = 0.000e+00
  loss sum w*s
    jax builtin       : nan= 0 max|g|=8.568e-01 max|g-fd|=1.891e-09
    jax builtin jit   : nan= 0 max|g|=8.568e-01 max|g-fd|=1.891e-09
    torch             : nan= 0 max|g|=8.568e-01 max|g-fd|=1.891e-09
    jax lorentz jit   : nan= 0 max|g|=8.568e-01 max|g-fd|=1.891e-09
  loss P1
    jax builtin       : nan=16 max|g|=nan max|g-fd|=nan
    jax builtin jit   : nan=16 max|g|=nan max|g-fd|=nan
    torch             : nan= 0 max|g|=2.664e-01 max|g-fd|=3.615e-10
    jax lorentz jit   : nan= 0 max|g|=2.664e-01 max|g-fd|=3.615e-10
  loss P3
    jax builtin       : nan= 9 max|g|=nan max|g-fd|=nan
    jax builtin jit   : nan= 9 max|g|=nan max|g-fd|=nan
    torch             : nan= 0 max|g|=7.210e-01 max|g-fd|=1.498e-09
    jax lorentz jit   : nan= 0 max|g|=7.210e-01 max|g-fd|=1.498e-09

# tenferro，默认 a.svd()
== diag(3,2,2,1): faer s = [3.0, 2.0, 2.0, 1.0]
  loss sum w*s  tenferro default: nan= 0 max|g|=1.000e0 max|g-fd|=3.273e-10
  loss P1       tenferro default: nan= 0 max|g|=6.419e-1 max|g-fd|=7.816e-11
  loss P3       tenferro default: nan= 0 max|g|=5.493e-1 max|g-fd|=6.610e-11
== rotated (3,2,2,1): faer s = [3.0, 2.0, 1.9999999999999998, 1.0]
  loss sum w*s  tenferro default: nan= 0 max|g|=8.568e-1 max|g-fd|=9.213e-10
  loss P1       tenferro default: nan= 0 max|g|=2.664e-1 max|g-fd|=2.345e-10
  loss P3       tenferro default: nan= 0 max|g|=7.210e-1 max|g-fd|=1.213e-7
  • 只用奇异值（Σ wᵢsᵢ）四家都对：ds = diag(Uᵀ dA V)，这个损失的梯度不经过 F。
  • 用到奇异向量，只要两个奇异值逐位相等，JAX 内置规则就是 NaN，jit 不 jit 一样；PyTorch 在对角矩阵上也是 NaN。tenferro 和 JAX 的展宽规则都和差分一致。
  • tenferro 和 JAX 展宽规则的梯度逐元素比：对角矩阵上差 0 到 4e-18；转过的矩阵上 Σ wᵢsᵢ、P₁ 差 1e-15 量级，P₃ 差 1.2e-7。后者出在 faer 算出的两个 2 差一个 ulp（2.2e-16）：s_j² − s_i² 约 9e-16，远小于 ε，F ≈ d/ε² ≈ 9e8，这个大系数很可能在对称的两项之间没有抵消干净；LAPACK 那边两个值逐位相等，F 恰好是 0。这个量级的分析见 线性代数 的动脑筋答案。
📝NOTE
没有蠢问题

问： 转过的矩阵上 PyTorch 没有 NaN，梯度还是对的，JAX 内置却是 NaN。PyTorch 是不是用了更稳的规则？

答： 不是规则的区别，是奇异值本身的区别。同一个矩阵，三家的分解给出：

torch s ['3.0', '2.000000000000001', '2.0', '0.9999999999999999'] diff 8.881784197001252e-16
jax s ['3.0', '2.000000000000001', '2.000000000000001', '1.0000000000000002'] diff 0.0
numpy s ['3.0', '2.000000000000001', '2.000000000000001', '1.0000000000000002']

PyTorch 的两个 2 差 8.9e-16，F 是一个约 2.8e14 的有限数，不是 inf，也就没有 inf · 0；结果和差分一致（误差 1e-9 量级）。JAX 拿到两个逐位相等的数，F 是 inf。所以用 1/d 规则时，出不出 NaN 取决于分解在最后一位上怎么舍入；对角输入这种奇异值精确相等的情形，JAX 和 PyTorch 都会撞上（上面对角矩阵的输出）。

JAX 里补上：custom_jvp

试验阶段的 JAX 参考实现用 jax.custom_jvp 包了一层 SVD，把 F 换成 tenferro 的展宽形式。规则只写 JVP，grad 由 JAX 转置得到（TRG 里的矩阵都是实方阵）：

EPS = 1e-12  # tenferro 的 DEFAULT_DECOMPOSITION_DERIVATIVE_EPS

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

@lorentz_svd.defjvp
def lorentz_svd_jvp(primals, tangents):
  (a,), (da,) = primals, tangents
  u, s, vh = lorentz_svd(a)
  v = vh.T
  dp = u.T @ da @ v
  ds = jnp.diagonal(dp)
  d = (s[None, :] - s[:, None]) * (s[None, :] + s[:, None])
  f = jnp.where(jnp.eye(s.shape[0]) > 0, 0.0, d / (d * d + EPS * EPS))
  dss = dp * s[None, :]
  sds = s[:, None] * dp
  du = u @ (f * (dss + dss.T))
  dv = v @ (f * (sds + sds.T))
  return (u, s, vh), (du, ds, dv.T)

tenferro 一侧什么都不用写：a.svd() 默认就是这个规则，ε 用 a.svd_with_options(SvdOptions::default().derivative_eps(e)) 改。上面这段比 tenferro 的规则少两样东西：复数时的规范修正项，和 m ≠ n 时的投影修正项，所以只适用于实方阵。

有两处写法不是随手的：

  • d 写成 (s_j − s_i)(s_j + s_i)，不写成 s_j**2 - s_i**2。原因在下一节。JAX 自己的内置规则也是这么写的（前面引的 s_diffs）。
  • 对角线用 where 强制为 0。F[i,i] = 0/(0 + ε²) 本来就是 0，但它对 d 的斜率是 1/ε²，反向套反向求二阶导数时会把结果带偏，tenferro 默认规则没有这一步，grad(grad) 偏 6%；这是 高阶导数 的话题。一阶导数上，where 还顺带挡住了下一节的问题。

jit 的 FMA 陷阱

把 d 写成平方差、对角线不屏蔽，F 的对角线应当是 0/(0 + ε²) = 0。在 jit 里，它不是 0。TRG 的梯度（β = 0.4）：

steps=5 fd=1.295053391531
  squares  masked=False jit grad=-3.050596362613e+20  eager grad=1.295053391807e+00
  product  masked=False jit grad=1.295053391807e+00  eager grad=1.295053391807e+00
  squares  masked=True  jit grad=1.295053391807e+00  eager grad=1.295053391807e+00
  product  masked=True  jit grad=1.295053391807e+00  eager grad=1.295053391807e+00
steps=20 fd=1.106045408322
  squares  masked=False jit grad=1.396314156080e+120
  product  masked=False jit grad=1.106045407614e+00
  squares  masked=True  jit grad=1.106045407614e+00
  product  masked=True  jit grad=1.106045407614e+00

平方差加 jit：5 步时梯度 −3e20，20 步时 1.4e120；不 jit 就是对的。换成乘积形式、或者屏蔽对角线，任何一个都够。

看对角线上到底是什么。在同一个 jit 里做 SVD 和平方差，把 d 的对角线和 s 一起返回：

def sq(a):
  s = jnp.linalg.svd(a, compute_uv=False)
  return s[None, :] ** 2 - s[:, None] ** 2, s

d_j, s_j = jax.jit(sq)(a)
print([math.fma(x, x, -(x * x)) for x in s_j])   # Python 3.14 的 math.fma

  s (jit)                        = [3.0348734488811724, 1.6885475440472182, 1.1555909458038691, 0.3182721502403828]  same as eager: True
  eager diag(s_j**2 - s_i**2)    = [0.0, 0.0, 0.0, 0.0]
  jit   diag(s_j**2 - s_i**2)    = [-5.712212290647009e-16, -1.2505230706026656e-16, 8.384298537456344e-17, -3.484010694261065e-18]
  math.fma(s, s, -(s*s))         = [-5.712212290647009e-16, -1.2505230706026656e-16, 8.384298537456344e-17, -3.484010694261065e-18]
  jit   diag((s_j-s_i)(s_j+s_i)) = [0.0, 0.0, 0.0, 0.0]
  eps=1e-12: jit diag F = d/(d^2+eps^2) = [-571221042.6788777, -125052305.10468869, 83842984.78517693, -3484010.6942187757]
  eps=1e-06: jit diag F = d/(d^2+eps^2) = [-0.0005712212290647009, -0.00012505230706026657, 8.384298537456344e-05, -3.4840106942610654e-06]

jit 下的对角线和 fma(s, s, −fl(s·s)) 逐位相同：XLA 把 s_j·s_j − s_i·s_i 编成了一条融合乘加，一边是精确的 s_j²，另一边是舍入过的 s_i²，差就是 s_i² 的舍入误差，1e-16 量级。ε = 1e-12 时 F 的对角线就是 d/ε²，到 5.7e8；正常的非对角元是 0.1 量级。这个大系数乘上 dP 的对角线进入 dU，一步步放大到 1e120。ε = 1e-6 时对角线只有 5.7e-4，前面实验里也只偏了一点。乘积形式 (s_j − s_i)(s_j + s_i) 的第一个因子在对角线上是 s − s，精确的 0，乘什么都是 0。

tenferro 的规则也是乘积形式（s_gap = s_sum * s_diff，见 线性代数），对角线上的 s_diff 是精确的 0。怎么读这一卷 里看到的 1 ulp 差别，就是同一种融合。

❗IMPORTANT
动脑筋：把 SVD 放在 jit 外面，只把 s 当参数传进 jit(lambda s: s[None, :] ** 2 - s[:, None] ** 2)，对角线还是那四个 1e-16 量级的数吗？

所以写自定义的 SVD 规则，凡是「数学上应当恰好为 0」的量，都写成精确的形式（s − s 再乘），或者用 where 明确置 0；不要指望 x*x - y*y 在 x == y 时得到 0。

动脑筋答案

  s as jit argument: diag        = [0.0, 0.0, 0.0, 0.0]
  • 是 0。同一个表达式，s 是 jit 的参数时 XLA 没有生成那条乘加，s 和平方差在同一个 jit 里由 SVD 算出来时才生成。融合不融合由编译器在整个计算里的安排决定，不能从这一行代码看出来。
  • 所以这个坑很难用小例子提前发现：把规则单独拿出来测，s 往往是参数，测不出问题；放进 TRG 的 jit 里才出现。写成乘积形式或者屏蔽对角线，就不依赖编译器怎么安排。

小结

  • TRG 第 5 步的 256×256 矩阵里，LAPACK 把 228 个舍入意义上的零奇异值算成同一个数，1/(s_j² − s_i²) 是 inf，乘被截断方向上的零余切得到 NaN。JAX 内置 jnp.linalg.svd 和 PyTorch 从这一步起梯度全是 NaN。
  • tenferro 默认的展宽 F = d/(d² + ε²)（ε = 1e-12）在 d = 0 时给 0，TRG 的梯度有限且和差分一致，和 JAX 补上同样规则后的结果 12 位相同。
  • 损失只用奇异值时，内置规则没事；用到奇异向量、又有逐位相等的奇异值时才 NaN。PyTorch 在转过的矩阵上没出 NaN，是因为它的两个奇异值碰巧差一个 ulp。
  • JAX 里用 custom_jvp 把 F 换成展宽形式；d 写成 (s_j − s_i)(s_j + s_i)，对角线用 where 置 0。写成 s_j**2 - s_i**2 时，jit 会生成 FMA，对角线变成 1e-16 量级，TRG 的梯度到 1e120。

💬 关联讨论

发起讨论
加载中...