Reflected Flow Matching

• Reflected Flow Matching is an extension of flow matching that integrates boundary reflection into continuous normalizing flows to strictly enforce hard domain constraints.
• It utilizes closed-form conditional velocity fields to eliminate biases and inefficiencies found in previous reflected diffusion models while ensuring Neumann boundary adherence.
• Empirical evaluations show that RFM achieves high-fidelity image generation with zero boundary violations and improved performance metrics compared to standard methods.

Reflected Flow Matching (RFM) is an extension of flow matching for continuous normalizing flows (CNFs) that enforces hard domain constraints—particularly in bounded or otherwise constrained domains such as the cube $[0, 255]^d$ relevant to image generation—by incorporating boundary reflection into both the model formulation and sampling process. RFM parameterizes the CNF velocity field to obey Neumann boundary conditions, preventing generated paths from exiting the legal sample space and thus avoiding artifacts such as oversaturated images. The methodology supports closed-form conditional velocity fields, eliminating bias and computational inefficiency present in earlier reflected diffusion models, while enabling high-fidelity generative modeling and principled constraint adherence (Xie et al., 2024).

1. Problem Setup and Motivation

Continuous normalizing flows are defined by time-dependent diffeomorphisms $\phi_t: \Omega \rightarrow \Omega$ parameterized by an ordinary differential equation (ODE)

$$\frac{d \phi_t(x)}{dt} = v_t(\phi_t(x)),\quad \phi_0(x)=x,$$

which transport a tractable prior $p_0$ supported on a domain $\Omega$ to an approximate data distribution $p_1 \approx p_{\rm data}$ via the pushforward densities $p_t$.

In standard flow matching (FM), a neural velocity model $v_\theta$ is trained to regress toward a prescribed target velocity field $v_t(x)$ by minimizing the mean squared error

$$\mathcal{L}_{\rm FM}(\theta) = \int_0^1 \mathbb{E}_{x \sim p_t} \| v_\theta(x, t) - v_t(x) \|^2 dt.$$

Simulation or regression errors in $\phi_t: \Omega \rightarrow \Omega$0, as well as inaccuracies in numerical ODE solvers, can push CNF state trajectories beyond constrained domains, resulting in invalid samples. This is particularly problematic in domains such as $\phi_t: \Omega \rightarrow \Omega$1 for image data. Recent reflected diffusion models have sought to address this but lacked closed-form scores and introduced additional statistical bias (Xie et al., 2024).

2. Mathematical Formulation of Reflected CNFs

To constrain flows within $\phi_t: \Omega \rightarrow \Omega$2, RFM augments the CNF ODE with a reflecting term:

$\phi_t: \Omega \rightarrow \Omega$3

where $\phi_t: \Omega \rightarrow \Omega$4 is a process of bounded variation reflecting any outward normal movement at the boundary $\phi_t: \Omega \rightarrow \Omega$5 (Definition 2.1).

The evolution of the density $\phi_t: \Omega \rightarrow \Omega$6 then satisfies the continuity equation with a Neumann boundary condition:

$\phi_t: \Omega \rightarrow \Omega$7

where $\phi_t: \Omega \rightarrow \Omega$8 is the outward unit normal at boundary points.

At the conditional level (for flows connecting $\phi_t: \Omega \rightarrow \Omega$9 to $$\frac{d \phi_t(x)}{dt} = v_t(\phi_t(x)),\quad \phi_0(x)=x,$$0), conditional reflected flows $$\frac{d \phi_t(x)}{dt} = v_t(\phi_t(x)),\quad \phi_0(x)=x,$$1 are constructed, inducing time-varying densities and velocities via inverse flow transformations:

$$\frac{d \phi_t(x)}{dt} = v_t(\phi_t(x)),\quad \phi_0(x)=x,$$2

Integration yields the marginal densities and velocities, both globally satisfying the Neumann condition (Theorem 2.1).

3. Reflected Flow Matching Objective and Training

The RFM objective is the marginal mean squared error

$$\frac{d \phi_t(x)}{dt} = v_t(\phi_t(x)),\quad \phi_0(x)=x,$$3

Alternatively, the conditional objective is

$$\frac{d \phi_t(x)}{dt} = v_t(\phi_t(x)),\quad \phi_0(x)=x,$$4

