跳转至

Root-Selecting Fixed-Point Inversion for Rectified Flows via Trajectory Straightness

会议: NeurIPS 2026
arXiv: 2606.17584
OpenReview: https://openreview.net/forum?id=WdrqbDVAXt
代码: https://github.com/seminkim/selfix
领域: 图像生成
关键词: 图像反演, Rectified Flow, 定点反演, 轨迹平直度, 图像编辑

一句话总结

针对 Rectified Flow 离散反演方程存在多个定点解导致的误差累积问题,SelFix 提出以轨迹平直度作为择根准则,通过带平直度锚点的 Halpern 迭代与解耦动量,在保证渐近精确反演的同时选出最平直的逆向轨迹。

研究背景与动机

随着文本到图像生成模型的快速发展,主流生成主干正在由传统扩散模型向以 Rectified Flow(如 FLUX 系列)为代表的流匹配架构演进。反演(Inversion)作为从真实图像逆向寻找生成该图像的初始噪声的关键技术,是无须微调的下游提示词图像编辑(training-free image editing)的基础支柱。然而,离散化误差长期以来是反演的核心瓶颈:在粗糙时间网格上简单地沿反向时间积分速度场,极易因轨迹弯曲而在回放生成过程中累积巨大的偏移,导致真实图像重建崩溃与编辑失真。

定点反演(Fixed-Point Inversion)通过固定前向离散求解器,将每一步反演形式化为求解该步求解器诱导的局部代数逆方程,在收敛意义下能够实现渐近精确反演。然而,既有定点反演方法(如 ReNoise、AIDI、GNRI)仅关注于快速收敛到方程的“任意一个”可行解。作者通过实证分析发现,神经网络学到的离散局部逆方程在实际计算中往往具有多个残差极其相近但空间位置显著不同的近似定点根。这些具有相当残差的不同解在后续前向回放中诱导的整体轨迹几何性质差异巨大,直接导致最终重建与背景保真度的显著分歧。

这一发现揭示出:定点反演不应仅仅是求解局部残差方程,本质上还必须是在多根流形中做出合理选择的“择根问题”(Root Selection)。对于 Rectified Flow 而言,其生成机制的天生优势在于速度场直线化使得离散求解器步长容忍度更高,弯曲轨迹则会急剧放大常速度局部近似带来的残差扰动。本文的核心 idea 是:将 Rectified Flow 的定点反演重新审视为基于轨迹平直度的择根优化,构建因果局部的平直度锚点,并设计带解耦动量的 Halpern 迭代,在渐近收敛至严格精确逆根的同时挑选出诱导最平直线轨迹的定点。

方法详解

整体框架

在离散 Rectified Flow 反演过程中,时间网格定义为 \(0 = t_0 < t_1 < \dots < t_N = 1\)。从真实图像潜在编码 \(x_0\) 出发,在第 \(j\) 步(对应从 \(t_{j-1}\) 到 \(t_j\) 的逆向步骤),算法需要恢复满足前向离散生成步 \(f_{t_j \to t_{j-1}}(x) = x_{j-1}\) 的状态 \(x_j\)。定义局部残差函数 \(r_j(x) := f_{t_j \to t_{j-1}}(x) - x_{j-1}\) 与对应的定点映射 \(P_j(x) := x - r_j(x)\),精确逆状态构成了非空局部闭凸解集 \(S_j^C = \operatorname{Fix}(P_j) \cap C_j\)。SelFix 在这一解集之上引入平直度指标作为根选择准则,从而构建受约束的择根极小化问题。

整体流程包含外层时间推进与内层定点求解:外层利用历史前几步已恢复的速度计算局部平直度锚点;内层则运行带平直度锚定项的 Halpern 迭代,并配合解耦动量平滑神经网络速度场输出的扰动,直至迭代收敛输出 \(x_j\)。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入真实图像潜在特征 x0 与时间网格"] --> B["外层反演时间步推进 j = 1 到 N"]
    B --> C["因果平直度锚点构建<br/>滑动窗口加权平均速度外推"]
    C --> D["平直度导向的定点择根表述<br/>投影优化求解"]
    D --> E["解耦动量 Halpern 迭代循环<br/>锚点导向 + 模型扰动平滑"]
    E --> F["输出当前步选根状态 xj 与估计速度"]
    F -->|递推直至 j=N| G["输出最终初始噪声 xN 与平直反演轨迹"]

