Score-Regularized Joint Sampling with Importance Weights for Flow Matching
This work addresses the problem of high-variance estimates in flow matching models for researchers and practitioners in generative modeling, representing an incremental improvement with novel regularization and weighting techniques.
The paper tackled the challenge of estimating expectations from flow matching models under limited sampling budgets by proposing a non-IID sampling framework that jointly draws multiple samples to cover diverse, salient regions, resulting in diverse, high-quality samples and accurate estimates of importance weights and expectations.
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.