Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis

1. 出发点 (Motivation)

2024 年的 REPA (Yu et al.) 给业内一个意外发现:flow / diffusion 模型如果把内部 feature 跟外部冻结的 image encoder (例如 86M 的 DINOv2-B) 对齐,生成质量能大幅提升。这意味着 flow 模型自己不会学强语义表征 — 必须借。整个 2025 年一票工作 (REPA, REPA-E, SRA 等) 都在这条路上。

但这条路有三个根本毛病:

  1. 反向 scaling. 换更强的 encoder 反而更差。DINOv2-B → DINOv2-L → DINOv3-B → DINOv3-H+ 上,FID 单调上升 (越大越差)。论文 §3.2 Fig 2a 实测出来的。
  2. 模型 scaling 也失效. REPA 把生成模型从 290M scale 到 1B,收益严重递减 — Self-Flow 的 625M 直接超过 REPA 的 1B。
  3. 跨模态崩溃. 视频对齐到 V-JEPA2 / Depth Anything 3,音频对齐到 MERT,都比 vanilla flow matching 还差。外部 encoder 的训练目标(识别 / 表征聚类)跟生成目标 (synthesis) 没对齐,迁移到非图像模态时主动有害。

诊断:把"表征学习"外包给一个跟生成目标不一致的、固定的外部模型,是个 bottleneck。应该把表征学习内嵌到生成目标里,让它自己学。

但 flow matching 训练目标 (denoise 一个均匀噪声的图) 本身没什么"必须学语义"的压力 — 局部相关性就能解大部分 denoising。Self-Flow 的核心 trick 是制造信息不对称:给同一张图的不同 token 加不同噪声,某些 token 很干净、某些很脏,模型为了把脏的还原回去必须用干净 token 的语义。这就强制 model 学全局表征。

Self-Flow hero result
Fig. 1 — 左:T2I 上 Self-Flow 收敛比 REPA 快约 2.8×,且 REPA 在 1.2M step 后开始 plateau,Self-Flow 仍在下降。右:在 vanilla flow matching 失败的 prompt 上("a parrot preens its brilliant feathers"、"LOVE" 写在指甲上、"woman doing push-ups"),Self-Flow 同时改善结构连贯性、文字渲染、时间一致性。

2. 方法 (Method)

核心思想 (类比)

把训练 flow 模型想象成考填空题。Flow matching 的填空题是:

Self-Flow 改成不对称填空题:

但只这样训练有train-inference gap:推理时每步噪声等级是均匀的,如果训练时全是不对称的,模型推理就崩。Self-Flow 的具体方案:采样两个时步 \(t, s\),用较干净的那个 \(\min(t, s)\) 给 mask 内的 token,另一个 \(t\) 给其他 token。这样每个 token 的边缘噪声分布仍跟 vanilla flow matching 一致,只是被 mask 的子集系统性地更干净 — 推理时不会有意外。

有了不对称输入,Self-Flow 做self-distillation:同一个模型扮 student (噪声混合输入) 和 teacher (EMA 副本, 用 min(t,s) 的更干净版本输入),让 student 的中间层 feature 跟 teacher 的对齐 — 跟 REPA 形式上一样,但 teacher 变成了模型自己更干净的副本而非外部 DINO。

Self-Flow architecture
Fig. 3 — 总体框架。给定干净输入 $x_0$(图像/视频/音频),采样两个时步 $t, s$ 和随机 mask $M$,按 mask 给每个 token 分配 $t$ 或 $s$ → 构造 $x_\tau$ (student 输入)。Teacher 输入用 $\tau_{\min} = \min(t, s)$,所有 token 都用较干净的时步,获得了 student 没有的信息。Student 同时做两件事:(1) 在 $x_\tau$ 上做 flow matching 的 denoising (生成损失 $\mathcal L_{\text{gen}}$),(2) 让自己的某个中间层 feature 跟 teacher 在 $x_{\tau_{\min}}$ 上的 feature 对齐(表征损失 $\mathcal L_{\text{rep}}$)。

2.1 Flow Matching + REPA 基础

Rectified flow (Liu et al. 2022). \(x_0\) 是干净数据 (N 个 token,每个维度 C),\(x_1 \sim \mathcal N(0, I)\) 是噪声。线性插值轨迹:

速度场是常向量 \(v_t = dx_t/dt = x_1 - x_0\)。生成损失(Eq 2):

REPA (Yu et al. 2024). 用一个冻结的外部 encoder \(r_\phi\) (DINOv2 等),把 flow 模型第 \(l\) 层 feature 跟 encoder 第 \(k\) 层 feature 对齐(Eq 3):

