SoFlow: Solution Flow Models for One-Step Generative Modeling

AI summary

Overview

Research area: Generative modeling — specifically few-step and one-step generation with diffusion and Flow Matching models.

Technical level: Advanced. The paper builds on Flow Matching theory, neural ODEs, and consistency models, and its derivations rely on Jacobian and partial differential equation manipulations.

Scope: The paper introduces SoFlow, a from-scratch training framework that learns the closed-form solution of the Flow Matching velocity ODE, enabling high-quality one-step image generation without per-step Jacobian-vector products.

What This Paper Is About

Diffusion and Flow Matching models generate images by iteratively solving an ordinary differential equation dozens of times, which is slow. Prior one-step methods either cannot cleanly support Classifier-Free Guidance (CFG), or they require expensive Jacobian-vector product (JVP) computations that deep learning frameworks like PyTorch do not handle efficiently. SoFlow solves both problems by training a neural network to represent the exact solution function of the velocity ODE, so a single forward pass yields a clean image.

Key Contributions

1. A solution-function formulation of generative modeling. The authors define the function f(x_t, t, s) that maps a noisy state at time t directly to its ODE-evolved state at time s, and derive two conditions (a boundary condition and a PDE identity) that any valid solution function must satisfy.

2. Two complementary training losses without JVPs. They introduce a Flow Matching loss (a special case at t = s) and a solution consistency loss (the s < t case) that together train the network, eliminating the Jacobian-vector product computations that bottleneck competing methods like MeanFlow.

3. Training-time Classifier-Free Guidance. Because the model separately predicts the unconditional velocity field, CFG can be baked in during training rather than applied only at inference, and a velocity mixing ratio m is used to suppress the variance introduced by guidance.

4. State-of-the-art 1-NFE results on ImageNet 256×256. With identical DiT architectures and training epochs, SoFlow improves 1-NFE FID-50K over MeanFlow at every model size (e.g., 2.96 vs. 3.43 for XL/2).

Main Findings

• Consistent improvement across model sizes: On ImageNet 256×256 with 1-NFE, SoFlow achieves FID-50K scores of 4.85 (B/2), 3.73 (M/2), 3.20 (L/2), and 2.96 (XL/2), outperforming MeanFlow by roughly 0.5–1.3 FID points at each scale.

• Solution consistency loss alone works, but combining losses is better: A 0% Flow Matching data ratio still produces one-step samples (FID 66.36 in ablation), but the best results come from a 75% Flow Matching ratio (FID 58.57).

• Robust loss weighting matters: Using a power-weighted MSE with p = 1 dramatically outperforms plain MSE (58.57 vs. 68.25), consistent with findings in prior consistency models.

• Linear noising with Euler parameterization is best: The linear schedule x_t = (1-t)x_0 + t x_1 paired with the Euler form f_θ = x_t + (s-t)F_θ gives FID 11.59, substantially better than trigonometric variants (17.50–19.01).

• Small velocity mixing ratio m is beneficial: m = 0.25 (FID 11.59) beats m = 1.0 (FID 17.22), confirming that substituting part of the high-variance stochastic target with the model's own guided velocity prediction stabilizes training.

• Training-time CFG is powerful: Increasing CFG strength w from 1.5 to 3.0 improves FID from 34.47 to 11.59, showing guided training is highly effective for this framework.

• l → t schedule choice is minor: Exponential, cosine, linear, and constant schedules produce similar FID (58.57–60.17), so the transition speed of the consistency target has limited impact.

Methodology in Plain English

The team starts from the observation that a Flow Matching model's iterative sampling process is just numerically integrating an ODE. Instead of teaching a network to take small ODE steps, they teach it the entire solution of that ODE as a single function of three arguments: the current noisy state, the current time, and the target time.

To make this learnable, they prove that a network qualifies as a valid solution if it satisfies two identities — one guaranteeing the network returns its input when the source and target times coincide, and one linking the network's derivatives to the underlying velocity field. The first identity is enforced by a special parameterization; the second becomes the training objective.

Two training losses follow. The "Flow Matching" loss handles the case where source and target times are equal, essentially recovering the standard velocity-prediction objective. The "solution consistency" loss handles the case where the target time is earlier — it asks f evaluated at (x_t, t, s) to match f evaluated at a slightly earlier state that the stop-gradient copy of the model produces. This target mimics a Taylor expansion of the ODE trajectory but requires no Jacobian-vector product. A small auxiliary neural network head predicts the unconditional velocity, letting the model estimate a CFG-guided velocity field during training.

For inference, the trained model is applied once with t = 1 and s = 0 to convert Gaussian noise into a clean latent sample, which the VAE decoder turns into an image.

Why This Matters

Impact on research: SoFlow shows that one-step generative modeling is achievable without the JVP computational overhead that has limited recent consistency and mean-flow approaches. Its solution-function framing also provides a clean theoretical bridge between Flow Matching, consistency training, and classical ODE solution theory, which may inform future work on fast samplers.

Real-world applications:

• Real-time image and video generation on consumer GPUs or in interactive design tools, where multi-step sampling is currently too slow.
• Mobile and edge deployment of high-quality generative models, since one network evaluation is far cheaper than dozens.
• Latency-critical pipelines such as AR/VR content creation, live avatar rendering, or on-device photo editing.
• Reduced-API-cost generative services, where inference time and energy scale almost linearly with the number of function evaluations.

Industry relevance: Inference cost is a major barrier to deploying diffusion models at scale. A from-scratch, one-step generator that preserves image quality and supports CFG is directly useful to any company shipping image or video generation as a product.

Future Directions

• Extending beyond image generation. The formulation is domain-agnostic, so applying it to video, 3D assets, audio, or molecule generation is a natural next step.

• Scaling to larger models. The paper reports promising trends up to DiT-XL/2; determining whether SoFlow retains its advantage at billion-parameter scale or with longer training schedules remains open.

• Better guidance during training. The velocity mixing ratio and variance-suppression tricks hint that more principled CFG training strategies could yield further gains.

• Theoretical improvements to the consistency loss. Since the current target comes from a first-order Taylor approximation, higher-order or adaptive variants might tighten the consistency objective.

• Combining with distillation. The framework is presented as from-scratch only, but using it to distill existing large multi-step models is an obvious practical application.

Target Audience

This paper is best suited for machine learning researchers and graduate students working on generative models, particularly those already familiar with diffusion, Flow Matching, or consistency models. Practitioners building latency-sensitive generative systems and engineers evaluating one-step samplers for deployment will also find the empirical results directly relevant. Readers unfamiliar with ODEs, Jacobians, or Flow Matching fundamentals will need to consult related work before fully engaging with the derivations.

Authors’ 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.