feat: 目标写法开关 mf_target(mf/lin/sq/mid)，默认 mf 且原地路径逐位不变；float32 接线用解析高斯场自证

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()