西湖大学破解Rectified Flow难题量子位

近年来，大规模生成模型在视觉内容生产领域的能力持续进阶，人们对 AI 的应用方式也逐步从 “纯生成” 向 “可控生成” 与 “智能编辑” 转变。

在创意设计、影视制作、社交媒体、数字人建模等实际应用场景中，用户往往期望模型不仅能从文本生成图像，更能对现有内容开展精细化修改，比如保持主体不变的局部编辑、风格调整、跨帧一致的视频修改以及现实照片的结构化增强等。

这类任务的共同前提是：模型必须具备稳定且可靠的反演（inversion）能力，能够将输入图像精准映射至自身的潜在表示空间。

在新一代生成架构中，Rectified Flow（RF）模型凭借高效的 ODE 推理、平滑的生成轨迹以及良好的训练稳定性，正成为替代扩散模型的重要方向。越来越多的主流模型开始采用流式生成结构，希望借助其速度优势、可控性与结构化潜力，构建起统一的生成系统。

但相较于其在正向生成上的出色表现，RF 模型在反演阶段仍面临着根本性的数值不稳定问题：

逆向路径对误差高度敏感，极易偏离前向轨迹，不同次反演得到的 latent 表示存在显著差异，噪声还会在逆向传播过程中不断被放大。

这些问题让 RF 模型难以在实际编辑任务中稳定重建输入图像，极大限制了其在真实场景中的实际可用性。

更关键的是，现有诸多改善反演效果的方法，往往需要依赖额外训练、增设模型组件或修改原有结构，这不仅抬高了工程实现成本，也降低了方法在不同架构中的泛化能力。

在大模型时代，这类高度依赖训练的优化方式，难以满足快速集成与跨模型复用的实际需求。相比之下，一种能直接作用于反演流程、无需任何训练就能显著提升稳定性，且可立即接入现有模型的即插即用（plug-and-play）机制，显然更契合实际部署的要求。

基于这一研究动机，西湖大学张驰研究团队提出了 PMI（Prox-Mean-Inversion）方法。PMI 的核心设计思路，是利用 RF 模型内部固有的平均流结构，在逆向 ODE 的每一步中加入一个轻量级的 proximal 修正步骤，从而在不改变模型原有结构、不引入额外参数的前提下，有效抑制逆向轨迹的偏移与发散问题。

相关研究论文已被 ICLR26 正式接收。

该论文的第一作者为西湖大学研究人员王晨如，指导老师是西湖大学 AGI 实验室负责人张驰助理教授，该实验室的核心研究方向聚焦于生成式人工智能与多模态机器学习领域。

由于 PMI 无需开展额外的模型训练，也不改变原有推理方式，因此可直接作为即插即用（plug-and-play）组件集成到各类 RF 模型中，为反演过程带来显著的稳定性提升。换言之，PMI 为 RF 模型的反演优化提供了一种真正意义上的 free lunch：无需付出额外成本、无需修改模型结构，却能在多个任务和不同模型中取得可观的性能收益。

这一特性让 PMI 对于希望快速落地 RF 编辑系统的科研与工业场景而言尤为重要，通过一个简单、通用且与训练无关的优化步骤，就能赋予模型更可靠的反演能力，进而让局部编辑、条件控制、视频一致性以及多模态交互等任务拥有更稳定的基础。

PMI：基于 Proximal Operator 的逆向轨迹稳定机制

现有 Rectified Flow（RF）模型在反演阶段面临的核心挑战，在于逆向 ODE 对微小误差高度敏感，极易沿着数值不稳定的方向偏离前向流形，最终导致轨迹发散、图像重建不一致、编辑效果不可控。为解决这一核心问题，研究团队提出了 PMI（Prox-Mean-Inversion）—— 一种针对 RF 反演稳定性的轻量化修正机制。

PMI 的核心思想是：

在每一步逆向积分完成后，引入一个基于 proximal operator 的修正项，让反演轨迹始终保持在模型内部的 “平均流” 结构（mean flow）附近，从而显著减少轨迹偏移与震荡现象。

具体而言，RF 模型的速度场在训练过程中，会形成具有代表性的平均方向，这一方向反映了模型生成路径上的核心结构特征。而逆向 ODE 在数值计算中极易偏离这一方向，因此 PMI 在每个时间步中，通过

将当前 latent 向平均流估计值进行轻量投影，使其贴近模型的高密度区域。

与传统的反演增强方法相比，PMI 具备以下显著特征：

无需额外训练：完全依托模型现有结构，不引入任何额外参数；

无需修改原路径或 ODE 结构：与模型推理流程完全解耦，保障原始模型的行为不受干扰；