关键设计

1. 平直度导向的定点择根表述:将逆方程求解升华为极值优化

传统定点反演只要求求解满足 \(r_j(x) = 0\) 的任意解,忽视了不同解对全局生成动态的影响。SelFix 首先建立离散轨迹平直度(Discrete Straightness, DS)度量: $\(DS(X) = \sum_{i=1}^T h_i \|v_i - \bar{v}\|^2\)$ 其中 \(v_i = (x_i - x_{i-1})/h_i\),\(\bar{v}\) 为整条轨迹的加权平均速度。实证数据表明反演轨迹的 DS 与重建 PSNR 呈现极强的负相关。然而整条轨迹的全局均值 \(\bar{v}\) 依赖于未发生的未来反演状态,在线反演过程中不可求。为此,SelFix 将全局均值替换为前序 \(W\) 个时间窗口内的加权平均速度: $\(\bar{v}^{W}_{j-1} := \frac{\sum_{i=1}^{W} h_{j-i} v_{j-i}}{\sum_{i=1}^{W} h_{j-i}}\)$ 并在时间步 \(j=1\) 时以模型初始反向速度作为冷启动回退。基于此,当前候选状态 \(x\) 的局部因果平直度选择子被定义为: $\(\phi^{\mathrm{str}}_{j,W}(x) := h_j \|v_j(x) - \bar{v}^{W}_{j-1}\|^2 = \frac{1}{h_j} \|x - a_{j,W}\|^2\)$ 这一代数恒等变形揭示出:极小化平直度选择子在几何上完全等价于向“平直度锚点” \(a_{j,W} := x_{j-1} + h_j \bar{v}^{W}_{j-1}\) 做欧氏投影: $\(x_j^\star \in \arg\min_{x \in S_j^C} \phi^{\mathrm{str}}_{j,W}(x) = \operatorname{proj}_{S_j^C}(a_{j,W})\)$ 该表述不引入任何额外的模型前向计算,且严格保留了精确逆根的约束集。

2. 锚定 Halpern 迭代:渐近精确性与择根偏置的理论统一

由于精确逆根集合 \(S_j^C\) 仅由非线性模型定点方程 \(x = P_j(x)\) 隐式给出,无法直接应用显式凸投影算子。SelFix 引入非扩张映射的 Halpern 迭代机制,以平直度锚点 \(a_{j,W}\) 作为吸引中心: $\(z^k = \alpha_k a_{j,W} + (1 - \alpha_k) P_j(z^{k-1})\)$ 其中衰减权重序列满足经典发散总和与渐减条件 \(\alpha_k \in (0, 1)\),\(\alpha_k \to 0\),\(\sum \alpha_k = \infty\)。作者具体采用有理递减调度方案: $\(\alpha_k = \alpha_1 \frac{\delta}{(k-1) + \delta}\)$ 在数学上,当迭代次数趋向无穷时,\(\alpha_k \to 0\) 保证了锚点外力完全衰减,迭代点渐近落入 \(P_j\) 的严格不动点集合;而在有限迭代步内,正权重的 \(a_{j,W}\) 持续施加向平直轨迹方向的几何引力,破除多根流形的对称性与任意性。理论证明在局部非扩张假设下,序列必强收敛于 \(S_j^C\) 中距离锚点最近的唯一极小化点,从根本上兼顾了反演残差的消除与轨迹几何的最优化。

3. 解耦动量加速:隔离恒定锚点与波动模型的更新噪声