\(h_\theta\) 是 student 的 MLP projection head,\(\mathrm{sim}\) 是 cosine similarity。关键问题:\(r_\phi\) 是 DINO,训练目标不是生成。把生成模型的 feature 拉向 DINO 的 feature manifold 会创造一个 bottleneck — 强模型反而被锁死在 DINO 的 representation 里。

2.2 REPA 反向 scaling — 论文最有力的实证

REPA scaling + scheduler comparison
Fig. 2 — (a) REPA 反向 scaling:DINOv2-B (蓝) FID 最好,DINOv2-L (黄) 接近,DINOv3-B (绿) 显著更差,DINOv3-H+ (红) 最差。encoder 越强,生成质量越差。(b) Scheduler 对比:Vanilla Flow Matching 是基线;Diffusion Forcing (每个 token 独立噪声) 和 Full Masking (随机置 t=1) 都严重恶化生成 — 因为推理时模型遇到均匀噪声,跟训练完全不一样;Dual-Timestep Scheduling (黄) 没用任何 self-supervised loss,就已经小幅超过 vanilla — 说明信息不对称本身就是有益的。

Fig 2a 是这篇 paper 的 founding observation。这条规律之前 Singh et al. 2025 也报告过,但 Self-Flow 是第一个把它扩到 DINOv3 系列、且作为 motivation 系统化论证的工作。

2.3 Dual-Timestep Scheduling (§3.3) — 核心 trick

怎么"给不同 token 加不同噪声"又不引入 train-inference gap?三种 naive 方案的失败:

Self-Flow 的方案 (Eq 4-5):

  1. 采样两个时步 \(t, s \sim p(t)\)。
  2. 采样一个 mask \(M = \{i \in \{1, \ldots, N\} \mid u_i \lt R_M\}\),其中 \(u_i \sim U(0, 1)\),masking ratio \(R_M \leq 0.5\)。
  3. 构造 per-token 时步 \(\tau \in \mathbb R^N\): $\(\tau^i = \begin{cases} s, & i \in M \\ t, & \text{otherwise} \end{cases}\)$
  4. 每个 token 按自己的 \(\tau^i\) 加噪: $\(x_\tau = \mathrm{diag}(\mathbf 1 - \tau)\,x_0 + \mathrm{diag}(\tau)\,x_1\)$

关键性质: 因为 mask 是随机且 \(s\) 和 \(t\) 都从同一分布 \(p(t)\) 采样,任何一个 token 的边缘噪声分布 仍然是 \(p(t)\) — 跟 vanilla flow matching 完全一样。所以推理时模型遇到的输入分布跟训练时每个 token 看到的分布没差异。

但是 batch 内的联合分布变了:同一张图的 token 之间有了正相关 (mask 内的都用 \(s\),mask 外的都用 \(t\))。Student 要 denoise 噪声较大的那一组,必须用噪声较小的那一组当线索 — 强制学全局表征。

Fig 2b 的黄线表明:就算不加 self-supervised loss (仅改 scheduler),Dual-Timestep Scheduling 单独就能改善生成。这是一个非平凡的副效应。

2.4 EMA Teacher + Self-Flow Loss (§3.4)