with gradients matching those of $$\frac{d \phi_t(x)}{dt} = v_t(\phi_t(x)),\quad \phi_0(x)=x,$$5. Training thus proceeds by sampling triplets $$\frac{d \phi_t(x)}{dt} = v_t(\phi_t(x)),\quad \phi_0(x)=x,$$6, constructing intermediate states $$\frac{d \phi_t(x)}{dt} = v_t(\phi_t(x)),\quad \phi_0(x)=x,$$7, and regressing $$\frac{d \phi_t(x)}{dt} = v_t(\phi_t(x)),\quad \phi_0(x)=x,$$8 toward the closed-form conditional velocities $$\frac{d \phi_t(x)}{dt} = v_t(\phi_t(x)),\quad \phi_0(x)=x,$$9.

4. Analytical Conditional Flow Construction and Bias Avoidance

For convex $p_0$0, RFM adopts an "optimal-transport-style" conditional flow:

$p_0$1

remaining strictly within $p_0$2 for all $p_0$3 by convexity. The resulting velocity field is

$p_0$4

For nonconvex domains, analogous piecewise-linear or polar flows guarantee constraint satisfaction (see Examples 2.1–2.2). Closed-form conditional velocities enable RFM to sidestep biases introduced by numerical approximations or eigenfunction expansions required in reflected diffusion models.

5. Algorithmic Framework

The RFM methodology divides into distinct procedures for training and sampling:

Phase	Key Steps	Distinctive Features
Training	Sample $p_0$5, compute $p_0$6, regress $p_0$7 to $p_0$8	No reflection; analytic flows
Sampling	Numerically solve ODE with $p_0$9; at each out-of-domain step, reflect at $\Omega$0 as needed per Algorithm 1	Reflection at boundary; ODE solver integration

During sampling, each ODE step that leads outside $\Omega$1 is truncated at the boundary, the velocity is reflected across the tangent, and integration resumes until the prescribed step length is consumed. This guarantees that all generated sample paths remain feasible.

Class-conditional generation employs weighted flow guidance:

$\Omega$2

inducing a Neumann boundary for the guided density.

6. Empirical Evaluation and Benchmarks

RFM has been evaluated on both synthetic and real-world settings:

• Toy domains (hypercube, simplex, half-annulus, cup): RFM matches or exceeds FM in KL divergence, completely eliminating constraint violations (Tables 1–2).
• Unconditional CIFAR-10 (32×32): RFM achieves FID 4.76 (NFE ≈ 139), outperforming FM (FID 6.35, NFE ≈ 142), and obtains zero boundary violations compared to significant violations in FM. RDM outperforms in FID (≈2.7) but at much higher NFE (2000) and does not ensure zero boundary violation (Table 3).
• Class-conditioned ImageNet 64×64: Under strong guidance ($\Omega$3), FM oversaturates 51–64% of pixels; RFM maintains all values within $\Omega$4. RFM achieves FID ≈ 26.7, versus FM (31.8) and RDM (≫35 for low NFE) (Figure 1, Tables 4–5).
• Ablation studies confirm RFM's performance is robust to the ODE solver, number of function evaluations (NFE), and the reflection enforcement.

7. Theoretical Insights, Generality, and Limitations

Existence and uniqueness of solutions for the reflected ODE in RFM are established under mild regularity and Lipschitz continuity of $\Omega$5 (Theorem 2.1, Pilipenko 2014). The error in the output distribution, measured in squared 2-Wasserstein distance, is bounded by the regression loss:

$\Omega$6

for convex $\Omega$7 and M-Lipschitz $\Omega$8 (Theorem 2.3).

RFM retains the architectural flexibility of FM, supporting arbitrary priors and plug-and-play design of conditional flows constrained to $\Omega$9. However, $p_1 \approx p_{\rm data}$0 must have a regular analytic boundary and admit a tractable family of differentiable conditional flows that remain entirely within $p_1 \approx p_{\rm data}$1 for all $p_1 \approx p_{\rm data}$2.

Potential extensions include application to higher-dimensional domains (e.g., high-resolution images), more intricate nonconvex domains, and other modalities such as point clouds and sequential motion data (Xie et al., 2024).