Importance-Weighted Non-IID Sampling for Flow Matching Models

TL;DR

This work tackles reliable estimation of functionals under flow-matching models when sampling budgets are tight, where IID sampling yields high variance. It introduces an importance-weighted non-IID joint sampling framework that jointly evolves multiple trajectories with a score-based diversity velocity to cover diverse, high-density regions while staying on the data manifold, and a learned residual velocity $r_\phi$ to model the non-IID marginal for unbiased weighting. The key contributions are (i) score-based regularization that preserves on-manifold diversity, (ii) the first method to learn a residual flow for unbiased importance weights in non-IID sampling, and (iii) comprehensive empirical validation showing improved diversity, sample quality, and accurate weight/expectation estimates on Gaussian mixtures and downstream tasks like text-to-image generation and image inpainting. The approach enables more reliable characterization of flow-matching model outputs under fixed budgets, with practical implications for downstream AI systems that rely on accurate distributional expectations.

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 with high-impact outcomes dominate the expectation. We propose an importance-weighted non-IID sampling framework that jointly draws multiple samples to cover diverse, salient regions of a flow's distribution while maintaining unbiased estimation via estimated importance weights. To balance diversity and quality, we introduce a score-based regularization for the diversity mechanism, 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. We further develop the first 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. 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.

Figures (5)

• Figure 1: 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.
• Figure 2: Illustration of our score-based regularization for the diversity mechanism.
• Figure 3: 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.
• Figure 4: 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.
• Figure S1: Joint ODE trajectories on a three-component Gaussian mixture. 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 circles). 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.

Theorems & Definitions (6)

• Definition 1: Original ODE
• Definition 2: Joint ODEs
• Definition 3: Marginal ODE
• Theorem 1: Evolution of the Importance Weight at Fixed Position
• Theorem 2: Evolution of the Importance Weight Along the Trajectory
• Corollary 1: Evolution under Rectified Flows