From 3b8682a34544f71334b3cce93b2db1a300e77463 Mon Sep 17 00:00:00 2001 From: cbb <2625274797@qq.com> Date: Tue, 20 Oct 2026 21:48:28 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E7=9B=AE=E6=A0=87=E5=86=99=E6=B3=95?= =?UTF-8?q?=E5=BC=80=E5=85=B3=20mf=5Ftarget(mf/lin/sq/mid)=EF=BC=8C?= =?UTF-8?q?=E9=BB=98=E8=AE=A4=20mf=20=E4=B8=94=E5=8E=9F=E5=9C=B0=E8=B7=AF?= =?UTF-8?q?=E5=BE=84=E9=80=90=E4=BD=8D=E4=B8=8D=E5=8F=98=EF=BC=9Bfloat32?= =?UTF-8?q?=20=E6=8E=A5=E7=BA=BF=E7=94=A8=E8=A7=A3=E6=9E=90=E9=AB=98?= =?UTF-8?q?=E6=96=AF=E5=9C=BA=E8=87=AA=E8=AF=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ml/model/meanflow.py | 41 +++++++++++++++++++---- ml/train.py | 3 +- tests/test_meanflow_targets.py | 60 ++++++++++++++++++++++++++++++++++ 3 files changed, 97 insertions(+), 7 deletions(-) diff --git a/ml/model/meanflow.py b/ml/model/meanflow.py index 4d4e3df..fd540dd 100644 --- a/ml/model/meanflow.py +++ b/ml/model/meanflow.py @@ -26,17 +26,46 @@ def sample_rt(B, device, dtype, p_neq=0.25, mu=-0.4, sigma=1.0, p_edge=0.0): return r, t -def meanflow_loss(model, mem, x, p_neq=0.25, ada_p=1.0, ada_c=1e-3, return_aux=False, p_edge=0.0): +def avg_velocity_target(model, mem, x, e, r, t, target='mf', lam=0.5): + """平均速度目标的三种写法（目标都 detach；u 带梯度）。target='mf' 与原实现逐位相同。 + mf ：Eq.4 单段自举 v − (t−r)·du/dt，收敛解 = 平均速度（解析高斯场已验证偏差 <1e−12） + lin ：v − (t−r)·[(1−λ)·du(z_t,c,t)/dt + λ·du(z_c,r,c)/dc]，系数和=1 ⇒ 收敛解仍是平均速度 + sq ：同上的 (1−λ)²/λ² 版（K-MF 论文摘要里的写法）⇒ 系数和 0.5，实测只线性收敛（结构性少修一半） + mid ：v − (t−r)·du(z_c,r,c)/dc，整条修正移到中点 + """ + fn = lambda z_, r_, t_: model.decode(mem, z_, r_, t_) + z = (1 - t)[:, None, None] * x + t[:, None, None] * e + v = e - x + if target == 'mf': + u, dudt = jvp(fn, (z, r, t), (v, torch.zeros_like(r), torch.ones_like(t))) + return u, (v - (t - r)[:, None, None] * dudt).detach(), z + with torch.no_grad(): + u = fn(z, r, t) # 这三种写法只要前向拿 u，导数全部进目标（sg） + c = (1 - lam) * t + lam * r + zc = (1 - c)[:, None, None] * x + c[:, None, None] * e + tv = (v, torch.zeros_like(r), torch.ones_like(t)) + + def jv(zz, rr, tt): + with torch.enable_grad(): + return jvp(fn, (zz, rr, tt), tv)[1].detach() + + if target == 'mid': + corr = jv(zc, r, c) + else: + d1 = jv(z, c, t) + d2 = jv(zc, r, c) + w = ((1 - lam) ** 2, lam ** 2) if target == 'sq' else ((1 - lam), lam) + corr = w[0] * d1 + w[1] * d2 + return u, (v - (t - r)[:, None, None] * corr).detach(), z + + +def meanflow_loss(model, mem, x, p_neq=0.25, ada_p=1.0, ada_c=1e-3, return_aux=False, p_edge=0.0, target='mf', lam=0.5): mem, x = mem.float(), x.float() with torch.autocast('cuda', enabled=False): B = x.shape[0] r, t = sample_rt(B, x.device, x.dtype, p_neq, p_edge=p_edge) e = torch.randn_like(x) - z = (1 - t)[:, None, None] * x + t[:, None, None] * e - v = e - x - fn = lambda z_, r_, t_: model.decode(mem, z_, r_, t_) - u, dudt = jvp(fn, (z, r, t), (v, torch.zeros_like(r), torch.ones_like(t))) - u_tgt = (v - (t - r)[:, None, None] * dudt).detach() + u, u_tgt, z = avg_velocity_target(model, mem, x, e, r, t, target, lam) err = (u - u_tgt).pow(2).mean(dim=(1, 2)) w = 1.0 / (err.detach() + ada_c).pow(ada_p) loss = (w * err).mean() diff --git a/ml/train.py b/ml/train.py index b5ff0bd..34a955b 100644 --- a/ml/train.py +++ b/ml/train.py @@ -161,7 +161,8 @@ def main(): l_mf, xh = wta_loss(model, mem, Y); mf_raw = l_mf.detach(); xh = xh.float() else: p_neq = tc['p_neq'] if step >= warm else 0.0 - l_mf, mf_raw, aux = meanflow_loss(model, mem, Y, p_neq=p_neq, return_aux=True, p_edge=tc.get('p_edge', 0.0)) + l_mf, mf_raw, aux = meanflow_loss(model, mem, Y, p_neq=p_neq, return_aux=True, p_edge=tc.get('p_edge', 0.0), + target=tc.get('mf_target', 'mf'), lam=tc.get('mf_lam', 0.5)) sel = ((aux['r'] == aux['t']) & (aux['t'] <= tc.get('fk_tmax', 0.3))).nonzero(as_tuple=True)[0] tt = aux['t'][sel][:, None, None] xh = aux['z'][sel] - tt * aux['u'][sel] # 配对去噪估计（fp32） diff --git a/tests/test_meanflow_targets.py b/tests/test_meanflow_targets.py index 20540f6..c6f500f 100644 --- a/tests/test_meanflow_targets.py +++ b/tests/test_meanflow_targets.py @@ -9,6 +9,7 @@ import unittest import torch +from torch.func import jvp LAM = 0.5 @@ -108,5 +109,64 @@ class MeanFlowTargetFixedPointTest(unittest.TestCase): 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()