Free Lunch for Stabilizing Rectified Flow Inversion

1. The Problem: Instability in Rectified Flow Inversion

Generative models based on Rectified Flow (RF) have emerged as high-performance alternatives to traditional diffusion models, offering high-quality generation and the ability to perform inversion tasks (mapping observed data back to their latent representation) without additional training. However, existing inversion methods (such as "vanilla" RF-based inversion) suffer from a critical problem:

• Accumulative Approximation Errors: During the inversion process, the latent is not exactly accessible, so it must be approximated using the observed data at each time step. This introduces estimation errors at every step.
• Velocity Field Instability: These errors accumulate over time steps, perturbing the velocity field, which should be nearly constant in an ideal RF.
• Consequences: This instability causes the inversion trajectory to deviate toward low-density regions of the data distribution, resulting in significant degradation in image reconstruction quality and fidelity in editing tasks (such as preserving the original structure and content).

2. Proposed Methodology

The authors propose two main methods, both training-free and designed to correct the perturbed velocity field:

A. Proximal-Mean Inversion (PMI)

Aimed at stabilizing the inversion and reconstruction process.

• Core Concept: Instead of following the pure velocity prediction (which contains errors), PMI guides the current velocity toward a running average of the velocities predicted from the start up to the current step.
• Correction Mechanism:
  1. Calculate the average of historical velocities.
  2. Apply a proximal update that pushes the current velocity toward the direction of this global average.
  3. Theoretical Constraint: To prevent the corrected trajectory from deviating into low-density regions (where reconstruction would fail), each correction step is restricted within a spherical Gaussian. The radius of this sphere is derived theoretically based on an instability analysis, ensuring that sampling points remain in high-density regions.
• Efficiency: It requires no additional neural function evaluations (NFE) or training and integrates easily into any ODE solver (Euler, Heun, RF-Solver, FireFlow).

B. Mimic-CFG (Editing)

Aimed at balancing structural consistency and editing capability in image modification tasks.

• Inspiration: It is based on the structure of Classifier-Free Guidance (CFG), but applied to velocity instead of log-likelihood scores.
• Mechanism:
  1. Calculate the average of past velocities during the editing phase.
  2. Project the current predicted velocity onto the direction of this historical average.
  3. Perform a linear interpolation between the original predicted velocity and its projection.
• Objective: This allows maintaining the structure of the original image (background preservation) while permitting sufficient deviation to achieve the desired semantic change, avoiding both excessive instability and a lack of editability.

3. Key Contributions

1. Proposal of PMI: A training-free method that corrects the perturbed velocity field via a gradient correction toward a running average, effectively reducing inversion errors in RF models.
2. Proposal of Mimic-CFG: A lightweight velocity correction strategy for editing tasks that mimics the CFG structure, interpolating between the current velocity and its projection onto the historical average direction.
3. Theoretical Foundation: Derivation of a correction radius based on the instability analysis of high-dimensional ODEs, guaranteeing that corrected trajectories remain in high-probability density regions.
4. SOTA Performance: Establishment of a new state-of-the-art (SOTA) on the PIE-Bench benchmark, achieving superiority in reconstruction and editing with fewer sampling steps and no additional NFE.

4. Experimental Results

Experiments were conducted on the PIE-Bench benchmark (700 images, 10 editing categories) using the Flux.1-dev model as the base.

• Reconstruction and Inversion:
  • Integrating PMI with base solvers (Euler, Heun, RF-Solver, FireFlow) consistently improved image quality metrics: PSNR (signal-to-noise ratio) and SSIM (structural similarity) increased, while MSE (mean squared error) and LPIPS (perceptual similarity) decreased.
  • Improvements were notable in "unconditional" configurations, demonstrating that the method stabilizes the trajectory even without text prompts.
• Image Editing:
  • The method significantly improved CLIP similarity (semantic editing fidelity) and background preservation (measured by PSNR/SSIM in unedited areas).
  • An optimal balance was achieved between the ability to edit content and maintaining the original structure.
• Efficiency:
  • The proposed methods allowed reaching or surpassing the performance of base methods with fewer sampling steps (reduced NFE). For example, FireFlow with correction achieved superior results with only 18 steps compared to base variants requiring more.
  • No additional computational inference costs were incurred (no extra NFE required).

5. Significance and Impact

This work addresses a fundamental limitation in rectified flow models: the inherent instability in high-performance inversion. By providing a theoretical and practical solution (PMI and Mimic-CFG) that is plug-and-play (easy to integrate) and training-free, the authors demonstrate that it is possible to stabilize inversion without sacrificing computational efficiency.

The importance of this study lies in:

• Enabling more robust and higher-quality image editing and reconstruction applications.
• Reducing the gap between the theory of rectified flows and their practical application in inverse tasks.
• Offering a theoretical framework to understand and mitigate instability in high-dimensional ODEs, which is crucial for the future development of efficient generative models.

In summary, the article presents a "free lunch" in the sense that it achieves substantial improvements in stability and quality without the cost of retraining models or increasing the computational inference load.