SubFlow: Sub-mode Conditioned Flow Matching for Diverse One-Step Generation

Abstract: Flow matching has emerged as a powerful generative framework, with recent few-step methods achieving remarkable inference acceleration. However, we identify a critical yet overlooked limitation: these models suffer from severe diversity degradation, concentrating samples on dominant modes while neglecting rare but valid variations of the target distribution. We trace this degradation to averaging distortion: when trained with MSE objectives, class-conditional flows learn a frequency-weighted mean over intra-class sub-modes, causing the model to over-represent high-density modes while systematically neglecting low-density ones. To address this, we propose SubFlow, Sub-mode Conditioned Flow Matching, which eliminates averaging distortion by decomposing each class into fine-grained sub-modes via semantic clustering and conditioning the flow on sub-mode indices. Each conditioned sub-distribution is approximately unimodal, so the learned flow accurately targets individual modes with no averaging distortion, restoring full mode coverage in a single inference step. Crucially, SubFlow is entirely plug-and-play: it integrates seamlessly into existing one-step models such as MeanFlow and Shortcut Models without any architectural modifications. Extensive experiments on ImageNet-256 demonstrate that SubFlow yields substantial gains in generation diversity (Recall) while maintaining competitive image quality (FID), confirming its broad applicability across different one-step generation frameworks. Project page: https://yexionglin.github.io/subflow.

• The paper introduces SubFlow that refines flow matching through explicit sub-mode conditioning, mitigating dominant mode bias and restoring intra-class diversity.
• It employs K-Means clustering on pre-trained DINOv3 features to partition data into semantic sub-modes, enabling efficient one-step ODE integration.
• Experimental results on ImageNet-256 show that SubFlow improves Recall and FID, validating its effectiveness in recovering lost minority modes.

SubFlow: Sub-mode Conditioned Flow Matching for Diverse One-Step Generation

Motivation and Problem Statement

Recent advancements in flow matching and diffusion-based generative models have enabled the generation of high-fidelity samples across multiple data modalities. To accelerate inference, few-step and single-step generative frameworks (e.g., MeanFlow, Shortcut Models, SoFlow) have emerged, allowing direct mapping from noise to data with reduced computational cost. However, a critical limitation persists: a marked degradation in generative diversity, particularly significant mode collapse, where samples concentrate on dominant modes and rare, yet semantically valid, sub-modes are systematically suppressed. This diversity loss is especially problematic in downstream tasks that require capturing infrequent but salient variations, and it substantially limits the practical deployment of rapid generative models.

The root of this phenomenon is traced to averaging distortion: when class-conditional flow matching is trained with an MSE objective, the learned vector field approximates a frequency-weighted mean over all sub-modes available within a class. Sub-modes with greater density dominate the learning process, causing rare sub-modes to be structurally neglected. This dominant-mode bias is entrenched in the learned vector field itself and is not mitigated by increasing the number of ODE integration steps during inference, as further confirmed by the empirical findings in this work.

Methodology: Sub-mode Conditioned Flow Matching (SubFlow)

SubFlow directly targets and resolves averaging distortion by reframing the conditioning granularity of flow matching from the entire class to fine-grained semantic sub-modes. The key methodological contribution is the augmentation of the conditioning variables: in addition to the class label $c$, SubFlow introduces a sub-mode index $k$, where each training sample is clustered into one of $K$ semantic partitions using features extracted from a pre-trained vision foundation model (DINOv3). Specifically, within each class, K-Means clustering is applied to the embedded feature space to identify K semantically meaningful sub-modes. Each sample is assigned a deterministic sub-mode label $k$ in this offline pre-processing stage.

During training, the vector field $v_\theta(x_t, t, c, k)$ is regressed to the exact transport direction of the matched data pair within class and sub-mode, thereby learning sub-mode-specific transport velocities. In contrast with the standard class-conditional formulation, each (class, sub-mode) sub-distribution is approximately unimodal, removing the ambiguity and averaging distortion present in the classical approach. The model is trained via a modified conditional flow matching objective, with k always retained (even when performing classifier-free guidance by dropping $c$), ensuring structural consistency.

At inference, given a desired class, a sub-mode index is sampled from the empirical prior $p(k | c)$ constructed from the training data, and the ODE is integrated in a single step (or multi-step, if desired). This mechanism induces sample diversity that accurately reflects the full intra-class variation, restoring coverage of minority sub-modes efficiently.

