E²PO: Embedding-perturbed Exploration Preference Optimization for Flow Models

Paper: arXiv:2605.15803v1 PDF: /Users/bytedance/ai-skills/papers/2605.15803v1.pdf Code search: 代码搜索未找到开源实现(arXiv 源码未给 GitHub 链接;GitHub REST / Web 查询标题、E2POEmbedding-perturbedAMAP E2PO Flow 未找到对应公开仓库)

1. Motivation (研究动机)

现有 RL-based text-to-image / flow-model alignment 方法把 GRPO 一类 group-relative 优化迁移到视觉生成:通过 SDE/随机采样或 DiffusionNFT 式 forward-process contrastive objective,让模型直接朝 reward model 偏好的图像区域更新。问题是这些方法的学习信号高度依赖 group 内 reward 的相对差异;训练推进后,样本逐渐收敛到相似模式,intra-group discriminative variance 快速衰减,导致 reward 标准差接近 ,优势/optimality 概率失去可分性。

作者要解决的具体目标不是单纯“提高 reward”,而是在 flow matching / diffusion-like 生成模型的 RL 微调中,持续维持 group 内可判别差异,使优化不会因为同质化样本而停滞,也不通过扩大 latent group size 或随机初始噪声来高成本、低稳定性地补探索。

Figure 1 解读:上半部分展示 baseline rollout 逐渐出现 Vanishing Discriminative Signal:多个样本 reward 接近,group 内方差消失;E²PO 则通过 semantic manifold 上的 embedding-level perturbation 保持不同轨迹。下半部分把这种机制和最终效果连在一起:左侧 reward alignment 更高,中间显式注入 semantic perturbation,右侧探索空间从狭窄局部模式扩展到更宽的语义区域。

这个问题值得研究,因为一旦 group variance 维持不住,视觉 RL 的 reward 提升会同时遭遇两个失败模式:第一,optimization stagnation,即相对优势近似消失;第二,reward hacking,即模型锁定少数 proxy-high-reward 但视觉/语义不可靠的模式。E²PO 的价值在于把探索源从无结构 latent noise 转移到文本 embedding semantic manifold,用更少 seed 和更多 semantic variants 组合出稳定的 discriminative signal。

2. Idea (核心思想)

核心 insight:对 flow model 的 RL 微调而言,真正缺的不是更多随机噪声,而是能在保持 prompt anchor 的同时产生“语义上可控差异”的 exploration source。E²PO 因此不盲目增加 noise seeds,而是在 frozen text encoder 的 embedding space 中学习 个 perturbation,让同一 prompt 形成多个语义近邻条件,从而持续产生 reward 可区分的 trajectories。关键创新可以压缩为三点:Embedding-Perturbed Mechanism 负责学习互相分散但仍接近原 prompt 的 semantic variants;

Noise-Aware Sampling Schedule 只在 early high-noise phase 强注入变体,later phase 回到 original anchor 以保真;Reference-Anchored Batching Strategy 把原始条件放回 batch,并用 original prompt 计算 reward / policy update,把“探索条件”和“优化锚点”解耦。与 Flow-GRPO / DiffusionNFT 的根本差异在于:这些方法主要依赖 latent group size 、stochastic noise 或 forward-process contrastive objective 产生相对信号;

E²PO 在同等或更小 latent budget 下引入 semantic variants ,把 group 从 改成 ,其中 。因此它不是简单“采更多样本”,而是改变探索维度。

3. Method (方法)

3.1 总体框架

E²PO 建立在 Rectified Flow / Flow Matching 与 DiffusionNFT 的 supervised RL 视角上。Flow Matching 将真实样本 与高斯噪声 做线性插值: 条件 velocity model 以文本/类别条件 回归目标速度 DiffusionNFT 把 online generated group 的 raw reward 归一化为 optimality probability: \begin{split} r(\boldsymbol{x}_0,\boldsymbol{c}) :=& \frac{1}{2}+\frac{1}{2}\operatorname{clip}\bigg[\frac{1}{Z_{\boldsymbol{c}}}\Big(r^{\mathrm{raw}}(\boldsymbol{x}_0,\boldsymbol{c}) \\ &-\mathbb{E}_{\pi^{\mathrm{old}}(\cdot|\boldsymbol{c})}r^{\mathrm{raw}}(\boldsymbol{x}_0,\boldsymbol{c})\Big),-1,1\bigg]. \end{split} 如果 group reward 标准差 ,这个归一化信号会变得不可判别或数值不稳定。E²PO 的整体框架就是先在 condition embedding 中构造稳定 semantic variants,再用这些 variants 采样,但把 policy update 锚回 unperturbed condition。