有了不对称输入,Self-Flow 把 REPA 的外部 encoder 换成同一模型的 EMA 副本。设 \(f_\theta\) 是 student,\(f_{\theta'}\) 是 EMA teacher (动量更新)。Teacher 输入用 \(\tau_{\min} = \min(\tau) \in \{t, s\}\) — 即所有 token 都用更干净的那个时步,得到 \(x_{\tau_{\min}}\)。

表征对齐目标 (Eq 6):

跟 REPA 的形式一样,只是把 \(r_\phi\) 换成了 \(f_{\theta'}\)。层选择规则跟 REPA 一致:\(l < k\) (student 的早期层对齐 teacher 的中后期层,因为 diffusion 模型的语义 feature 集中在中后期)。

总目标 (Eq 7):

\(\gamma\) 是平衡系数。Teacher 通过 EMA 慢慢追 student:

\(\alpha \approx 0.999\) 量级。关键洞察:这条 self-distillation 之所以能 work,是因为 (a) teacher 看的输入本质上更容易 (噪声整体更小),(b) student 看的输入缺信息,(c) 强迫 student 在缺信息的情况下匹配 teacher 的 feature,就是在强迫 student 学 "如何从少量线索推全图" 的语义能力。

跟 REPA 比 — REPA 的 teacher (DINO) 看的图没有噪声且模型完全不同,这条 chain 走不通;Self-Flow 的 teacher 跟 student 是同一模型同一架构,差距只在"看到的输入有多干净",所以表征 chain 走得顺。这正是为什么 Self-Flow 可以 scale (Fig 6) 而 REPA 不行。

2.5 与代码对照

repo (BFL 官方) 只发布 inference 代码 + ImageNet 256 checkpoint,训练代码未公开。下面引的是模型架构里的 per-token 时步实现 — 这是 Dual-Timestep Scheduling 在前向时唯一需要的修改。

(a) per-token 时步 embedding — Eq 4 的 \(\tau \in \mathbb R^N\) 在 transformer 里怎么用?

repo/src/model.py:L455-L491 — SelfFlowDiT._forward (per-token timestep conditioning)

def _forward(self, x, t, y, return_features=False, return_raw_features=False):
    """Forward with per-token timestep conditioning."""
    assert not (return_raw_features and return_features)

    x = self.x_embedder(x) + self.pos_embed
    batch_size, seq_len, hidden_dim = x.shape

    # Handle timestep embedding - per-token or broadcast
    if t.ndim == 1:
        # Standard mode: same timestep for all tokens (compat. with vanilla flow matching)
        t_emb = self.t_embedder(t).unsqueeze(1).expand(-1, seq_len, -1)
    elif t.ndim == 2:
        # Self-Flow mode: each token has its own timestep (Dual-Timestep Scheduling)
        t_flat = t.reshape(-1)
        t_emb_flat = self.t_embedder(t_flat)
        t_emb = t_emb_flat.reshape(batch_size, seq_len, -1)
    else:
        raise ValueError(f"Timesteps must be 1D or 2D, got shape {t.shape}")

    # Class embedding (broadcast to per-token)
    y_emb = self.y_embedder(y, self.training).unsqueeze(1).expand(-1, seq_len, -1)

    # Combine embeddings
    c = t_emb + y_emb

    # Apply per-token blocks
    for i, block in enumerate(self.blocks):
        x = block(x, c)
        if (i + 1) == return_features:
            zs = self.projector(x)   # MLP projection head h_θ for L_rep
        elif (i + 1) == return_raw_features:
            zs = x

    # Apply per-token final layer
    x = self.final_layer(x, c)

    if return_features or return_raw_features:
        return x, zs
    else:
        return x

三个对照点:

(b) per-token adaLN 块 — 标准 DiT block 的 adaLN 是 per-sequence (一个 batch 一个 \(c\)),Self-Flow 改成 per-token (每个 token 都有自己的 \((\gamma, \beta)\))。

repo/src/model.py:L171-L189 — PerTokenDiTBlock

class PerTokenDiTBlock(DiTBlock):
    """DiT block that handles per-token conditioning (N, T, D) instead of (N, D)."""

    def forward(self, x, c):
        """
        Args:
            x: (N, T, D) tokens
            c: (N, T, D) per-token conditioning      # ← key change
        """
        batch_size, seq_len, hidden_dim = c.shape
        c_flat = c.reshape(-1, hidden_dim)
        modulation_flat = self.adaLN_modulation(c_flat)
        modulation = modulation_flat.reshape(batch_size, seq_len, -1)

        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
            modulation.chunk(6, dim=-1)
        )
        # Both shift/scale/gate are now per-token tensors (N, T, D)
        x = x + gate_msa * self.attn(modulate_per_token(self.norm1(x), shift_msa, scale_msa))
        x = x + gate_mlp * self.mlp(modulate_per_token(self.norm2(x), shift_mlp, scale_mlp))
        return x

对比标准 DiTBlock (model.py:L139-L168): 那里 shift/scale/gate 是 (N, D),被 unsqueeze(1) 广播到所有 token。Self-Flow 这版直接保留 (N, T, D) 维度,每个 token 用自己的 (γ, β)。这是适配 per-token 时步的必要架构改动 — 但它本质上无损兼容 vanilla flow matching(只要 \(\tau\) 在所有 token 上相同就行)。

(c) 训练循环 (didactic) — 论文 Algorithm 1 等价,repo 未发布训练代码。

Didactic — not Self-Flow official code; reconstructed from paper §3

def self_flow_train_step(student, teacher_ema, x0, y, gamma=0.5, mask_ratio=0.25,
                          repa_layer_student=8, repa_layer_teacher=21):
    B, N, C = x0.shape

    # 1) Sample two timesteps t, s from the flow matching schedule
    t = sample_timestep(B)        # (B,)
    s = sample_timestep(B)        # (B,)

    # 2) Sample mask M (per-token), masking_ratio ≤ 0.5
    u = torch.rand(B, N, device=x0.device)
    M = u < mask_ratio            # (B, N) bool

    # 3) Construct per-token tau (Eq 4): tau[i] = s if i ∈ M else t
    tau = torch.where(M, s.unsqueeze(-1), t.unsqueeze(-1))   # (B, N)

    # 4) Noise input per-token (Eq 5)
    x1 = torch.randn_like(x0)
    x_tau = (1 - tau).unsqueeze(-1) * x0 + tau.unsqueeze(-1) * x1

    # Teacher input: same min(t,s) for ALL tokens — uniformly cleaner
    tau_min = torch.minimum(t, s)                            # (B,)
    x_tau_min = (1 - tau_min).view(B, 1, 1) * x0 + tau_min.view(B, 1, 1) * x1
    tau_min_per_token = tau_min.unsqueeze(-1).expand(-1, N)

    # 5) Generation loss (Eq 2) — student predicts velocity field
    v_pred, h_student = student(x_tau, tau, y, return_features=repa_layer_student)
    L_gen = F.mse_loss(v_pred, x1 - x0)

    # 6) Representation loss (Eq 6) — student feature aligns to EMA teacher
    with torch.no_grad():
        _, h_teacher = teacher_ema(x_tau_min, tau_min_per_token, y,
                                   return_raw_features=repa_layer_teacher)
    L_rep = -F.cosine_similarity(h_student, h_teacher.detach(), dim=-1).mean()

    # 7) Total loss
    loss = L_gen + gamma * L_rep
    loss.backward()
    optimizer.step()

    # 8) Update teacher via EMA (~0.999)
    update_ema(teacher_ema, student, decay=0.999)
    return loss

关键 invariant:teacher 永远看的是 student 同一张图但整体更干净的版本。teacher 的 feature 自然更"语义化",student 用更少信息匹配这些 feature 就被迫学全局表征。

3. 结论 (Key Findings)

ImageNet 256×256 (类条件 SiT-XL):

Table 1: ImageNet results
Tab. 1 — Self-Flow 在不用外部表征的设置下,4M step 拿 FID 5.70 / sFID 4.97 — 超过 REPA (DINOv2,4M step,FID 5.89)。这是 DINOv2 在 ImageNet 上训练的天然主场,REPA 还有 IS 优势 (157.66 vs 151.40),但 FID/Recall 上 Self-Flow 都更好。配合 RAE (Representation Autoencoder) 后 FID 进一步降到 2.95 (vs RAE baseline 3.24)。

文本到图像、视频、音频 — 全模态 Pareto frontier:

Quantitative across modalities
Fig. 5 — 4 个模态全部 SOTA。T2I FID: Self-Flow 3.61 vs REPA 3.92 vs SRA 3.70 vs Vanilla FM 4.08。T2I FD-DINOv2: Self-Flow 167.98 — 超过 REPA 的 173.35,尽管 REPA 直接拿 DINOv2 当 teacher 训。Video FVD: Self-Flow 47.81 vs REPA-DINOv2 49.59 vs REPA-V-JEPA2 53.55 (V-JEPA2 比 vanilla 还差)。Audio FAD: Self-Flow 145.6 — REPA-MERT 几乎没收益。

scale 行为 — REPA 失效,Self-Flow 不失效:

Scaling behavior
Fig. 6 — 模型从 290M (depth=8) → 420M (14) → 625M (21) → 1B (28),REPA 跟 Self-Flow 的差距随 scale变宽。Self-Flow 的 625M 直接超过 REPA 的 1B(实线红 vs 虚线红)。FLOPs 维度也一样:REPA 在 ~$10^{20.5}$ FLOPs 后 plateau,Self-Flow 仍在涨。

多模态训练 — 一个模型同时学三种模态:

Multi-modal training
Fig. 8 — (a) 一个模型同时训 T2I + T2V + T2A,通过 modality 损失权重在三模态间权衡。Self-Flow (阴影外区域) 在所有权重配置下都比 vanilla FM (内圈) 全面胜出。三模态 simultaneously 改善。(b) 视频-动作联合任务 (RT-1 机器人 + SIMPLER 仿真),Self-Flow 在复杂多对象任务 (Move Near, Open and Place) 上显著领先。
SIMPLER rollouts
Fig. 7 — SIMPLER 机器人仿真:Vanilla FM 在 "Move Near" 和 "Place Apple in Closed Top Drawer" 这种需要多步骤、多对象推理的任务上失败 (左),Self-Flow 成功完成 (右)。这是 self-supervised 学到的"全局表征"在 embodied AI 上的具体收益。

关键数字汇总:

4. 实现细节 (Implementation Notes)

5. 批判性总结 (Critical Assessment)

5.1 优点

5.2 不足 / 疑点

5.3 适用 vs 不适用

5.4 进一步阅读

讨论 / Comments

评论托管在本仓库的 GitHub Discussions, 需 GitHub 账号。