SoFlow: Direct Generative ODE Modeling

• SoFlow is a generative modeling framework that directly learns the closed-form solution of the velocity ODE underlying diffusion-based models, enabling efficient one- or few-step data generation.
• It utilizes a novel parameterization with complementary Flow Matching and Solution Consistency losses and leverages Diffusion Transformers in VAE latent space for superior performance.
• The method significantly reduces GPU memory usage and training time while achieving competitive FID scores compared to traditional multi-step diffusion and GAN approaches.

Solution Flow Models (SoFlow) constitute a generative modeling framework that directly learns the closed-form solution of the velocity ordinary differential equation (ODE) underlying diffusion-based models, enabling efficient one-step or few-step data sample generation. By explicitly modeling the mapping from a latent prior to data in a single network pass, SoFlow overcomes the inefficiency of traditional multi-step denoising approaches. The approach is characterized by a novel parameterization of the generative ODE’s solution, a pair of complementary loss functions—Flow Matching and Solution Consistency—and an architecture leveraging Diffusion Transformers (DiT) in VAE latent space, achieving state-of-the-art performance among one-step generative models on ImageNet 256×256 (Luo et al., 17 Dec 2025).

1. Mathematical Structure of the SoFlow Framework

SoFlow starts from a continuous interpolant (“noising process”) bridging data $x_0\sim p_{\rm data}$ and a tractable prior $x_1\sim\mathcal{N}(0,I)$: $$x_{t} = \alpha_{t}\,x_{0} + \beta_{t}\,x_{1},\qquad t\in[0,1]$$ with $\alpha_0=1$, $\beta_0=0$, $\alpha_1=0$, $\beta_1=1$, $\alpha,\beta\in C^1$. This yields a marginal velocity field

$$v(x_{t},t) = \mathbb{E}_{p(x_0,x_1\mid x_t)}\left[\alpha_t' x_0 + \beta_t' x_1\right]$$

defining the generative ODE,

$$\frac{dX(t)}{dt} = v\big(X(t),t\big),\qquad X(1)=x_{1}\sim\mathcal{N}(0,I)$$

to be solved backward in time. Rather than numerically integrating this ODE, SoFlow directly learns its solution function: $x_1\sim\mathcal{N}(0,I)$0 satisfying

$x_1\sim\mathcal{N}(0,I)$1

for $x_1\sim\mathcal{N}(0,I)$2. Thus $x_1\sim\mathcal{N}(0,I)$3 instantaneously maps any $x_1\sim\mathcal{N}(0,I)$4 to $x_1\sim\mathcal{N}(0,I)$5 in closed-form, fundamentally distinguishing SoFlow from velocity-based diffusion and flow-matching models (Luo et al., 17 Dec 2025).

2. Loss Function Design: Flow Matching and Solution Consistency

The parametric solution map is

$x_1\sim\mathcal{N}(0,I)$6

where $x_1\sim\mathcal{N}(0,I)$7 and $x_1\sim\mathcal{N}(0,I)$8 are known (e.g., Euler or trigonometric parameterizations), $x_1\sim\mathcal{N}(0,I)$9 is a neural network, and $$x_{t} = \alpha_{t}\,x_{0} + \beta_{t}\,x_{1},\qquad t\in[0,1]$$0, $$x_{t} = \alpha_{t}\,x_{0} + \beta_{t}\,x_{1},\qquad t\in[0,1]$$1. Flow Matching loss anchors the network’s instantaneous velocity: $$x_{t} = \alpha_{t}\,x_{0} + \beta_{t}\,x_{1},\qquad t\in[0,1]$$2 with $$x_{t} = \alpha_{t}\,x_{0} + \beta_{t}\,x_{1},\qquad t\in[0,1]$$3 an adaptive weight, $$x_{t} = \alpha_{t}\,x_{0} + \beta_{t}\,x_{1},\qquad t\in[0,1]$$4 sampled logit-normal. The correct instantaneous velocity is enforced via analytic differentiation of the solution map at $$x_{t} = \alpha_{t}\,x_{0} + \beta_{t}\,x_{1},\qquad t\in[0,1]$$5.

Classifier-Free Guidance (CFG) is incorporated by interpolating class-conditional and unconditional velocity estimates during loss computation.

