SoFlow: Solution Flow Models for One-Step Generative Modeling

TL;DR

SoFlow is a one-step generative model trained from scratch that learns the solution function of the velocity ODE — a bi-time map taking a state at time t to its evolved state at time s — using a Flow Matching loss plus a “solution consistency” loss. Crucially the consistency loss is designed to avoid Jacobian-vector products (JVPs), which are the main training bottleneck in prior JVP-based one-step methods (MeanFlow et al.) because they are poorly optimized in PyTorch. The Flow Matching branch also lets SoFlow incorporate Classifier-Free Guidance during training, which is rare for one-step-from-scratch models. On ImageNet 256×256 with matched DiT backbones and equal training epochs, SoFlow beats MeanFlow in 1-NFE FID across all DiT-{B,M,L,XL}/2 sizes (e.g. 2.96 vs 3.43 at XL/2).

Key claims

The bi-time solution function f_θ(x_t, t, s) can be trained directly to map noise at t=1 to clean data at s=0 in a single forward pass, parameterized via Euler or trigonometric forms that enforce the boundary f_θ(x_t, t, t) = x_t [§3].
A solution consistency loss using a Taylor expansion of f_θ around nearby times — with a stop-gradient on the target branch — enforces self-consistency without requiring the Jacobian-vector product (JVP) that recent one-step methods need [§3].
A Flow Matching loss applied to the implicit velocity field of f_θ enables Classifier-Free Guidance to be applied during training via a mixed guided velocity target, yielding a one-step-from-scratch model that still benefits from CFG at inference [§3].
On ImageNet 256×256 with matched DiT architectures and equal training epochs, SoFlow’s 1-NFE FID-50K beats MeanFlow at every model size: B/2 4.85 vs 6.17; M/2 3.73 vs 5.01; L/2 3.20 vs 3.84; XL/2 2.96 vs 3.43 [Table].
Increasing to 2-NFE further improves quality (XL/2 reaches 2.66 FID), so the same model can trade a second forward pass for fidelity when needed [§4].
On CIFAR-10 (pixel-space, U-Net backbone), 1-NFE results are competitive with other fast-sampling methods, indicating the approach is not specific to latent-space DiT [§4].

Method

SoFlow learns a single network f_θ(x_t, t, s) representing the solution of the velocity ODE — i.e. for a noised sample at time t, f_θ directly outputs the corresponding sample at time s. Two parameterizations are studied: an Euler form (linear-in-time blend between identity and a learned residual) and a trigonometric form, both of which bake in the boundary condition f_θ(x_t, t, t) = x_t.

Training combines two losses. (1) A Flow Matching loss regresses the implicit velocity (∂f_θ/∂s)|_{s=t} against the marginal velocity field of the noising process; during conditional training this target is the guided marginal velocity (a mix of conditional and unconditional fields), with a velocity mixing ratio m trading off variance vs. guidance signal. (2) A solution consistency loss enforces f_θ(x_t, t, s) ≈ f_θ(x_t, t, s') when s and s’ are close, using a first-order Taylor expansion around a stop-gradient target — the algebra cancels the JVP that MeanFlow-style consistency requires, so training fits cleanly inside vanilla PyTorch.

Inference for image generation is a single forward pass: sample x_1 ~ N(0, I) and output f_θ(x_1, 1, 0). Optional 2-NFE sampling takes one step to an intermediate time and re-noises before a final jump, trading latency for FID.

Results

ImageNet 256×256 class-conditional, DiT backbone, SD-VAE latents, equal training epochs vs. MeanFlow baseline (1-NFE FID-50K, lower better): B/2 4.85 vs 6.17, M/2 3.73 vs 5.01, L/2 3.20 vs 3.84, XL/2 2.96 vs 3.43. The gap is largest at smaller scales (~1.3 FID at B/2) and persists at XL/2 (0.47 FID). Two-step sampling drops XL/2 further to 2.66. CIFAR-10 unconditional (U-Net, 32×32 pixel space) is competitive with prior 1-NFE methods. The paper emphasizes that training is more memory- and time-efficient than JVP-based baselines because the consistency loss is JVP-free.

Why it’s interesting

One-step generation from scratch is the cleanest test of whether a model can compress all of the multi-step denoising work into a single learned map, and SoFlow is currently the leading clean baseline against MeanFlow in that regime — with two pragmatic wins (no JVPs, supports CFG during training) that matter for anyone who has tried to scale MeanFlow-style training in PyTorch. The solution-function framing is also conceptually different from consistency models: rather than enforcing endpoint agreement, it directly parameterizes the bi-time flow map and uses Taylor-based local consistency as the only auxiliary objective. For Luma, the relevance is twofold: (a) the JVP-free recipe transfers straight to DiT-based video/image stacks; (b) training-time CFG without the usual two-pass overhead is a meaningful efficiency lever for class- or text-conditional one-step generators.