Free Lunch for Stabilizing Rectified Flow Inversion

TL;DR

This paper introduces Proximal-Mean Inversion and mimic-CFG, novel gradient correction methods for rectified flow models that improve stability, reconstruction, and editing quality without additional training.

Contribution

The paper proposes training-free gradient correction techniques that stabilize rectified flow inversion, enhancing reconstruction and editing performance with theoretical guarantees.

Findings

Significantly improved inversion stability and reconstruction quality.

Enhanced editing fidelity with reduced neural evaluations.

Achieved state-of-the-art results on PIE-Bench.

Abstract

Rectified-Flow (RF)-based generative models have recently emerged as strong alternatives to traditional diffusion models, demonstrating state-of-the-art performance across various tasks. By learning a continuous velocity field that transforms simple noise into complex data, RF-based models not only enable high-quality generation, but also support training-free inversion, which facilitates downstream tasks such as reconstruction and editing. However, existing inversion methods, such as vanilla RF-based inversion, suffer from approximation errors that accumulate across timesteps, leading to unstable velocity fields and degraded reconstruction and editing quality. To address this challenge, we propose Proximal-Mean Inversion (PMI), a training-free gradient correction method that stabilizes the velocity field by guiding it toward a running average of past velocities, constrained within a theoretically derived spherical Gaussian. Furthermore, we introduce mimic-CFG, a lightweight velocity correction scheme for editing tasks, which interpolates between the current velocity and its projection onto the historical average, balancing editing effectiveness and structural consistency. Extensive experiments on PIE-Bench demonstrate that our methods significantly improve inversion stability, image reconstruction quality, and editing fidelity, while reducing the required number of neural function evaluations. Our approach achieves state-of-the-art performance on the PIE-Bench with enhanced efficiency and theoretical soundness.