fix(train): 首次正式预训练(21850 步闭环到达 0%)排查后修正——FK/平滑改加在配对去噪估计上(原加在随机噪声一步采样上=回归条件均值,多峰塌成单峰,实测候选标准差 0.030/幅度 0.151);时间嵌入最高角频率 1000→50、MF 分支 fp32、p_neq 前 5000 步为 0(实测 du/dt 均方根 6–12 vs v 1.1);增广 40%/3°/12 步收回;新增 dagger 阶段(策略闭环状态请老师重规划 ×2 轮微调);开训前老师轨迹回放自检;评估加开环误差与候选多样性,best/早停按到达率+开环误差
This commit is contained in:
+24
-13
@@ -1,6 +1,12 @@
|
||||
"""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))。"""
|
||||
目标 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
|
||||
|
||||
@@ -14,18 +20,23 @@ def sample_rt(B, device, dtype, p_neq=0.25, mu=-0.4, sigma=1.0):
|
||||
return r, t
|
||||
|
||||
|
||||
def meanflow_loss(model, mem, x, p_neq=0.25, ada_p=1.0, ada_c=1e-3):
|
||||
B = x.shape[0]
|
||||
r, t = sample_rt(B, x.device, x.dtype, p_neq)
|
||||
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)
|
||||
return (w * err).mean(), err.mean().detach()
|
||||
def meanflow_loss(model, mem, x, p_neq=0.25, ada_p=1.0, ada_c=1e-3, return_aux=False):
|
||||
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)
|
||||
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):
|
||||
|
||||
+9
-2
@@ -55,10 +55,17 @@ class DecBlock(nn.Module):
|
||||
return x + self.ff(self.n3(x))
|
||||
|
||||
|
||||
TEMB_WMAX = 50.0 # 时间嵌入最高角频率(rad / 单位 t)
|
||||
|
||||
|
||||
def time_emb(t, dim):
|
||||
"""2026-10-18 首次正式预训练后改(step 21850 闭环到达 0%):原实现照扩散模型写成 sin(1000·t·f),f∈[1e-3, 1],
|
||||
最高角频率 1000 —— t∈[0,1] 上网络对 t 极其敏感,实测 last.pt 的 ∂u/∂t 均方根 6–12,而 v 只有约 1.1,
|
||||
MeanFlow 目标 u_tgt = v − (t−r)·du/dt 被自激的 JVP 项主导:MF 残差 2 万步一直停在 7–8、17k–19k 尖峰到 96、梯度 100% 被裁剪。
|
||||
改为角频率在 [0.5, 50] 几何分布(最快一维在 [0,1] 上约转 8 圈,足够分辨 t,导数量级降 20 倍)。"""
|
||||
half = dim // 2
|
||||
fr = torch.exp(-math.log(1000.0) * torch.arange(half, device=t.device, dtype=t.dtype) / half)
|
||||
a = t[:, None] * 1000.0 * fr[None]
|
||||
w = torch.exp(torch.linspace(math.log(0.5), math.log(TEMB_WMAX), half, device=t.device, dtype=t.dtype))
|
||||
a = t[:, None] * w[None]
|
||||
return torch.cat([torch.sin(a), torch.cos(a)], -1)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user