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

TL;DR

SubFlow introduces a method to improve diversity in one-step generative models by decomposing classes into sub-modes, addressing averaging distortion and enhancing mode coverage without architectural changes.

Contribution

It proposes SubFlow, a plug-and-play approach that conditions flow models on sub-mode indices to restore diversity and coverage in one-step generation.

Findings

SubFlow significantly improves generation diversity (Recall) on ImageNet-256.

It maintains competitive image quality (FID) while enhancing diversity.

SubFlow integrates seamlessly into existing models without architectural modifications.

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.