Efficient Constraint-Aware Flow Matching via Randomized Exploration

2025-08-18

AI Summary

This paper addresses constrained generation in flow matching, focusing on two constraint modeling settings: (1) constraints defined by a differentiable distance function, and (2) constraints accessible only via black-box membership queries. To this end, we propose a randomized exploration framework that requires no convexity assumptions or explicit barrier functions. Our method adopts a two-stage pipeline: first, coarse-grained constraint guidance is achieved via distance-regularized and randomized-mean-flow learning; second, fine-grained correction is performed using numerical optimization coupled with black-box membership queries. The framework unifies treatment of both explicit and implicit constraints, significantly improving constraint satisfaction rates while preserving high-fidelity alignment with the target distribution. We validate its effectiveness and generalizability across multiple synthetic benchmarks and a challenging real-world task—black-box adversarial example generation under hard-label constraints.

Technology Category

• Constraint Satisfaction and Optimization: Constraint Learning and Acquisition
• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Adversarial Learning & Robustness

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs

Abstract

We consider the problem of generating samples via Flow Matching (FM) with an additional requirement that the generated samples must satisfy given constraints. We consider two scenarios, viz.: (a) when a differentiable distance function to the constraint set is given, and (b) when the constraint set is only available via queries to a membership oracle. For case (a), we propose a simple adaptation of the FM objective with an additional term that penalizes the distance between the constraint set and the generated samples. For case (b), we propose to employ randomization and learn a mean flow that is numerically shown to have a high likelihood of satisfying the constraints. This approach deviates significantly from existing works that require simple convex constraints, knowledge of a barrier function, or a reflection mechanism to constrain the probability flow. Furthermore, in the proposed setting we show that a two-stage approach, where both stages approximate the same original flow but with only the second stage probing the constraints via randomization, is more computationally efficient. Through several synthetic cases of constrained generation, we numerically show that the proposed approaches achieve significant gains in terms of constraint satisfaction while matching the target distributions. As a showcase for a practical oracle-based constraint, we show how our approach can be used for training an adversarial example generator, using queries to a hard-label black-box classifier. We conclude with several future research directions. Our code is available at https://github.com/ZhengyanHuan/FM-RE.

Problem

• Generating samples via Flow Matching with constraint satisfaction
• Handling constraints via differentiable distance or membership oracle
• Achieving computational efficiency and high constraint satisfaction rates

Innovation

• Penalizes distance to constraint set via differentiable function
• Uses randomization to learn mean flow satisfying constraints
• Two-stage approach improves computational efficiency with constraints