Figure 3 解读:图中有三个关键环节:(a) 只对内容 token 的 embedding 注入可学习扰动,形成多个 semantic variants;(b) 用 noise-aware schedule 控制 early stage 使用 perturbed condition、late stage 回到 original condition;(c) 在 batch 中保留 ,并相对 original prompt 计算 reward / gradient,避免 semantic drift。

3.2 Intra-group Discriminative Variance 分析

Figure 2a–2c 解读:三张子图分别跟踪 zero standard deviation group ratio、GenEval reward standard deviation(log scale)和 PickScore reward standard deviation(log scale)。baseline 的 reward std 随训练显著下降;E²PO 保持更稳定的 variance level,说明 embedding perturbation 比单纯扩大 更能维持 group 内可判别差异。

作者把探索集合重写为: 其中 。这意味着在同样总样本预算下,E²PO 将一部分预算从 latent noise seeds 转换成 semantic variants。直觉上,latent noise 是无结构随机性,容易仍落入相同 high-density mode;embedding perturbation 沿语义流形移动,更容易生成 reward 和结构均有差异但仍与 prompt 一致的样本。

3.3 Embedding-Perturbed Mechanism

给定 prompt ,先得到 token sequence 和 static embeddings 。方法定义有效内容 token index set ,排除 padding、[SOS][EOS] 等边界/非内容 token,令

正文用 个 semantic variants 表示 embedding 扰动集合;在 Algorithm 1 的 batching 语境下,其中一个条件保留为 unperturbed anchor,因此实际优化的是 个 perturbation tensors。论文公式写作: 初始化为: 扰动只作用于有效内容 token: 随后 frozen text encoder 输出 contextualized features;用 pooled embedding 表示第 个变体和 anchor: 多样性项最小化 variants 之间的 mean pairwise cosine similarity: anchor 约束把每个 variant 与原 prompt 的 cosine similarity 拉向目标 ,并使用 -insensitive loss: 最终 embedding objective: 这里的关键不是让 prompt 任意漂移,而是在 diversity 与 semantic consistency 之间做 constrained exploration。 防止 个 variants 退化为重复条件; 防止 variants 离原 prompt 太远,从而降低 reward hacking 与语义漂移风险。

def optimize_embedding_variants(text_encoder, token_embeddings, content_indices,
                                K, steps=300, sigma_init=1e-4,
                                lambda_div=50.0, mu=0.80, eps=0.01, lr=1e-3):
    E = token_embeddings.detach()
    L, d = len(content_indices), E.shape[-1]
    delta = torch.randn(K - 1, L, d, device=E.device) * sigma_init
    delta = torch.nn.Parameter(delta)
    opt = torch.optim.AdamW([delta], lr=lr)
    anchor = pool(text_encoder(E[None, ...])).detach()
 
    for _ in range(steps):
        variants = []
        for k in range(K - 1):
            E_tilde = E.clone()
            E_tilde[content_indices] = E_tilde[content_indices] + delta[k]
            variants.append(pool(text_encoder(E_tilde[None, ...])))
        e = torch.cat(variants, dim=0)
        sim = F.cosine_similarity(e[:, None, :], e[None, :, :], dim=-1)
        div_loss = (sim.sum() - sim.diag().sum()) / ((K - 1) * (K - 2))
        anchor_sim = F.cosine_similarity(e, anchor.expand_as(e), dim=-1)
        anc_loss = torch.clamp((anchor_sim - mu).abs() - eps, min=0).sum()
        loss = lambda_div * div_loss + anc_loss
        opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            delta.clamp_(min=-0.05, max=0.05)
    return delta.detach()

