Importance-Weighted Non-IID Sampling for Flow Matching Models

2025-11-21

AI Summary

Under limited sampling budgets, expectation estimation in flow matching models suffers from high variance due to rare, high-impact events under independent sampling. This work proposes an unbiased importance-weighted non-i.i.d. sampling framework—the first to integrate importance weighting into the flow matching generative process. Our method learns a residual velocity field guided by the score function, jointly reconstructing the target marginal distribution and estimating sample importance weights via diversity regularization. A score-based regularization term further enforces moderate separation of samples in high-density regions, mitigating off-manifold drift. Experiments demonstrate that the approach preserves estimator unbiasedness while significantly improving sample diversity and quality. Consequently, it yields more accurate and robust expectation estimates, enhancing both interpretability and reliability of flow matching model outputs.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Search and Optimization: Sampling/Simulation-based Search
• Computer Vision: Diffusion Models for Vision

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Estimating expectations of flow-matching model outputs under limited sampling budgets
• Reducing high variance in estimates caused by rare high-impact outcomes
• Balancing diversity and quality in non-IID sampling while maintaining unbiased estimation

Innovation

• Importance-weighted non-IID sampling for unbiased estimation
• Score-based regularization ensures diversity within high-density regions
• Residual velocity field learning for importance weight estimation