Notably, SubFlow is architecturally non-invasive: it can be implemented on top of any flow-matching-based one-step model without modifying the backbone network, ODE solver, or base pipeline. All that is required is augmenting the conditioning pathway with an embedding for $k$, rendering SubFlow strictly plug-and-play.

Experimental Evaluation

The empirical assessment is conducted on the ImageNet-256 benchmark, integrating SubFlow into three state-of-the-art one-step frameworks: MeanFlow, Shortcut Models, and SoFlow. The evaluation metrics include FID (Fréchet Inception Distance) for overall image quality, and separate Precision/Recall scores to measure fidelity and diversity.

Strong quantitative results are observed:

• MeanFlow-B/2+SubFlow achieves a Recall of 48.84% (vs. 43.45% for MeanFlow) and a FID improvement (5.86 vs. 6.17).
• SoFlow-B/2+SubFlow and Shortcut-B/2+SubFlow both exhibit consistent Recall and FID improvements.
• Compared to adversarial and multi-step methods, SubFlow-augmented one-step models achieve comparable diversity with a fraction of the sampling cost and without adversarial training.
• Increasing the number of integration steps in MeanFlow actually further entrenches the dominant-mode bias and does not materially improve Recall, confirming that the root of diversity loss is in the learned vector field, not the ODE discretization.

Qualitative analyses demonstrate that conditioning on different sub-mode indices yields visually distinct and semantically meaningful samples from the same noise initialization, confirming that SubFlow not only increases diversity scores but genuinely recovers lost intra-class variability.

Ablation studies validate key components:

• K-Means for sub-mode discovery outperforms alternatives (such as DBSCAN), due to more stable and even partitioning.
• Semantic clustering is required; random or uninformative sub-mode assignments confer no diversity benefit.
• Empirical prior-based sub-mode sampling achieves the best trade-off between Recall and FID.
• Retaining the sub-mode index during classifier-free guidance is crucial; dropping it erodes the advantages of sub-mode conditioning.
• Recall improvements stabilize as K increases, with quality degradation for excessively large K.

Theoretical and Practical Implications

The identification of averaging distortion as the central failure mode for diversity in few/few-step flow matching has major implications. It clarifies why accelerating generative inference by naïvely reducing the number of steps structurally sacrifices minority modes—a limitation not resolvable by improved ODE solvers or increased computational budgets at inference.

SubFlow provides a modular and computationally efficient solution, requiring only a modest offline clustering phase and negligible online overhead. Plug-and-play compatibility, architectural invariance, and strictly improved diversity at parity of fidelity make SubFlow an attractive and versatile addition for a wide variety of generative modeling contexts.

These findings suggest that, as generative modeling migrates toward extremely fast, one-step architectures—especially in class-conditional and prompt-driven image, audio, and video synthesis—explicit conditional signal refinement (e.g., sub-mode decomposition) is essential to avoid entrenched generative bias and mode collapse. More broadly, SubFlow exemplifies a class of solutions where simple conditioning augmentation corrects for profound statistical mismatches induced by standard objectives.

Future Directions

The sub-mode conditioning strategy outlined by SubFlow is general and extensible:

• It can be integrated with other semantic partitioning algorithms, potentially allowing hierarchical or data-adaptive sub-mode assignment.
• The principle extends beyond class conditioning, suggesting applications to multi-attribute or structured generative conditioning scenarios.
• Although currently demonstrated for image synthesis, the approach should generalize to audio, video, and multimodal generation, contingent on the existence of adequate feature extractors.
• There is scope for dynamic, adaptive sub-mode discovery which evolves as the base model improves or as more data become available.

Unifying sub-mode conditioning with advanced guidance techniques, latent-space architectures, and controllable generation opens new research trajectories for both theoretical understanding and practical deployment in real-world high-throughput, high-diversity media generation systems.

Conclusion

SubFlow isolates and resolves a major obstacle for diversity in modern one-step generative models. By augmenting the conditioning set with semantic sub-modes discovered via unsupervised clustering on robust feature extractors, SubFlow enforces lossless intra-class variation and restores sample diversity without architectural or computational penalty. The alignment between theoretical motivation and empirical gains underscores SubFlow’s utility for future generative modeling endeavours, emphasizing the centrality of structural conditioning in high-speed, high-quality, and high-diversity generation (2604.12273).