3.4 Noise-Aware Sampling Schedule

作者认为 diffusion / flow sampling 对 condition 的敏感性随 timestep 变化:早期 high-noise phase 更适合决定结构/布局,后期 low-noise phase 更适合保真细节。因此 E²PO 不在全程使用 perturbed condition,而用归一化 sampling time (从 衰减)动态插值: 其中: 高时, 接近 ,trajectory 被 semantic variant 推向不同结构;当 低时, 接近 ,condition 回到 original anchor,避免 late-stage 语义漂移或视觉 artifact。

def noise_aware_condition(C_orig, C_opt, sigma_t, rho=0.4):
    gamma = torch.clamp((sigma_t - (1.0 - rho)) / rho, min=0.0, max=1.0)
    while gamma.ndim < C_orig.ndim:
        gamma = gamma[..., None]
    return gamma * C_opt + (1.0 - gamma) * C_orig

3.5 Reference-Anchored Batching 与 policy update

Reference-Anchored Batching 构造: 采样时使用 perturbed guidance ,但 reward 和 policy update 都以 为条件。这样高 reward trajectory 会把 original prompt 的分布推向高质量结构区域,低 reward trajectory 被负向项惩罚;同时,因为 update 不绑定到漂移后的 prompt,模型不会把“扰动语义”误当成用户原始意图。

DiffusionNFT 的 positive / negative implicit velocity 为: 对应 weighted contrastive loss: \begin{split} \mathcal{L}_{\text{NFT}}(\theta)=\mathbb{E}_{t,c,x_0,x_1}\big[&r\|v_\theta^+(x_t,t,c)-v\|^2\\ &+(1-r)\|v_\theta^-(x_t,t,c)-v\|^2\big]. \end{split} 论文还给出 self-normalized -regression 版本:

def e2po_train_iteration(ref_policy, old_policy, policy, text_encoder, prompts, reward_fn,
                         K, beta, rho, optimizer):
    buffer = []
    for prompt in prompts:
        E, content_indices = embed_prompt(prompt)
        delta = optimize_embedding_variants(text_encoder, E, content_indices, K=K)
        C_orig = text_encoder(E[None, ...]).detach()
        C_opts = build_optimized_conditions(text_encoder, E, content_indices, delta)
        C_batch = [C_orig] + C_opts[:K - 1]
 
        samples = []
        for C_opt in C_batch:
            x = torch.randn_like_latent(prompt)
            for t, sigma_t in sampler_timesteps():
                C_t = C_orig if C_opt is C_orig else noise_aware_condition(C_orig, C_opt, sigma_t, rho)
                x = flow_sampler_step(old_policy, x, t, C_t)
            samples.append(decode_latent(x))
 
        raw_reward = reward_fn(samples, prompt)
        r = normalize_to_optimality_probability(raw_reward)
        buffer.extend((prompt, image, prob) for image, prob in zip(samples, r))
 
    for prompt, x0, r in buffer:
        t = torch.rand(())
        eps = torch.randn_like(x0)
        xt = alpha(t) * x0 + sigma(t) * eps
        v_target = target_velocity(x0, eps, t)
        C_orig = encode_original_condition(text_encoder, prompt)
        v_old = old_policy(xt, t, C_orig).detach()
        v_new = policy(xt, t, C_orig)
        v_pos = (1.0 - beta) * v_old + beta * v_new
        v_neg = (1.0 + beta) * v_old - beta * v_new
        loss = r * F.mse_loss(v_pos, v_target) + (1.0 - r) * F.mse_loss(v_neg, v_target)
        optimizer.zero_grad(); loss.backward(); optimizer.step()

代码搜索未找到开源实现;以上 pseudocode 是按论文 Algorithm 1、Eq. (embedding perturbation / noise-aware / NFT) 重构,用于理解方法流程,不是 released code mapping。

Figure 4 解读:定性对比覆盖 spatial reasoning、attribute binding、text rendering 和 counting 等失败模式。作者指出 baseline 会把 cow 放进 TV、漏掉 black sink、生成乱码文字或在计数任务中重复/过生成;E²PO 更好地保持指令约束和结构多样性,说明 semantic exploration 不只是提高 reward,而是在减少 proxy reward hacking。

