52 lines
3.1 KiB
Python
52 lines
3.1 KiB
Python
"""MeanFlow 损失（Geng et al. 2025；MP1 配方：25% 样本 r≠t，自适应损失加权）。
|
||
约定：t=0 为数据、t=1 为噪声；z_t = (1−t)·x + t·ε，瞬时速度 v = ε − x；网络输出平均速度 u(z,r,t)。
|
||
目标 u_tgt = v − (t − r)·du/dt，du/dt = v·∂_z u + ∂_t u 用 torch.func.jvp 一次前向得到（切向量 (v, 0, 1)）。
|
||
2026-10-18 首次正式预训练后改：
|
||
- 整个 MF 分支（解码器前向 + JVP）在 fp32 下算：目标里含网络自己的导数，bf16 下 JVP 误差直接变成目标噪声；
|
||
- return_aux=True 时返回 (z, t, r, u)，训练用 r==t 且 t 小的样本做"配对去噪估计" x̂ = z − t·u 来加 FK / 平滑损失。
|
||
原来把 FK 加在 one_step(随机噪声) 上、对齐这一条样本的老师 Y：对任意噪声都回归同一个 Y，最优解是条件均值 E[Y|X]，
|
||
多峰头被训成单峰（实测 last.pt 8 个候选间标准差 0.030，动作幅度 0.151），左绕/右绕取平均正好穿柱。
|
||
x̂ 由 z_t 里带着的那一条 x 决定（t 小时 z_t 几乎就是 x），不会把不同模式拉到一起。"""
|
||
import torch
|
||
from torch.func import jvp
|
||
|
||
|
||
def sample_rt(B, device, dtype, p_neq=0.25, mu=-0.4, sigma=1.0, p_edge=0.0):
|
||
"""2026-10-19（审查 review_model_train P0-3）：部署只在 (r=0, t=1) 推理一次，而 logit-normal 取 max/min 后
|
||
P(t>0.95) 约 0.08%、P(r<0.05) 约 0.5%，这个角几乎从未被训练到（一步残差与噪声相关 −0.86 即外推失准）。
|
||
p_neq > 0 时另以 p_edge 的概率把样本直接放在 (r=0, t=1)；预热期（p_neq=0）不放，那时一步映射本就不训。"""
|
||
t = torch.sigmoid(torch.randn(B, device=device, dtype=dtype) * sigma + mu)
|
||
r = torch.sigmoid(torch.randn(B, device=device, dtype=dtype) * sigma + mu)
|
||
t, r = torch.maximum(t, r), torch.minimum(t, r)
|
||
eq = torch.rand(B, device=device) >= p_neq
|
||
r = torch.where(eq, t, r)
|
||
if p_neq > 0 and p_edge > 0:
|
||
edge = torch.rand(B, device=device) < p_edge
|
||
t = torch.where(edge, torch.ones_like(t), t); r = torch.where(edge, torch.zeros_like(r), r)
|
||
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):
|
||
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()
|
||
err = (u - u_tgt).pow(2).mean(dim=(1, 2))
|
||
w = 1.0 / (err.detach() + ada_c).pow(ada_p)
|
||
loss = (w * err).mean()
|
||
if return_aux:
|
||
return loss, err.mean().detach(), {'z': z, 't': t, 'r': r, 'u': u}
|
||
return loss, err.mean().detach()
|
||
|
||
|
||
def one_step(model, mem, e):
|
||
B = e.shape[0]
|
||
one = torch.ones(B, device=e.device, dtype=e.dtype)
|
||
return e - model.decode(mem, e, torch.zeros_like(one), one)
|