在实际运行中,受限于计算预算,内层定点迭代步数 \(K\) 通常较小(如 4 到 10 步)。直接使用标准 Halpern 迭代易受到神经网络速度场局部非线性扰动的影响,导致收敛缓慢或震荡。如果直接对整体更新向量施加传统动量,会将随时间恒定的锚点 \(a_{j,W}\) 与波动剧烈的模型输出 \(P_j\) 混合滑动平均,削弱平直度导向信号。SelFix 提出了“解耦动量”(Decoupled Momentum)机制:仅对由模型前向计算主导的定点校正项维持动量缓冲,而锚点项直接以瞬时权重介入: $\(M^k = \mu M^{k-1} + (1 - \mu) P_j(z^{k-1})\)$ $\(z^k = \alpha_k a_{j,W} + (1 - \alpha_k) M^k\)$ 其中 \(\mu \in [0, 1)\) 为动量系数。该设计将不变的先验几何信号与含噪的观测校正信号完全解耦。论文严格证明了该解耦动量递推可以重构为满足 Cuntavepanit-Panyanak 条件的修正 Halpern 迭代形式,确保在引入动量平滑加速的同时,渐近极限依旧严格收敛于平直度极小化精确根。

实验关键数据

主实验

实验基于 FLUX.1-dev(12B 参数 Rectified Flow 主干)在权威基准 PIE-Bench 上展开,所有对比方法的函数评估次数(NFE)均在公平预算下对齐(重建实验对齐为 ~165-168 NFE,包含 15 步反演与对应前向生成;编辑实验对齐为 ~75-78 NFE)。重建实验考察图像保真度(PSNR、SSIM、LPIPS)以及轨迹平直度(离散平直度 DS 与累计选择子 \(\sum \phi^{\mathrm{str}}_{j,W}\))。编辑任务采用特征注入管线,评估背景保真度、结构距离(DINO 自相似度)与文本对齐度(CLIP Similarity)。

方案类别 方法 PSNR ↑ SSIM ↑ LPIPS ↓ 轨迹平直度 DS ↓ 累计代理 \(\sum \phi^{\mathrm{str}}_{j,W}\) ↓
朴素反向积分 ReFlow 19.54 0.6965 0.2382 0.181 0.176
定制高阶求解器 RF-Solver 21.31 0.7400 0.1997 0.172 0.162
定制高阶求解器 FireFlow 26.98 0.8408 0.1037 0.185 0.180
传统定点反演 ReNoise 22.97 0.8159 0.1155 0.177 0.157
传统定点反演 AIDI 23.00 0.7862 0.1533 0.187 0.165
传统定点反演 朴素 FPI 15.05 0.6317 0.3059 0.208 0.186
本文方法 SelFix (Ours) 28.29 0.8714 0.0742 0.170 0.151

在图像编辑任务中,SelFix 显著改善了背景保留与指令编辑之间的权衡曲线:

方法 结构距离 (DINO) ↓ 背景 PSNR ↑ 背景 SSIM ↑ 背景 LPIPS ↓ 全图 CLIP 相似度 ↑
ReFlow 0.0657 18.38 0.7204 0.2329 27.16
RF-Solver 0.0525 19.65 0.7527 0.1992 26.85
FireFlow 0.0629 18.33 0.7305 0.2218 27.25
ReNoise 0.0530 18.64 0.7449 0.1943 27.04
AIDI 0.0505 18.90 0.7581 0.1943 26.96
朴素 FPI 0.0570 18.51 0.7482 0.1830 26.09
SelFix (Ours) 0.0457 19.79 0.7748 0.1767 27.29

消融实验

消融实验深入验证了锚点引导机制与核心超参数配置:

引导机制 / 超参数变体 设定详情 重建 PSNR ↑ 轨迹平直度 DS ↓ 机制效果分析
常数调度 \(\alpha\) \(\alpha = 0.5\) (固定) 15.53 0.127 锚点偏置过强,破坏定点方程收敛,反演崩溃
常数调度 \(\alpha\) \(\alpha = 0.0068\) (固定) 27.86 0.170 锚点影响微弱,择根导向作用不足
显式梯度下降 \(\text{lr} = 0.1\) 19.56 0.157 梯度更新干扰残差归零过程
显式梯度下降 \(\text{lr} = 0.01\) 27.81 0.170 效果次优,难以自适应衰减
SelFix (完整) 衰减调度 + 解耦动量 28.29 0.170 兼顾精确反演约束与平直度最优择根
动量机制消融 \(\mu = 0.0\) (无动量) 16.77 0.204 有限步内缺少扰动平滑,震荡无法收敛
动量机制消融 耦合动量 (\(\mu = 0.5\)) 27.54 0.170 锚点被平滑稀释,表现劣于解耦动量
动量系数敏感度 解耦动量 \(\mu = 0.25\) 23.71 0.180 平滑强度不足
动量系数敏感度 解耦动量 \(\mu = 0.75\) 26.93 0.126 动量过重导致迭代更新迟滞
窗口大小 \(W\) \(W = 1\) (局部单步) 28.29 0.1702 局部恒速近似最稳健,避免历史累积误差
窗口大小 \(W\) \(W = 4\) (多步滑动) 28.17 0.1697 平直度微增但受早期轻微误差累积干扰

