跳转至

DirectEdit: Step-Level Accurate Inversion for Flow-Based Image Editing

会议: ICML 2026
arXiv: 2605.02417
代码: https://desongyang.github.io/Directedit (有)
领域: 扩散模型 / 图像编辑 / Rectified Flow
关键词: 流模型反演、训练无关图像编辑、步级精确重建、特征注入、Mask 引导融合

一句话总结

DirectEdit 通过在 Rectified Flow 反演过程中记录每一步的潜变量残差 \(\Delta\mathbf{Z}_t\) 并在前向路径中提前注入,让重建路径与反演轨迹严格逐步对齐,从而在不增加任何 NFE 的前提下实现"步级精确重建",并结合 MLLM+SAM 多分支 mask 噪声融合与注意力 Value 注入,在 PIE-Bench 上以综合排名 4.0 (FLUX) / 2.43 (SD3.5) 显著优于 RF-Inversion、FireFlow、FTEdit、DNAEdit 等所有现有训练无关方法。

研究背景与动机

领域现状:Rectified Flow (RF) 已成为 SD3.5、FLUX 等大规模 T2I 模型的主流框架。基于这些预训练流模型做训练无关图像编辑通常遵循"反演→去噪重建+编辑双路径"的范式:先用 Euler 反演把干净图像映射到噪声潜变量,再在前向去噪时同时跑一条 source 重建路径和一条 target 编辑路径,并把 source 路径的特征(attention KV、潜变量等)注入 target 路径以保持原图信息。

现有痛点:Euler 反演的核心近似 \(\mathbf{Z}^{inv}_t = \mathbf{Z}^{inv}_{t+1} - (\sigma_{t+1}-\sigma_t)\,v_\theta(\mathbf{Z}^{inv}_{t+1})\)\(t+1\) 时刻的速度近似 \(t\) 时刻速度,单步误差虽小但会沿去噪步累积,导致重建路径完全偏离反演轨迹。后续工作(高阶 ODE 求解器 RFEdit、定点迭代 FTEdit、插值噪声优化 DNAEdit)只能缓解轨迹整体漂移,每一步内部的 step-level 误差仍然存在——这意味着每一步注入到编辑路径的 source 特征都是"漂移的特征",最终导致背景失真和编辑伪影;FlowEdit 类 inversion-free 方法则因为依赖随机噪声插值反而牺牲了 fidelity。作者在 Table 2 实测:标准 Stepwise Correction 的 step-level MSE 平均 0.2857、最大 11.73;FTEdit 也有 0.0881。

核心矛盾:现有方法都在"修正反演路径让其逼近重建路径"这条思路上做文章,但 RF 的 Euler 反演对单步近似误差比 DDIM 反演敏感得多,再怎么优化反演侧,前向去噪那一步用的速度 \(v_\theta(\mathbf{Z}_t)\) 和反演时记录的速度 \(v_\theta(\mathbf{Z}^{inv}_{t+1})\) 本质上就不一样,step-level 误差无法根除。

本文目标:在不增加任何额外神经网络评估次数 (NFE) 的前提下,实现 step-level 的精确重建——即让前向路径每一步的潜变量都严格等于反演路径对应步的潜变量 \(\mathbf{Z}_{t+1} = \mathbf{Z}^{inv}_{t+1}\)

切入角度:反过来想——既然修正反演路径走不通,那为什么不直接对齐前向路径?只要能让前向去噪那一步用的速度恰好等于反演那一步用的速度 \(v_\theta(\mathbf{Z}_t) = v_\theta(\mathbf{Z}^{inv}_{t+1})\),逐步对齐就自动成立。而 \(\mathbf{Z}^{inv}_{t+1}\) 在反演阶段是完全可访问的,只要在反演时把每一步的潜变量残差 \(\Delta\mathbf{Z}_t = \mathbf{Z}^{inv}_{t+1} - \mathbf{Z}^{inv}_t\) 缓存下来即可。

核心 idea:用反演时缓存的潜变量残差 \(\Delta\mathbf{Z}_t\) 在前向路径预测速度前临时加到当前 latent 上得到 \(\hat{\mathbf{Z}}_t = \mathbf{Z}_t + \Delta\mathbf{Z}_t\),再用 \(v_\theta(\hat{\mathbf{Z}}_t)\) 去更新 \(\mathbf{Z}_t\),从而无代价地强制速度对齐反演轨迹。

方法详解

整体框架

