Score-Regularized Joint Sampling with Importance Weights for Flow Matching

TL;DR

This paper introduces a score-regularized joint sampling method with importance weights for flow matching models, improving the diversity, quality, and expectation estimation of generated samples under limited sampling budgets.

Contribution

It proposes a non-IID joint sampling framework with score-based regularization and importance weighting, enhancing flow matching models' sampling diversity and estimation accuracy.

Findings

Produces diverse, high-quality samples

Accurately estimates importance weights and expectations

Advances reliable characterization of flow models

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.