To guarantee correct solution mapping over finite intervals, SoFlow introduces a Solution Consistency loss: $$x_{t} = \alpha_{t}\,x_{0} + \beta_{t}\,x_{1},\qquad t\in[0,1]$$6 where $$x_{t} = \alpha_{t}\,x_{0} + \beta_{t}\,x_{1},\qquad t\in[0,1]$$7 is a stopped-gradient network copy and $$x_{t} = \alpha_{t}\,x_{0} + \beta_{t}\,x_{1},\qquad t\in[0,1]$$8 is an adaptive weight. The final objective is a weighted sum

$$x_{t} = \alpha_{t}\,x_{0} + \beta_{t}\,x_{1},\qquad t\in[0,1]$$9

with $\alpha_0=1$0. Notably, $\alpha_0=1$1 involves no Jacobian-vector products, yielding superior training efficiency compared to flow-anchored objectives (Luo et al., 17 Dec 2025).

3. Model Architecture and Implementation Protocols

SoFlow adopts the Diffusion Transformer (DiT) backbone, operating in VAE latent space (32×32×4) for ImageNet 256×256 generation. Model variants include B/2 (131M), M/2 (308M), L/2 (459M), XL/2 (676M), with patch size 2×2 for B/2 upwards. Training from scratch uses batch size 256 and 240 epochs, AdamW optimizer (lr=$\alpha_0=1$2, betas=(0.9,0.99)), no weight decay or lr decay, and EMA 0.9999.

Hyperparameters:

• Time sampling: logit-normal for Losses ($\alpha_0=1$3 for $\alpha_0=1$4, $\alpha_0=1$5, $\alpha_0=1$6 for $\alpha_0=1$7).
• Noising schedule: linear ($\alpha_0=1$8) with Euler parameterization.
• CFG strength $\alpha_0=1$9 and velocity-mix $\beta_0=0$0 tuned per model size; $\beta_0=0$1 decays from 2.5/2.0 to 1.0 for large $\beta_0=0$2.
• CIFAR-10 experiments use a U-Net backbone, RAdam, batch size 1024, 800K steps, analogous settings (Luo et al., 17 Dec 2025).

4. Empirical Performance and Benchmarks

On ImageNet 256×256, SoFlow sets new FID-50K standards among one-step generative models across all tested DiT model scales:

Method	Params	1–NFE FID
MeanFlow B/2	131 M	6.17
SoFlow B/2	131 M	4.85
MeanFlow M/2	308 M	5.01
SoFlow M/2	308 M	3.73
MeanFlow L/2	459 M	3.84
SoFlow L/2	459 M	3.20
MeanFlow XL/2	676 M	3.43
SoFlow XL/2	676 M	2.96

For two function evaluations, SoFlow XL/2 achieves 2.66 FID (vs. 2.93 for MeanFlow XL/2). These results are competitive with multi-step diffusion, autoregressive, and GAN methods at comparable or lower NFE, with SoFlow’s performance realized at significantly reduced inference cost. SoFlow consistently outperforms MeanFlow, the previous strongest one-step baseline (Luo et al., 17 Dec 2025).

5. Sampling, Inference Efficiency, and Practical Implications

SoFlow’s generative sampling proceeds as:

1. Sample $\beta_0=0$3;
2. Produce $\beta_0=0$4 via a single forward pass.

Optional few-step sampling is possible by perturbing $\beta_0=0$5 and recursively invoking the solution map. The Solution Consistency loss eliminates JVPs, reducing peak GPU memory by 31% and enabling 23% faster training than MeanFlow on H100 GPUs. SoFlow inherits the computational efficiency of state-of-the-art attention kernels used in the DiT backbone (Luo et al., 17 Dec 2025).

6. Limitations, Extensions, and Forward-Looking Directions

SoFlow is currently best suited to scenarios demanding minimal NFE, with a tradeoff remaining at ultra-low NFE budgets compared to deep multi-step diffusion. Several extension directions are highlighted:

• Optimization of noising/interpolant schedules.
• Improved weighting schemes or variance reduction within the loss.
• Application to text-to-image, video, or hybrid few-step regimes.
• Empirical exploration of higher NFE hybrids (2–4 steps) for further FID improvements.

SoFlow’s formulation offers a platform for rapid progress in efficient generative modeling, with a unified framework that supports precise velocity field learning, CFG integration, and closed-form mapping from prior to data (Luo et al., 17 Dec 2025).

---

For the SoFlow “Semi-dilute-Flow” model relevant to polymer compression in Couette flow, see Dunstan (Dunstan, 2014). That distinct SoFlow framework models coil compression, not generative learning.