即插即用（plug-and-play）：可直接嵌入任意 RF 模型的反演过程，不依赖模型的规模与细节设计；

不增加推理成本：prox 修正步骤仅为轻量级算子，不会改变模型的时间复杂度；

显著提升反演稳定性：有效减少多次反演之间的结果差异，提高轨迹一致性与编辑可控性。

通过 PMI 的优化，RF 模型能在无额外开销的前提下，获得更平滑、更鲁棒的反演轨迹，为后续的局部编辑、图像增强、风格迁移等任务，提供更稳定的基础表示。

Mimic-CFG：解决编辑中的 “身份危机”

在利用反演技术进行图像编辑时，通常会引入无分类器引导（CFG）机制。CFG 虽能增强模型对编辑指令的遵循度，却往往会大幅改动潜空间路径，导致生成的图像虽符合文本描述，却丢失了原图的身份特征（ID），或出现严重的伪影问题。

针对这一矛盾，PMI 创新性地提出了 mimic-CFG 策略，通过两大核心设计实现平衡：

方向投影：将带引导的速度场投影到历史平均方向上；

动态补偿：借助该投影机制，保留原图中本质的结构信息，同时仅在 “必要” 的维度上响应编辑指令。

这一策略的实际效果十分显著：模型能够在完成 “给猫戴上帽子” 这类大幅度编辑的同时，确保主体仍是原本的对象，背景纹理也不会出现崩坏。

简言之，mimic-CFG 解决的是图像编辑中 “编辑效果达成后，未编辑区域是否保持一致” 的关键问题。

PMI 与 mimic-CFG 的结合，让 RF 模型在反演阶段同时具备数值稳定性与语义可控性，进而显著提升可控编辑任务的可靠性。

实验结果：刷新多项指标，实现 “零失真” 重建与编辑

研究团队在图像编辑权威基准数据集 PIE-Bench 上，对 PMI 开展了全面的性能评测，并将其与 Vanilla RF-Inversion、FireFlow 以及传统的 Euler 采样器进行了多维度对比分析。

定量分析：重建精度与感知质量的双重领先

在图像重建（Reconstruction）任务中，PMI 展现出统治级的性能表现。实验数据显示，在相同采样步数（如 20 步）的前提下：

重建保真度：PMI 在 PSNR（峰值信噪比）指标上显著高于各类对比方法，意味着其能近乎无损地还原原始图像的像素级细节；

感知相似度：在 LPIPS 指标上，PMI 取得了最低的数值，证明其生成的图像在人类视觉感知层面与原图高度一致，彻底解决了传统反演中常见的 “纹理模糊” 问题；

效率优势：相较于需要大量迭代的优化类方法，PMI 能在更短的推理时间内实现收敛。

零成本适配：即插即用的通用性

值得关注的是，PMI 的实验验证并非局限于单一模型。研究结果证明，PMI 具备极强的普适性，可直接应用在 Flux.1-schnell/dev、Stable Diffusion 3 等基于 Rectified Flow 的主流模型上，无需进行任何额外的微调，也不会产生额外的显存开销。

随着 Flow 模型在生成式人工智能领域的影响力不断扩大，反演（inversion）已成为各类编辑、控制和增强任务的基础能力。但 Rectified Flow 模型在逆向推理过程中普遍存在的不稳定性，使其在实际应用中面临显著限制，而 PMI 的提出，为这一核心问题提供了高效、简洁且无需额外训练的解决方案。

PMI 通过一个基于 proximal operator 的轻量级机制，在不改变模型结构与推理流程的前提下，让逆向轨迹稳定贴近模型的平均流结构，从而显著提升反演的一致性与可复现性。这一方法不仅有效增强了图像重建质量，也让后续的局部编辑、条件控制等任务更可靠、更易用，为 Flow 模型向实际编辑系统的落地过渡奠定了重要基础。

实验中，PMI 在多个 Rectified Flow 模型及不同反演场景中，均展现出显著的性能优势，包括轨迹稳定性的提升、编辑可控性的增强，以及跨任务的一致表现。其无需训练、即插即用的核心特点，使其能直接应用于不同架构、不同规模的 Flow 系列模型，具备良好的泛化性与工程实用价值。

展望未来，Flow 模型的发展将继续朝着更高效、更可控的方向推进。研究团队认为，以 PMI 为代表的稳定反演技术，将成为推动 Flow 模型落地应用的重要组件之一。随着模型规模的持续扩大、应用场景的不断丰富，稳定的反演机制将在视频编辑、多模态融合、三维生成等复杂场景中发挥更关键的作用，也有望进一步推动 Flow 模型构建起真正意义上的统一生成与可控编辑框架。