Reflected Flow Matching

TL;DR

Reflected Flow Matching (RFM) extends continuous normalizing flows to constrained domains by introducing a reflection term in the governing ODEs, yielding reflected CNFs that stay inside the domain. It trains the velocity model via simulation-free conditional velocity-field matching (CRFM), using analytically derived conditional flows to avoid bias and improve stability on boundaries. The method achieves competitive or superior results on low-dimensional constrained tasks and image benchmarks (CIFAR-10 and ImageNet64), while enforcing zero boundary violations under high guidance weights. This approach offers efficient, boundary-consistent generative modeling for data with domain constraints and broad applicability to constrained-generation problems.

Abstract

Continuous normalizing flows (CNFs) learn an ordinary differential equation to transform prior samples into data. Flow matching (FM) has recently emerged as a simulation-free approach for training CNFs by regressing a velocity model towards the conditional velocity field. However, on constrained domains, the learned velocity model may lead to undesirable flows that result in highly unnatural samples, e.g., oversaturated images, due to both flow matching error and simulation error. To address this, we add a boundary constraint term to CNFs, which leads to reflected CNFs that keep trajectories within the constrained domains. We propose reflected flow matching (RFM) to train the velocity model in reflected CNFs by matching the conditional velocity fields in a simulation-free manner, similar to the vanilla FM. Moreover, the analytical form of conditional velocity fields in RFM avoids potentially biased approximations, making it superior to existing score-based generative models on constrained domains. We demonstrate that RFM achieves comparable or better results on standard image benchmarks and produces high-quality class-conditioned samples under high guidance weight.

Figures (9)

• Figure 1: Illustration of the conditional velocity fields on two nonconvex domains: half annulus (left) and cup (right). The black solid curve is the designed conditional velocity field within the domain. The OT conditional velocity field is represented by the blue dashed segment which violates the domain constraint.
• Figure 2: The histplots of samples obtained by different methods compared to the ground truth on the two-dimensional hypercube, simplex, and cup data set. Samples out of the constrained domain are plotted with red dots. The total sample size is 100,000.
• Figure 3: Class-conditioned guided samples from the CNFs trained with FM (left, $\textrm{NFE}=769$) and the reflected CNFs trained with RFM (right, $\textrm{NFE}=723$) with a high guidance weight $w=15$ on ImageNet ($64\times 64$). The samples from CNFs are clipped to $[0,255]$ after the ODE simulation.
• Figure 4: The histplots of samples obtained by different methods compared to the ground truth on the two-dimensional hypercube, simplex, and cup data set. We generate the samples from CNFs and reflected CNFs using the Dormand–Prince method dormand1980dopri with absolute and relative tolerances of $10^{-5}$. Samples out of the constrained domain are plotted with red dots. The total sample size is 100,000.
• Figure 5: Velocity field in reflected CNFs learned by RFM for different $t$s on two-dimensional generation tasks.

Theorems & Definitions (9)

• Theorem 3.1
• Remark 3.2
• Theorem 3.3
• proof
• Theorem 3.4
• Theorem 3.5: Wasserstein Bound
• Example 3.6: Convex Domain
• Example 3.7: Half Annulus
• Theorem 3.8