Constraint-Aware Flow Matching via Randomized Exploration

1. Problem Statement

The paper addresses the challenge of constrained generative modeling using Flow Matching (FM). While standard FM models effectively learn to transform a source distribution (e.g., Gaussian noise) into a target data distribution, they often fail to satisfy sample-wise constraints (e.g., an image must have specific attributes, or a generated fluid dynamic state must remain within physical bounds).

The core difficulty arises because:

1. Constraint Violations: Vanilla generative models produce samples that may fall outside the valid constraint set $C$, even if the target distribution is theoretically supported within $C$.
2. Gradient Issues: In many practical scenarios, the constraint set $C$ is defined only by a membership oracle (a black-box function that returns 1 if $x \in C$ and 0 otherwise). Since the indicator function $1_C(x)$ is non-differentiable almost everywhere, standard gradient-based optimization cannot directly enforce constraints during training.
3. Limitations of Existing Methods: Current approaches often require convex constraints, knowledge of barrier functions, or reflection mechanisms that rely on explicit boundary normals, limiting their applicability to complex, non-convex, or black-box constraints.

The authors formalize the problem under two specific scenarios:

• Case (a): A differentiable distance function $d(x, C)$ to the constraint set is available.
• Case (b): Only a membership oracle $1_C(x)$ is available (no distance function or boundary gradient).

---

2. Methodology

The paper proposes two distinct methods to address these scenarios, collectively termed Constraint-Aware Flow Matching (CAFM).

A. FM-DD: Constraint-Aware Flow Matching with Differentiable Distance

• Applicability: Used when a differentiable distance function $d(x, C)$ is known (e.g., convex sets or smooth manifolds).
• Mechanism: The authors modify the standard FM objective by adding a penalty term proportional to the distance between the generated sample and the constraint set.
• Objective Function:
  $$\min_\theta \int_0^1 \mathbb{E}[\|u_\theta(X_t, t) - \dot{\Psi}_t(X_0, X_1)\|^2] dt + \lambda \mathbb{E}[d(X_1^\theta, C)]$$
  Here, $u_\theta$ is the learned velocity field, $\Psi_t$ is the linear interpolant, and $\lambda$ is a weighting hyperparameter. This approach penalizes samples that are far from the constraint set, encouraging the flow to stay within $C$.

B. FM-RE: Constraint-Aware Flow Matching via Randomized Exploration

• Applicability: Used when only a membership oracle is available (the most general and challenging case).
• Core Insight: Since the gradient of the indicator function is zero almost everywhere, the authors introduce randomization into the velocity field to create a differentiable expectation of constraint satisfaction.
• Two-Stage Joint Optimization:
  1. Stage 1 (Deterministic): A standard FM model is trained for the time interval $t \in [0, t_0]$ to learn the base flow dynamics.
  2. Stage 2 (Stochastic): For $t \in [t_0, 1]$, the velocity field is perturbed with Gaussian noise: $\tilde{U}_t = u_\theta(X_t, t) + \sigma_t W$, where $W \sim \mathcal{N}(0, I)$.
• Gradient Estimation: By treating the final sample $X_1$ as a random variable dependent on the noise $W$, the constraint term becomes $\mathbb{E}[1_C(X_1)]$. Using the REINFORCE trick (score function estimator), the gradient with respect to parameters $\theta$ and noise scale $\sigma$ is estimated as:
  $$\nabla \mathbb{E}[1_C(X_1)] \approx \mathbb{E}_T \left[ \sum \nabla \log \pi(\tilde{U}_t | X_t) \cdot 1_C(X_1) \right]$$
  This allows the model to learn a "mean flow" that statistically maximizes the probability of satisfying the constraint.
• Efficiency: The two-stage approach (fixing the first stage and optimizing the second) is shown to be more computationally efficient than a single-stage approach where randomization occurs from $t=0$.

---

3. Key Contributions

1. Formalization of Constrained FM: The paper rigorously defines constrained flow matching under both differentiable distance and membership oracle settings, bridging a gap in the literature.
2. Randomized Exploration Strategy (FM-RE): A novel method to handle non-differentiable constraints by introducing stochasticity into the flow, enabling gradient-based optimization via the membership oracle.
3. Two-Stage Optimization: The derivation of a computationally efficient two-stage training procedure that separates the learning of the base distribution dynamics from the constraint exploration, reducing training costs while maintaining performance.
4. Broad Applicability: Unlike reflection-based or projection-based methods that require convexity or explicit boundary geometry, FM-RE works with non-convex, disconnected, and empty-interior constraints (e.g., subspaces) and black-box oracles.
5. Empirical Validation: Extensive experiments on synthetic data, MNIST attribute generation, and adversarial example generation against black-box classifiers.

---

4. Experimental Results

The authors evaluated their methods against baselines including standard FM, Reflected Flow Matching (RFM), Mirror Diffusion Models (MDM), and Projected Diffusion Models (PDM).

• Synthetic Experiments:

  • Metrics: Sliced Wasserstein Distance (SWD) for distribution matching and Constraint Violation Rate ($P(X_1 \notin C)$).
  • Findings: FM-DD and FM-RE achieved significantly lower constraint violation rates than vanilla FM while maintaining competitive SWD scores.
  • Complex Constraints: In "2 boxes" (disconnected, non-convex) and "Subspace" (empty interior) scenarios, methods like RFM and MDM failed (N/A) or performed poorly. FM-RE successfully satisfied constraints in these difficult settings.
  • Hyperparameter Sensitivity: Increasing the constraint weight $\lambda$ reduced violations but slightly degraded distribution matching. Starting randomization later ($t_0 > 0$) reduced training time without significant performance loss.

• MNIST Attribute Generation:

  • Tasks: Generating digits with specific "Brightness" (number of white pixels) and "Thickness" constraints.
  • Results: FM-RE achieved a constraint violation rate of 1.12% (vs. 9.14% for FM) for brightness and 5.40% (vs. 23.2% for FM) for thickness, with minimal impact on Fréchet Inception Distance (FID).

• Adversarial Example Generation:

  • Task: Generating adversarial examples for hard-label black-box classifiers (LeNet-5 on MNIST, ResNet-50 on CIFAR-10) where the constraint is "predicted label $\neq$ ground truth."
  • Results: FM-RE successfully generated imperceptible perturbations that reduced classifier accuracy from ~99% to 18.7% (MNIST) and 28.2% (CIFAR-10), demonstrating its ability to handle complex, non-differentiable, black-box constraints.

---

5. Significance and Conclusion

This paper makes a significant contribution to the field of constrained generative modeling by:

• Generalizing Constraints: Moving beyond the requirement for convexity or explicit boundary knowledge, enabling the application of flow matching to real-world problems where constraints are defined by complex oracles (e.g., physical simulations, black-box classifiers).
• Bridging Theory and Practice: Providing a practical, gradient-based solution for non-differentiable constraints through randomized exploration, a technique previously more common in reinforcement learning.
• Efficiency: Demonstrating that a two-stage approach can balance computational cost with constraint satisfaction, making it viable for high-dimensional problems.

The proposed FM-RE framework is particularly notable for its ability to handle black-box constraints, opening new avenues for applications in robust AI (adversarial training), scientific simulation (physics constraints), and controlled content generation. The code is publicly available, facilitating further research and adoption.