113 lines
5.7 KiB
Python
113 lines
5.7 KiB
Python
"""目标公式的不动点判决（解析高斯场，float64，CPU，与任何已训模型无关）。
|
||
|
||
1-D：数据 x~N(0,σ²)、噪声 ε~N(0,1)、路径 z_t=(1−t)x+tε ⇒
|
||
s²(t)=(1−t)²σ²+t²，v(z,t)=E[ε−x|z]=z·(t−(1−t)σ²)/s²(t)。
|
||
因为 a(t)=v/z=(s²)′/(2s²)，PF-ODE 的流图有闭式 z_r = z_t·sqrt(s²(r)/s²(t)) ⇒ 平均速度真值也是闭式。
|
||
把"精确场"代入四种目标写法，看谁能收敛回真值。判据用**收敛阶**而不是绝对值：
|
||
有结构性偏差的写法，区间减半时误差只能减半（线性）；没有偏差的应当减到 1/4（二阶）。
|
||
"""
|
||
import unittest
|
||
|
||
import torch
|
||
|
||
LAM = 0.5
|
||
|
||
|
||
class Field:
|
||
def __init__(self, sig):
|
||
self.sig2 = sig ** 2
|
||
|
||
def s2(self, t):
|
||
return (1 - t) ** 2 * self.sig2 + t ** 2
|
||
|
||
def v(self, z, t):
|
||
return z * (t - (1 - t) * self.sig2) / self.s2(t)
|
||
|
||
def u(self, z, r, t):
|
||
"""平均速度真值：状态在时刻 t 的 z，要走到时刻 r"""
|
||
return (z - z * torch.sqrt(self.s2(r) / self.s2(t))) / (t - r)
|
||
|
||
def total_du(self, z, r, t, tv):
|
||
"""du(z,r,t) 沿切向量 tv 的全导数 = ∂_z u·tv + ∂_t u（r 槽固定）"""
|
||
zz = z.clone().requires_grad_(True)
|
||
tt = t.clone().requires_grad_(True)
|
||
rr = r.clone().requires_grad_(False)
|
||
gz, gt = torch.autograd.grad(self.u(zz, rr, tt), (zz, tt), grad_outputs=torch.ones((), dtype=z.dtype))
|
||
return gz * tv + gt
|
||
|
||
def targets(self, z, r, t, lam=LAM):
|
||
r = torch.as_tensor(r, dtype=torch.float64); t = torch.as_tensor(t, dtype=torch.float64)
|
||
z = torch.as_tensor(z, dtype=torch.float64)
|
||
v = self.v(z, t)
|
||
c = (1 - lam) * t + lam * r
|
||
xh = z * (1 - t) * self.sig2 / self.s2(t) # E[x|z_t]，eh = E[ε|z_t]；耦合速度 eh−xh = v
|
||
eh = z * t / self.s2(t)
|
||
z_c = (1 - c) * xh + c * eh
|
||
d_mf = self.total_du(z, r, t, v)
|
||
d1 = self.total_du(z, c, t, v) # 方案里的第一项：r 槽挪到 c，对 t 求导
|
||
d2 = self.total_du(z_c, r, c, v) # 第二项：状态与第二个时刻槽都在中点
|
||
span = t - r
|
||
return {'mf': v - span * d_mf,
|
||
'sq': v - span * ((1 - lam) ** 2 * d1 + lam ** 2 * d2),
|
||
'lin': v - span * ((1 - lam) * d1 + lam * d2),
|
||
'mid': v - span * d2}
|
||
|
||
def abs_rel(self, z, r, t):
|
||
"""返回 (相对真值的偏差, 真值本身)。σ=1 且 (r,t)=(0,1) 时 s²(0)=s²(1) ⇒ 流图是恒等、真值为 0，
|
||
那是个退化点，只能用绝对偏差，别拿它算相对值。"""
|
||
z_ = torch.as_tensor(z, dtype=torch.float64); r_ = torch.as_tensor(r, dtype=torch.float64); t_ = torch.as_tensor(t, dtype=torch.float64)
|
||
u = self.u(z_, r_, t_)
|
||
den = u.abs() if float(u.abs()) > 1e-9 else torch.tensor(1.0, dtype=torch.float64)
|
||
return {k: float((val - u).abs() / den) for k, val in self.targets(z, r, t).items()}, float(u)
|
||
|
||
|
||
def abs_err(self, z, r, t):
|
||
"""绝对偏差（除以固定尺度 |v(z,t)|，不除以会趋 0 的 |u_true|）"""
|
||
z_ = torch.as_tensor(z, dtype=torch.float64); r_ = torch.as_tensor(r, dtype=torch.float64); t_ = torch.as_tensor(t, dtype=torch.float64)
|
||
u = self.u(z_, r_, t_); sc = self.v(z_, t_).abs()
|
||
return {k: float((val - u).abs() / sc) for k, val in self.targets(z, r, t).items()}
|
||
|
||
|
||
class MeanFlowTargetFixedPointTest(unittest.TestCase):
|
||
def test_mf_identity_is_exact(self):
|
||
"""先证仪表：把精确场代 Eq.4 必须逐位还原平均速度，否则后面的数都白读（用绝对偏差，退化点也能过）"""
|
||
for sig in (0.5, 1.0, 2.0):
|
||
for (r, t) in ((0.0, 1.0), (0.0, 0.3), (0.2, 0.9)):
|
||
z_ = torch.tensor(0.7, dtype=torch.float64)
|
||
f = Field(sig)
|
||
tgt = f.targets(0.7, r, t)
|
||
with self.subTest(sig=sig, r=r, t=t):
|
||
self.assertLess(float((tgt['mf'] - f.u(z_, torch.tensor(r, dtype=torch.float64),
|
||
torch.tensor(t, dtype=torch.float64))).abs()), 1e-12)
|
||
|
||
def test_convergence_order_separates_biased_from_consistent(self):
|
||
"""区间减半：sq 只能线性收敛（系数和 0.5 ⇒ 结构性少修一半），lin/mid 应当二次收敛"""
|
||
for sig in (1.0, 2.0):
|
||
f = Field(sig)
|
||
rel = [f.abs_rel(0.7, 1.0 - s, 1.0)[0] for s in (0.2, 0.1, 0.05)]
|
||
print('σ=%.1f 相对偏差 span=0.2/0.1/0.05：' % sig,
|
||
[{k: round(v, 4) for k, v in d.items()} for d in rel])
|
||
for name, want in (('sq', 0.5), ('lin', 0.25), ('mid', 0.25)):
|
||
ratios = [rel[1][name] / rel[0][name], rel[2][name] / rel[1][name]]
|
||
with self.subTest(sig=sig, name=name):
|
||
for rt in ratios:
|
||
self.assertLess(rt, want * 1.5)
|
||
self.assertGreater(rt, want * 0.5)
|
||
|
||
def test_full_span_step_is_unreliable_for_every_variant(self):
|
||
"""跨整段 (0,1) 时除 Eq.4 外都偏得很厉害；跨半段就明显变小 ⇒ "一步跨完"本身是难的，
|
||
不是训练出来的（σ=2，偏差按 |v| 归一，避免除以会趋 0 的 |u_true|）"""
|
||
f = Field(2.0)
|
||
full = f.abs_err(0.7, 0.0, 1.0)
|
||
half = f.abs_err(0.7, 0.5, 1.0)
|
||
print('σ=2 整段 (0,1) 偏差：', {k: round(v, 3) for k, v in full.items()})
|
||
print('σ=2 半段 (0.5,1) 偏差：', {k: round(v, 3) for k, v in half.items()})
|
||
self.assertLess(full['mf'], 1e-12)
|
||
self.assertGreater(full['sq'], 0.5)
|
||
self.assertLess(half['sq'], full['sq'])
|
||
self.assertLess(half['mid'], full['mid'])
|
||
|
||
|
||
if __name__ == '__main__':
|
||
unittest.main()
|