4. Experimental Setup (实验设置)

4.1 数据集与评估协议

  • GenEval:用于 compositional reasoning / text-image alignment;论文将 GenEval 既作为 reward model / direct task metric,也用其评估高探索设置与高效设置。论文未详细说明本实验使用的 prompt 数量。
  • PickScore validation set:用于 Human Preference Alignment;论文从 PickScore validation set 随机采样 prompts 做 user study,但未给出采样数量。
  • DivGenBench:用于 diversity benchmark,报告 IDS、ASC、SDI、PVS;论文未详细说明该 benchmark 在本文中的样本规模。
  • Unseen rewards / transfer metrics:Aesthetic Score、ImageReward、HPSv2.1,用于检查是否只 overfit PickScore proxy reward。

4.2 Baselines 与模型配置

Backbone 是 Stable Diffusion 3.5 Medium (SD3.5-M),所有实验在 8× H20 GPUs 上运行。对比方法包括 SD3.5-M base、Flow-GRPO、DiffusionNFT;user study 还包含 DanceGRPO。作者说明 baseline 使用 official open-source repositories,并在 reward setting、training steps 和 sampling budget 上做标准化;同时 SD3.5-M baseline 以 CFG-free 方式评估,以隔离 policy optimization 的贡献。

E²PO 三组主要配置:

  • High-Exploration Compositional Generation:GenEval,,训练 300 steps
  • Efficient Compositional Generation:GenEval,,训练 850 steps
  • Human Preference Alignment:PickScore,,训练 750 steps

4.3 训练超参数

  • General Settings:Random Seed 42;Optimizer AdamW;Mixed Precision fp16;Distributed Mode DDP;LoRA Rank ;LoRA Alpha ;Resolution ;Use EMA True
  • Embedding-perturbed Exploration:Optimize Steps ;Embedding LR ;Init. Std Dev ;Max Norm Clip ;Diversity Weight ;Anchor Target ;Anchor Margin ;Perturbation Range
  • Policy Optimization:Policy LR ;Weight Decay ;KL Penalty ;Grad Clip Norm 1.0;Grad Accumulation Per Epoch 1;Train Sample Steps 10;Train Guidance 1.0;Inference Steps 40

4.4 指标含义

GenEval 与 PickScore 是 in-domain reward / direct task metrics;Aesthetic、ImageReward、HPSv2.1 用来衡量对未见 reward 的 transferability。DivGenBench 的 IDS、ASC、SDI、PVS 用来衡量 diversity,其中表格中 IDS 越低越好,ASC/SDI/PVS 越高越好。论文没有在正文中展开每个 DivGenBench 子指标的完整定义。

5. Experimental Results (实验结果)

5.1 GenEval 主结果

High-Exploration () on GenEval

MethodGenEval ↑IDS ↓ASC ↑SDI ↑PVS ↑
SD3.5-M---0.2630.0440.1430.4580.392
+Flow-GRPO40.7240.0510.1260.4620.321
+Flow-GRPO240.7760.0640.1230.4220.318
+DiffusionNFT40.8680.0770.1260.4580.260
+DiffusionNFT240.9220.0540.1180.4180.259
+DiffusionNFT480.9170.0510.1090.4630.196
+E²PO (Ours)40.9320.0480.1270.4670.322

Efficient Generation () on GenEval

MethodGenEval ↑IDS ↓ASC ↑SDI ↑PVS ↑
SD3.5-M---0.2630.0440.1430.4580.392
+Flow-GRPO20.6780.0790.1180.4520.195
+Flow-GRPO80.7690.0810.1250.4470.190
+DiffusionNFT20.8220.0760.0420.4500.194
+DiffusionNFT80.9150.0930.1280.4170.188
+E²PO (Ours)20.9170.0750.1360.4530.195

结论:E²PO 在 下用比 DiffusionNFT 更小的 latent group size 达到最高 GenEval 0.932,同时保持最优 IDS/ASC/SDI/PVS;在 的高效设置中,GenEval 0.917 接近 DiffusionNFT 0.915,但 diversity 指标更稳。

