SoFlow: Solution Flow Models for One-Step Generative Modeling (2512.15657v1)

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.

• The paper introduces SoFlow, a novel framework that directly learns the solution function of the generative ODE to enable one-step synthesis and bypass iterative denoising.
• It combines adaptive flow matching and solution consistency losses with an Euler parameterization to enhance training stability and support classifier-free guidance for high-quality conditional generation.
• Empirical results demonstrate that SoFlow outperforms prior methods on benchmarks like ImageNet and CIFAR-10, achieving improved FID scores with significant computational efficiency gains.

Solution Flow Models for One-Step Generative Modeling: Technical Summary

Motivation and Foundations

The advent of Diffusion and Flow Matching models substantially advanced deep generative modeling, primarily via iterative denoising or velocity-field transport between prior and data distributions. However, both paradigms are fundamentally bottlenecked by multi-step sampling processes, which result in significant latency in generation tasks. Recent work focused on few-step or one-step generation frameworks—such as Consistency Models and related teacher-free training approaches—has alleviated sampling cost but introduced optimization instability and limited compatibility with standard conditional generation techniques like Classifier-Free Guidance (CFG). SoFlow introduces a new one-step generative framework, directly learning the solution function of the velocity ODE underlying Flow Matching, and circumvents the computational inefficiency and instability induced by previous approaches.

Solution Flow Model Formulation

Central to SoFlow is the direct modeling of the solution function $f(x_t, t, s)$, which maps an arbitrary state $x_t$ at time $t$ to its evolved state $x_s$ at time $s$ under the velocity ODE. The neural parameterization is constructed to satisfy critical boundary and PDE constraints:

$f_\theta(x_t, t, t) = x_t$

$$\partial_1 f_\theta(x_t, t, s)v(x_t,t) + \partial_2 f_\theta(x_t, t, s) = 0$$

Various parameterizations (Euler, trigonometric) are explored for $f_\theta$, with the linear/Euler form demonstrating empirically superior performance.

Training Objectives

Training combines two losses:

• Flow Matching Loss aligns the velocity predicted by the network at the boundary $t=s$ to ground-truth velocity via a robust, adaptively-weighted MSE, designed to avoid poorly-optimized JVPs and promote stable gradients.
• Solution Consistency Loss exploits Taylor expansion to relate $f_\theta(x_t, t, s)$ and $f_\theta(x_t + v(x_t, t)(l-t), l, s)$, driving local consistency without costly second-order derivatives.

This dual-objective formulation underpins efficient optimization and obviates reliance on ODE solvers for sampling.

Figure 1: Diagram of the solution consistency loss, showing the trajectory and direct alignment of evolved and reference states under the model.

Classifier-Free Guidance in SoFlow

A major practical advancement of SoFlow is its native support for CFG during training. By mixing conditional and unconditional velocity predictions within training batches—and employing a velocity mix ratio $m$ to interpolate between ground-truth and model-predicted guided velocities—SoFlow enables enhanced conditional sample quality in the strict 1-NFE regime. Empirical ablation demonstrates CFG significantly boosts FID measures, and tuning the guidance strength $w$ and mixing $m$ regulates variance to maximize generative performance.

Empirical Results

ImageNet 256$\times$256 (Latent VAE Space, DiT Backbone)

SoFlow achieves consistent improvements over MeanFlow across all tested model sizes when trained from scratch under matched conditions. For DiT-XL/2 architecture, SoFlow yields an FID-50K of 2.96 at 1-NFE and 2.66 at 2-NFE, outperforming MeanFlow-XL/2's 3.43 and 2.93, respectively. Ablation shows exponential ($l\to t$) schedule, high Flow Matching ratio ($\lambda=0.75$), linear-noising/Euler parameterization, and robust adaptive weighting $p=1.0$ yield optimality.

Figure 2: Qualitative comparison of 1-NFE SoFlow samples and quantitative FID improvements over MeanFlow on ImageNet 256$\times$256.

Figure 3: FID-50K as a function of model size, displaying monotonic gains as parameters increase.

Figure 4: Curated class-conditional 1-NFE samples by the XL/2 SoFlow model on ImageNet 256$\times$256, illustrating high fidelity single-step generation.

CIFAR-10 (Pixel Space, U-Net Backbone)

On CIFAR-10, SoFlow matches or exceeds strong baselines (e.g., iCT, ECT, MeanFlow, IMM) with FID-50K of 2.86 at 1-NFE, indicating the generality of the method beyond VAE-latent transformers.

Figure 5: Uncurated 1-NFE unconditional samples by the SoFlow-trained U-Net model on CIFAR-10, demonstrating competitive one-step synthesis.

Computational Efficiency

SoFlow's consistent loss forms circumvent JVP calculations, enabling compatibility with memory-efficient attention backends in PyTorch. Measured training speed increases by 23% and memory footprint drops by 31% compared to MeanFlow under matched implementation.

Theoretical Analysis

A rigorous bounding argument shows that SoFlow’s local consistency objectives ensure the global error between the model prediction and the true solution function is at most proportional to the residual and the training step schedule. Further, the ODE error for the predicted solution is controlled via the PDE error, under mild smoothness and Lipschitz assumptions, providing strong theoretical correctness guarantees for one-step sampling.

Practical and Theoretical Implications

SoFlow advances efficient, high-quality one-step generative models, mitigating computational overhead, and expanding compatibility with conditional sampling. From the perspective of future generative modeling frameworks:

• Direct neural solution of generative ODEs may be reliably scaled to even higher dimensions or multimodal spaces
• Native train-time CFG broadens utility in conditional generation tasks
• Avoidance of second-order derivatives positions SoFlow for integration with emerging acceleration techniques and hardware
• Theoretical foundation opens avenues for further reduction in sampling steps without loss of model capacity

Conclusion

SoFlow delineates a robust and efficient solution to one-step generative modeling, integrating structural consistency-based objectives with Flow Matching, and yielding superior empirical and computational outcomes relative to prior art. The paradigm is particularly promising for applications demanding high-fidelity single-step synthesis, low-latency inference, and native conditional guidance, and sets the stage for further explorations in neural ODE solution modeling for generative learning.