输入:源图像 \(\mathbf{I}_{src}\)、源 prompt \(\psi_{src}\)、目标 prompt \(\psi_{tar}\)\(T\) 步时间表 \(\{\sigma_0,\dots,\sigma_T\}\)。流程分两阶段:

  1. 反演阶段:VAE 编码得到 \(\mathbf{Z}_T\),按标准 Euler 反演 \(T\) 步得到反演轨迹 \(\{\mathbf{Z}^{inv}_t\}\)关键是缓存每一步的潜变量残差 \(\Delta\mathbf{Z}_t\);同时调用 MLLM 解析编辑类型 \(O \in \{\text{Local, Background, Global, Other}\}\) 与感兴趣区域坐标 \(P,Q\),结合 SAM 生成 mask \(\mathcal{M}\)
  2. 编辑阶段:从 \(\mathbf{Z}^{inv}_0\) 出发同时跑 src 和 tar 两条前向路径,每一步先用 \(\Delta\mathbf{Z}_t\) 把 latent 推到 \(\hat{\mathbf{Z}}_t\) 上预测速度(Direct Alignment),tar 路径的 self-attention 把 Value 替换为 src 路径的 Value(Attention Injection),更新后用 mask 把 src 与 tar latent 融合(Noise Blending)。最终 VAE 解码得到编辑结果。

整个过程额外开销只有缓存一份 \(\{\Delta\mathbf{Z}_t\}\)没有引入任何额外的 NFE——这是相对 FTEdit (120 NFE)、RFEdit (120 NFE) 的关键效率优势。

关键设计

  1. Direct Alignment for Accurate Inversion(步级直接对齐):

    • 功能:让前向重建路径在每一步都严格等于反演轨迹对应步的潜变量,从根源消除 step-level 重建误差。
    • 核心思路:由 Euler 更新公式可知 \(\mathbf{Z}_{t+1} = \mathbf{Z}^{inv}_{t+1}\) 等价于 \(v_\theta(\mathbf{Z}_t) = v_\theta(\mathbf{Z}^{inv}_{t+1})\),而后者中的 \(\mathbf{Z}^{inv}_{t+1}\) 反演阶段已知。因此在反演时记录潜变量残差 \(\Delta\mathbf{Z}_t = \mathbf{Z}^{inv}_{t+1} - \mathbf{Z}^{inv}_t\);前向去噪时先构造对齐潜变量 \(\hat{\mathbf{Z}}_t = \mathbf{Z}_t + \Delta\mathbf{Z}_t\) 用于速度预测,再以 \(\mathbf{Z}_{t+1} = \mathbf{Z}_t + (\sigma_{t+1}-\sigma_t)\,v_\theta(\hat{\mathbf{Z}}_t)\) 更新。注意速度预测的输入是 \(\hat{\mathbf{Z}}_t\) 但更新的是原 \(\mathbf{Z}_t\),这一巧妙错位让速度场恰好与反演时一致。
    • 设计动机:彻底跳出"修反演路径"的旧思路,转而修正前向路径的速度输入;不额外调用 \(v_\theta\) 因此 NFE 与 vanilla Euler 相同(60 NFE vs FTEdit 120 NFE),但 step-level MSE 从 0.0881 量级降到 0.0006(Table 2,降低两个数量级)。
  2. Multi-branch Mask-guided Noise Blending(多分支 mask 引导噪声融合):

    • 功能:根据编辑类型自动生成空间 mask,在 latent 空间逐步融合 src 与 tar 路径以精确保护背景。
    • 核心思路:先用 MLLM 输入 \((\mathbf{I}_{src}, \psi_{src}, \psi_{tar})\) 输出编辑类型 \(O\) 和感兴趣框 \((P, Q)\);再根据 \(O\) 走不同分支生成 mask——Local 用 SAM 分割得到对象 mask,Background 取其补集,Global(如风格迁移)整图为 1,其他(如物体添加)直接用矩形框 \(\mathcal{B}(P,Q)\)。每个去噪步结束后融合 \(\mathbf{Z}^{tar}_{t+1} \leftarrow \mathbf{Z}^{src}_{t+1} \odot (\mathbf{1}-\mathcal{M}) + \mathbf{Z}^{tar}_{t+1} \odot \mathcal{M}\),配合 DirectEdit 的精确重建,非编辑区域可做到几乎零失真。
    • 设计动机:单一矩形或单一分割掩码无法适配所有编辑类型(背景编辑取反、风格迁移要整图、物体添加 SAM 没法分割不存在的对象),多分支机制让 mask 生成自动跟编辑意图匹配;融合在 latent 空间逐步进行可避免最终一次性融合带来的边界伪影。消融显示 w/o multi-branch(全部退回矩形框)相比 full 在 PSNR 上掉 4.77 dB(27.86→32.63)。
  3. Attention Feature Injection(注意力 Value 注入):

    • 功能:在 self-attention 层把 src 路径的 Value 注入 tar 路径,保留被编辑对象的细粒度语义细节。
    • 核心思路:tar 路径每个 MM-DiT block 的自注意力输出在前 \(t_{inj}\) 步用 \(\text{Attention}(\mathbf{Q}^{tar}_t, \mathbf{K}^{tar}_t, \mathbf{V}^{src}_t)\) 替换,之后回归标准 self-attention;得到的特征 \(\hat{\mathbf{F}}^{tar}_t\) 用于刷新 tar 速度 \(\hat{v}^{tar}_t\)。架构层面 SD3.5 对所有 MM-DiT block(除最后一个)注入,FLUX 只注入 single blocks(因为这些块同时处理图像和文本特征,能更好保留 source 信息)。\(t_{inj}=3\) 在多数场景下表现最佳,更小的 \(t_{inj}\) 更贴目标 prompt,更大的 \(t_{inj}\) 更像源图。
    • 设计动机:Mask blending 解决了背景,但编辑区域内部如果完全自由生成会丢失原对象的纹理与身份信息;Value 注入既保留原对象的"内容"(V)又允许新的"关系"(Q、K 来自 tar),加上 DirectEdit 提供了未漂移的 V,这种注入才是"干净"的。

