SoFlow: Solution Flow Models for One-Step Generative Modeling

TL;DR

SoFlow introduces a one-step generative modeling framework that improves efficiency and performance by leveraging velocity and solution function relationships, along with novel loss functions, outperforming existing models on ImageNet.

Contribution

The paper proposes Solution Flow Models (SoFlow), a novel one-step generative approach with new loss functions that enhance training efficiency and generation quality.

Findings

Achieves better FID-50K scores than MeanFlow on ImageNet 256x256.

Uses a flow matching loss for velocity estimation during training.

Employs a solution consistency loss that avoids Jacobian-vector product calculations.

Abstract

The multi-step denoising process in diffusion and Flow Matching models causes major efficiency issues, which motivates research on few-step generation. We present Solution Flow Models (SoFlow), a framework for one-step generation from scratch. By analyzing the relationship between the velocity function and the solution function of the velocity ordinary differential equation (ODE), we propose a Flow Matching loss and a solution consistency loss to train our models. The Flow Matching loss allows our models to provide estimated velocity fields for Classifier-Free Guidance (CFG) during training, which improves generation performance. Notably, our consistency loss does not require the calculation of the Jacobian-vector product (JVP), a common requirement in recent works that is not well-optimized in deep learning frameworks like PyTorch. Experimental results indicate that, when trained from scratch using the same Diffusion Transformer (DiT) architecture and an equal number of training epochs, our models achieve better FID-50K scores than MeanFlow models on the ImageNet 256x256 dataset.