173 lines
8.8 KiB
Python
173 lines
8.8 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
|
||
from torch.func import jvp
|
||
|
||
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'])
|
||
|
||
|
||
class AvgVelocityTargetWiringTest(unittest.TestCase):
|
||
"""把三种写法接进 ml/model/meanflow.py 之后，再用同一个解析场验一遍接线：
|
||
浮点从 float64 换到 float32，量级判据放宽，但"mf 无偏 / sq 只线性收敛"必须照样成立，
|
||
否则说明 z_c、切向量、r 槽换位这几处接错了。"""
|
||
SIG = 2.0
|
||
|
||
class Stub:
|
||
def __init__(self, sig2):
|
||
self.sig2 = sig2
|
||
|
||
def s2(self, t):
|
||
return (1 - t) ** 2 * self.sig2 + t ** 2
|
||
|
||
def decode(self, mem, z, r, t):
|
||
return z * (1 - torch.sqrt(self.s2(r) / self.s2(t)))[:, None, None] / (t - r)[:, None, None]
|
||
|
||
def bias(self, target, r, t, lam=0.5):
|
||
from ml.model.meanflow import avg_velocity_target
|
||
f = Field(self.SIG)
|
||
z = torch.tensor([[0.7] * 4, [0.7] * 4], dtype=torch.float32) # [2,4,6]? 形状自定，逐元素即可
|
||
z = torch.full((2, 4, 6), 0.7)
|
||
rt = torch.tensor([r, r]); tt = torch.tensor([t, t])
|
||
xh = (z * (1 - tt)[:, None, None] * self.SIG ** 2 / f.s2(tt)[:, None, None])
|
||
eh = (z * tt[:, None, None] / f.s2(tt)[:, None, None])
|
||
u, tgt, _ = avg_velocity_target(self.Stub(self.SIG ** 2), None, xh, eh, rt, tt, target, lam)
|
||
truth = f.u(torch.tensor(0.7, dtype=torch.float32), torch.tensor(r), torch.tensor(t))
|
||
return float(((tgt - truth).abs() / truth.abs()).mean()), float(truth)
|
||
|
||
def test_wiring_keeps_mf_unbiased_and_sq_first_order(self):
|
||
b_mf, _ = self.bias('mf', 0.8, 1.0)
|
||
self.assertLess(b_mf, 1e-4)
|
||
rel = []
|
||
for s in (0.2, 0.1, 0.05):
|
||
rel.append({k: self.bias(k, 1.0 - s, 1.0)[0] for k in ('sq', 'lin', 'mid', 'mf')})
|
||
print('float32 接线后的偏差：', [{k: round(v, 4) for k, v in d.items()} for d in rel])
|
||
r_sq = rel[1]['sq'] / rel[0]['sq']
|
||
r_mid = rel[1]['mid'] / rel[0]['mid']
|
||
self.assertGreater(r_sq, 0.3) # 只线性收敛 ⇒ 结构性偏差
|
||
self.assertLess(r_mid, 0.35) # 二次收敛 ⇒ 只剩积分误差
|
||
self.assertLess(rel[0]['lin'], rel[0]['sq'])
|
||
|
||
def test_default_path_is_bitwise_identical(self):
|
||
"""target='mf' 必须与重构前的原地写法逐位相同（默认训练路径不许被这次改动碰到）"""
|
||
from ml.model.meanflow import avg_velocity_target
|
||
torch.manual_seed(0)
|
||
m = self.Stub(2.0)
|
||
x = torch.randn(3, 5, 6); e = torch.randn(3, 5, 6)
|
||
r = torch.rand(3) * 0.4; t = 0.6 + torch.rand(3) * 0.4
|
||
u_new, tgt_new, z_new = avg_velocity_target(m, None, x, e, r, t, 'mf', 0.5)
|
||
fn = lambda a, b, c: m.decode(None, a, b, c)
|
||
z = (1 - t)[:, None, None] * x + t[:, None, None] * e
|
||
v = e - x
|
||
u0, dudt = jvp(fn, (z, r, t), (v, torch.zeros_like(r), torch.ones_like(t)))
|
||
tgt0 = (v - (t - r)[:, None, None] * dudt).detach()
|
||
for name, a_, b_ in (('u', u_new, u0), ('u_tgt', tgt_new, tgt0), ('z', z_new, z)):
|
||
with self.subTest(name=name):
|
||
torch.testing.assert_close(a_, b_, rtol=0, atol=0)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
unittest.main()
|