5.2 PickScore 主结果

MethodPickScore ↑Aesthetic ↑ImgRwd ↑HPSv2.1 ↑IDS ↓ASC ↑SDI ↑PVS ↑
SD3.5-M---19.935.600-0.500.2030.0440.1430.4580.392
+Flow-GRPO822.606.2231.270.3150.2200.1370.4110.280
+Flow-GRPO2422.726.2731.300.3240.2220.1310.4100.276
+DiffusionNFT823.316.5351.250.3230.2170.1370.4090.282
+DiffusionNFT2423.346.5141.270.3240.2390.1330.3900.298
+E²PO (Ours)823.386.5381.290.3250.1670.1380.4120.301

PickScore 结果说明 E²PO 不只优化 in-domain reward:它在 PickScore 23.38、Aesthetic 6.538、HPSv2.1 0.325 和所有 DivGenBench diversity 指标上为最好;ImageReward 1.29 略低于 Flow-GRPO 1.30,但整体 reward-transfer-diversity trade-off 最强。尤其 IDS 从 DiffusionNFT 0.239 降到 0.167,支持作者关于缓解 mode collapse / reward hacking 的主张。

5.3 Ablation

Figure 5 解读:固定总预算 时,极端配置 都不充分;高预算 下, 显著优于只依赖一个维度的探索。该图支持“noise-level diversity + semantic-level diversity 需要平衡”的结论。

Figure 6 解读:第 150 training step 的样本显示,静态全程使用 optimized variant 容易造成 semantic drift 或 artifact;Noise-Aware Schedule 只在 high-noise phase 使用 perturbation,后期回到 ,因此能同时保持结构多样与视觉保真。

Figure 7 解读:Reference anchor embeddings 带来更稳定的收敛与更高的 asymptotic performance。没有 约束时,优化更容易偏离原 prompt 的语义邻域;加入 anchor 后,exploration 和 optimization 被分离,reward 提升更稳。

Figure 9 解读:只扰动 primary encoder(CLIP-L)比同时扰动 CLIP-L + OpenCLIP-G 更稳定。双 encoder 扰动扩大了优化空间,但也更难限制 semantic drift,尤其在 PickScore optimization 中导致性能退化;因此论文选择 CLIP-L-only 作为默认扰动范围。

ablation 的精确结果:

GenEval GenEval PickScore
300.9240.90223.14
500.9320.91623.38
1000.9300.91722.86

解释:dense candidate set(GenEval )和 PickScore 偏好适中 ;sparse candidate set(GenEval )需要更大的 repulsive coefficient 100 才能逃离局部模式。

5.4 User study 与额外定性结果

Figure 8 解读:user study 比较 SD3.5-M、DanceGRPO、Flow-GRPO、DiffusionNFT 与 E²PO,维度包括 Detail Preservation、Color Consistency、Image-Text Alignment、Overall Quality。论文报告 E²PO 在 Image-Text Alignment 上 preference rate 为 33.1%,Overall Quality 为 31.8%,均领先对比方法。

Figure 10 解读:附录的 GenEval 定性图展示了 E²PO 在 compositional prompts 下更少遗漏对象、更少空间关系错误。结合主表 0.932/0.917 的 GenEval 分数,作者用该图说明 semantic variants 能提升 compositional reasoning,而不是只提高审美 reward。

Figure 11 解读:附录的 PickScore 定性图强调视觉质量与 human preference alignment。与只追求 proxy reward 的方法相比,E²PO 在质量、多样性和 prompt adherence 上更均衡,呼应 PickScore 23.38、Aesthetic 6.538、HPSv2.1 0.325 与 IDS 0.167

5.5 局限与风险

作者没有单独列出 limitations section;从实验和 ablation 可见的限制是:E²PO 需要为每个 prompt 额外优化 embedding perturbations(默认 ),并引入 、anchor target 等超参数。论文也没有公开该方法的实现代码,因此当前笔记无法核对 released code 与 paper formula 是否存在差异;training-config 数字只能来自论文附录表,而非实际 launch config。