损失函数 / 训练策略

完全训练无关,无任何 loss 与反向传播。推理设置:FLUX.1-dev / SD3.5-medium 作骨干,30 步去噪,CFG=1(反演)/CFG=2(编辑),\(t_{inj}=3\)

实验关键数据

主实验

PIE-Bench (700 张图、9 类编辑) 上的训练无关编辑对比(综合排名越小越好):

Backbone 方法 Structure↓ PSNR↑ LPIPS↓ MSE↓ CLIP-Whole↑ Avg. Rank↓
FLUX RF-Inversion 41.17 20.86 187.01 120.12 25.08 13.57
FLUX RFEdit 25.15 24.33 121.59 56.98 25.57 9.14
FLUX FireFlow 27.40 23.11 128.46 70.75 26.13 9.43
FLUX FlowEdit 27.83 21.96 112.15 94.94 25.26 10.57
FLUX DNAEdit 16.81 25.20 86.68 48.35 24.81 7.71
FLUX DirectEdit 17.94 32.63 35.45 25.05 25.39 4.00
SD3.5 FTEdit 21.06 23.49 90.25 61.78 25.21 9.29
SD3.5 FlowEdit 23.13 23.29 92.81 69.09 26.71 7.29
SD3.5 DNAEdit 11.03 27.71 60.51 26.28 25.20 5.14
SD3.5 DirectEdit 14.65 31.82 31.36 21.64 25.64 2.43

重建误差对比(FLUX,60 NFE vs 120 NFE,关键指标 Step-Level MSE):

方法 NFE↓ PSNR↑ Step-Level MSE Avg↓ Step-Level MSE Max↓
VAE (理论下限) - 34.38 - -
Vanilla Euler 60 14.59 1177.73 39511.72
Stepwise Correction 60 34.38 0.2857 11.73
FTEdit 120 34.38 0.0881 14.82
RFEdit 120 21.92 231.72 20156.25
DirectEdit 60 34.38 0.0006 0.0757

DirectEdit 用一半 NFE 把 step-level MSE 平均/最大同时降低 2~4 个数量级,触及 VAE 重建下限。

消融实验 (FLUX, PIE-Bench)

配置 Struct.↓ PSNR↑ LPIPS↓ MSE↓ CLIP-Whole↑
Vanilla 75.95 16.81 276.29 332.65 23.57
w/o alignment (回退 Stepwise Correction) 29.22 31.12 53.16 48.17 25.24
w/o attention 23.75 31.93 39.60 33.97 25.60
w/o mask 21.93 24.70 102.92 56.76 25.89
w/o multi-branch (回退矩形框) 19.15 27.86 60.92 38.94 25.71
DirectEdit (full) 17.94 32.63 35.45 25.05 25.39

关键发现

  • Direct Alignment 是绝对核心:去掉后所有指标全面崩塌(PSNR 从 32.63→31.12,Structure 距离从 17.94→29.22),印证"step-level 重建误差是注入特征漂移的根因"这一论断。
  • Mask blending 主导背景保真度:w/o mask 时 PSNR 直接掉 7.93 dB(32.63→24.70);多分支 mask 进一步贡献约 4.77 dB。
  • Attention injection 在保细节与跟 prompt 之间做 trade-off:去掉后 CLIP 反而略升(25.60 vs 25.39),但视觉上会丢原对象的细粒度纹理,\(t_{inj}=3\) 是经验甜点。
  • 效率优势显著:60 NFE 跑出比 120 NFE 的 FTEdit 更好的 step-level MSE(0.0006 vs 0.0881),意味着既快又准。

