IWN-Flow: Importance-Weighted Non-IID Sampling for Flow Matching Models

Paper Code
Illustration of importance-weighted non-IID sampling.
Illustration of importance-weighted non-IID sampling. Under IID sampling, both samples are likely drawn from the same dominant mode. In contrast, diversity velocity encourages samples to diverge along their trajectories, leading to coverage of multiple modes. To correct the resulting sampling bias, importance weights are required. Intuitively, $w_1 > 1 > w_2$, since non-IID sampling draws the second sample from a minor mode.

Abstract

Flow matching models effectively represent complex distributions, yet estimating expectations of functions of their outputs remains challenging under limited sampling budgets. Independent sampling often yields high-variance estimates, especially when rare but high-impact outcomes dominate the expectation. We propose a non-IID sampling framework that jointly draws multiple samples to cover diverse, salient regions of a flow matching model's generative distribution. To balance diversity and quality, we introduce a score-based regularization for the diversity mechanism (SR), which uses the score function, i.e., the gradient of the log probability, to ensure samples are pushed apart within high-density regions of the data manifold, mitigating off-manifold drift. To enable unbiased estimation when desired, we further develop an approach for importance weighting of non-IID flow samples by learning a residual velocity field that reproduces the marginal distribution of the non-IID samples and by evolving importance weights along trajectories. Empirically, our method produces diverse, high-quality samples and accurate estimates of both importance weights and expectations, advancing the reliable characterization of flow matching model outputs. Our code will be publicly available on GitHub.

Method

Our approach addresses both (G1) diversity with quality and (G2) unbiasedness by modifying the sampling of a pre-trained flow. We generate a batch of $n$ samples jointly, $X^{(1:n)} = (X^{(1)}, \dots, X^{(n)})$, from velocity $v(x,t)$ with base distribution $p_0(x)$. We encourage diversity while preserving on-manifold quality and assign an importance weight to each sample so that the expectation estimator remains unbiased. We achieve this with two components: (1) score-based diversity velocity regularization that pushes samples apart primarily along high-density directions, and (2) a residual velocity for importance weighting added to $v$ to model each sample's marginal distribution under the joint sampler.

Qualitative Results

Illustration of importance-weighted non-IID sampling.
Qualitative results for text-to-image generation. Although diverse, Sample 1 from DPP appears unreasonable. Adding SR (hard) preserves diversity while making it reasonable. For Sample 2, SR (hard) refines the cat's eyes.
Illustration of importance-weighted non-IID sampling.
Qualitative results for image inpainting. Two samples are shown from the joint samples, with all methods sharing identical initialization. In Sample 1, DPP introduces artifacts (highlighted by the black rectangle) that are removed by SR (hard). While enhancing quality, DPP+SR (hard) retains the diversity of DPP, as illustrated in Sample 2.
Joint ODE trajectories on a three-component Gaussian mixture (video). All methods start from identical initial states. "IID" is the baseline without a diversity objective and collapses to the highest-weight mode. "Non-IID (DPP)" uses a determinantal point process objective to encourage sample diversity, enabling the three trajectories to discover all three modes. "SR" denotes our score-based diversity-velocity regularization (the soft version is used in this example). Combining DPP + SR pulls samples toward the underlying modes while preserving coverage (white rectangle). "IW" denotes importance-weight estimation. The full method (DPP + SR + IW) further assigns an importance weight to each trajectory; larger markers indicate higher estimated weight (white circle). For visualization, arrow lengths are rescaled nonlinearly while preserving their relative ordering. The background shows the relative value of $\log p(x)$, where $p$ is the target density corresponding to the IID ODE; yellow indicates higher density and purple indicates lower density.