Straighten Viscous Rectified Flow via Noise Optimization

TL;DR

This paper introduces VRFNO, a novel noise optimization framework for rectified flow models that improves image generation quality and speed by addressing limitations in previous methods like Reflow.

Contribution

VRFNO integrates an encoder and neural velocity field with noise reparameterization, enhancing trajectory prediction and coupling accuracy for better image synthesis.

Findings

VRFNO outperforms Reflow in synthetic and real datasets.

Achieves state-of-the-art results in one-step and few-step generation.

Effectively mitigates errors caused by previous rectified flow methods.

Abstract

The Reflow operation aims to straighten the inference trajectories of the rectified flow during training by constructing deterministic couplings between noises and images, thereby improving the quality of generated images in single-step or few-step generation. However, we identify critical limitations in Reflow, particularly its inability to rapidly generate high-quality images due to a distribution gap between images in its constructed deterministic couplings and real images. To address these shortcomings, we propose a novel alternative called Straighten Viscous Rectified Flow via Noise Optimization (VRFNO), which is a joint training framework integrating an encoder and a neural velocity field. VRFNO introduces two key innovations: (1) a historical velocity term that enhances trajectory distinction, enabling the model to more accurately predict the velocity of the current trajectory, and (2) the noise optimization through reparameterization to form optimized couplings with real images which are then utilized for training, effectively mitigating errors caused by Reflow's limitations. Comprehensive experiments on synthetic data and real datasets with varying resolutions show that VRFNO significantly mitigates the limitations of Reflow, achieving state-of-the-art performance in both one-step and few-step generation tasks.