亮点与洞察

  • "反着想"的反演哲学:所有前人都在修反演路径让它逼近重建路径,作者反过来——把反演路径当 pivot,让前向路径去对齐它。一旦换了视角,解法竟然简单到只需缓存一份残差并加回去,这是工程师审美的极致体现。
  • 零 NFE 代价实现两个数量级精度提升:纯靠潜变量层面的代数恒等(\(v_\theta(\mathbf{Z}_t + \Delta\mathbf{Z}_t) = v_\theta(\mathbf{Z}^{inv}_{t+1})\))就达成 step-level 严格对齐,没有额外神经网络前向,没有定点迭代,没有高阶求解器——这种"几乎免费"的改进特别可遇不可求。
  • MLLM+SAM 形成的"语义化 mask 路由器":把 mask 生成从"二选一"(要么 SAM 要么矩形框)升级为按编辑类型多分支路由,这个模式可以直接迁移到所有需要空间约束的训练无关编辑(视频编辑、3D 编辑、布局到图像等)。
  • 可复用 trick:缓存反演侧中间量并在前向侧"补偿"的范式,可推广到任何用 Euler/RK 离散求解 ODE 的反演任务,例如音频/视频 RF 编辑、3D 高斯编辑等。

局限与展望

  • 作者承认:编辑质量天花板被骨干 T2I 模型的先验锁死;对尺寸变化、空间操控、复杂上下文推理这类需要"重新组织内容"的编辑仍然吃力。
  • 隐含局限:MLLM 解析编辑类型与坐标的准确性直接决定 mask 质量,论文未量化 MLLM 误判对最终编辑的影响;当 MLLM 把 Local 误判为 Background 时整个融合方向会反掉。
  • 隐含局限:需要缓存 \(\{\Delta\mathbf{Z}_t\}\)\(T\) 份 latent),对长 schedule 和高分辨率的 FLUX 在显存上有压力,论文未给峰值显存数据。
  • 改进思路:(1) 把 Direct Alignment 与一阶高阶 solver 结合,进一步用更少步数达到同样精度;(2) 把 mask 路由从离散四分支扩展为可学习的连续 mask 生成器;(3) 与 inversion-free 思路(FlowEdit)做混合——前几步用 inversion-free 探索大幅编辑,后几步用 DirectEdit 精修细节。

相关工作与启发

  • vs Stepwise Correction (Direct Inversion, 2023):两者都意识到反演与重建路径要严格对齐,但前者通过"强制 \(\mathbf{Z}_t = \mathbf{Z}^{inv}_t\)"在重建后修正轨迹,每步内部速度仍然不一致;DirectEdit 通过"事先用残差对齐速度输入"在重建前就保证速度一致,更彻底且 step-level MSE 量级低两个数量级(0.2857→0.0006)。
  • vs FTEdit / DNAEdit:FTEdit 用定点迭代 + 重建轨迹修正、DNAEdit 用插值速度估计,本质上都是"减小但消不掉"step-level 误差,且都把 NFE 翻倍到 120;DirectEdit 在 60 NFE 下直接做到 step-level 误差量级 1e-4,且实测 PSNR 和综合 rank 全面领先。
  • vs FlowEdit (inversion-free):FlowEdit 通过随机噪声插值绕过反演,简单粗暴但 fidelity 差(FLUX 上 PSNR 仅 21.96);DirectEdit 走的是"把反演做到极致"的另一极端,证明了精确反演路线在 RF 上依旧是更优选择。
  • vs RF-Inversion:RF-Inversion 用 LQR 控制理论构造辅助速度场,理论漂亮但实测 PSNR 仅 20.86;DirectEdit 用最朴素的"残差缓存+对齐"达到 32.63,再次说明工程上简单方案+正确视角往往胜过复杂理论。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ "对齐前向而非反演"的视角切换 + 零 NFE 残差缓存技巧,简单到看完想拍大腿,属于真正改变思路的小创新。
  • 实验充分度: ⭐⭐⭐⭐ PIE-Bench 全指标 + 双骨干 (FLUX/SD3.5) + step-level MSE 单独验证 + 完整消融,覆盖训练无关编辑 SOTA,但缺少长 schedule 和高分辨率显存数据。
  • 写作质量: ⭐⭐⭐⭐ 推导清晰、对比图 (Fig.2) 和 algorithm 1 帮助理解,但部分关键证明放 Appendix A 未展开。
  • 价值: ⭐⭐⭐⭐⭐ 在 RF-based 训练无关编辑上同时把"效率(60 vs 120 NFE)"和"精度(step MSE 降 2 个数量级)"刷到新水位,可直接被下游研究复用作为反演基础组件。