关键发现

  • 平直度与反演质量的高度一致性:实验中 SelFix 的轨迹平直度指标 DS(0.170)与累积代理值(0.151)达到全场最低,直接推动重建 PSNR 突破 28 dB,大幅领先次优定点方法 ReNoise/AIDI 近 5 dB,证实了流平直性能够有效遏制前向回放时的数值发散。
  • 动量解耦是有限步计算的核心使能器:完全剔除动量(\(\mu=0\))导致 PSNR 从 28.29 断崖式下跌至 16.77,表明在有限预算下模型输出的噪声抑制至关重要;而解耦设计相比传统耦合动量获得 0.75 dB 的稳定优势。
  • 局域单步外推最优性:速度滑动窗口 \(W=1\) 表现优于更长窗口(\(W=4\)),说明在快速演化的流场中,紧邻的前序速度即可提供最优的切线外推,过长的历史积分反而容易引入早期数值漂移。

亮点与洞察

  • 多根视角的范式转换:首次系统论证了离散流匹配中定点反演存在大量局部低残差伪根的普遍现象,将“求逆根”问题升华为“择优根”的几何优化问题,为逆向生成提供了崭新的理论视角。
  • 代数形式与几何物理的无缝闭环:巧妙地将离散平直度代理化简为欧氏距离二次项,使得物理上的“保持匀速直线”天然等价于数学上的“向锚点正交投影”,从而直接解锁了经典 Halpern 迭代在流匹配反演中的理论落地。
  • 无模型开销的即插即用特性:平直度锚点与解耦动量完全构建在反演中间态的历史差分之上,无需额外反向传播、无辅助变量优化,亦不增加任何额外的模型前向推理开销(Zero Extra NFE)。

局限与展望

  • 依赖主干网络的固有能力:作者承认反演与编辑质量深度绑定于底层预训练生成模型(如 FLUX);当基座模型面对复杂布局或生僻 prompt 的指令遵循能力不足时,反演噪声依然难以挽救下游编辑失败。
  • 理论渐近性与有限步截断的落差:强收敛定理建立在严格的局部非扩张性假设及无穷步迭代上限之上;在极端紧凑的计算预算(如 \(K \le 2\))下,残差无法完全归零,仍存在局部漂移风险。
  • 未来探索方向:探索将平直度择根思想从时间步方向推广至空间结构自适应分配,或结合更高阶常微分方程求解器(如 Heun、Runge-Kutta 族)构建高阶定点修正算子。

相关工作与启发

  • vs ReNoise / AIDI / GNRI (传统定点反演):既有方法聚焦于设计更快的收敛步进(如加噪平滑或安德森加速)以求得任意不动点,忽略了多根流形的选择倾向;SelFix 在相同的残差容限下主动引导算法收敛至最具几何稳定性的平直根,显著抑制回放漂移。
  • vs RF-Solver / FireFlow (流专用高阶求解器):高阶数值求解器试图通过外推预测抵消局部曲率,但在粗粒度步长下难以保证步级逆变换的精确代数相消;SelFix 坚守定点反演的严格离散求逆性质,通过几何先验消解误差,取得了更高的重建精度与更佳的背景保持。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次揭示定点反演的多根现象并以平直度为准则提出因果择根机制。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 FLUX-dev 上的重建与编辑评测,严密对齐 NFE,消融与理论证明兼备。
  • 写作质量: ⭐⭐⭐⭐⭐ 动机严谨,数学推导清晰流畅,从实证分析到算法闭环逻辑自洽。
  • 价值: ⭐⭐⭐⭐⭐ 对 Rectified Flow 的无训练精确反演与下游高质量图像编辑具有重要实用价值。