---
title: Reflected Flow Matching
url: https://www.emergentmind.com/topics/reflected-flow-matching
type: topic
---

# Reflected Flow Matching

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 [2405.16577].

## 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 $v_\theta$, 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 $[0, 255]^d$ for image data. Recent reflected diffusion models have sought to address this but lacked closed-form scores and introduced additional statistical bias [2405.16577].

## 2. Mathematical Formulation of Reflected CNFs

To constrain flows within $\Omega$, RFM augments the CNF ODE with a reflecting term:
$$
d\phi_t(x) = v_t(\phi_t(x))\,dt + dL_t, \qquad \phi_0(x) = x,
$$
where $L_t$ is a process of bounded variation reflecting any outward normal movement at the boundary $\partial \Omega$ (Definition 2.1).

The evolution of the density $p_t$ then satisfies the continuity equation with a Neumann boundary condition:
$$
\partial_t p_t(x) = -\nabla \cdot (p_t(x) v_t(x)), \qquad p_t(x) v_t(x) \cdot n(x)\big|_{x \in \partial \Omega} = 0,
$$
where $n(x)$ is the outward unit normal at boundary points.

At the conditional level (for flows connecting $x_0 \sim p_0$ to $x_1 \sim p_{\mathrm{data}}$), conditional reflected flows $\phi_t(\cdot|x_1)$ are constructed, inducing time-varying densities and velocities via inverse flow transformations:
$$
v_t(x|x_1) = \frac{\partial}{\partial t}\, \phi_t(\phi_t^{-1}(x|x_1)), \qquad
p_t(x|x_1) = p_0(\phi_t^{-1}(x|x_1)) \left| \det \nabla_x \phi_t^{-1}(x|x_1) \right|.
$$
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
$$
\mathcal{L}_{\rm RFM}(\theta) = \int_0^1 \mathbb{E}_{x \sim p_t} \left\| v_\theta(x, t) - v_t(x) \right\|^2 dt.
$$
Alternatively, the conditional objective is
$$
\mathcal{L}_{\rm CRFM}(\theta) = \int_0^1 \mathbb{E}_{x_1 \sim p_{\mathrm{data}},\, x \sim p_t(\cdot|x_1)} \left\| v_\theta(x, t) - v_t(x|x_1) \right\|^2 dt,
$$
with gradients matching those of $\mathcal{L}_{\rm RFM}$. Training thus proceeds by sampling triplets $(x_0, x_1, t)$, constructing intermediate states $x_t = \phi_t(x_0\,|\,x_1)$, and regressing $v_\theta(x_t, t)$ toward the closed-form conditional velocities $v_t(x_t\,|\,x_1)$.

## 4. Analytical Conditional Flow Construction and Bias Avoidance

For convex $\Omega$, RFM adopts an "optimal-transport-style" conditional flow:
$$
\phi_t(x|x_1) = (1 - (1 - \sigma_{\min}) t)x + (1 - \sigma_{\min})t x_1,
$$
remaining strictly within $\Omega$ for all $t$ by convexity. The resulting velocity field is
$$
v_t(x|x_1) = \frac{1 - \sigma_{\min}}{1 - (1 - \sigma_{\min}) t} (x_1 - x).
$$
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 $(x_0, x_1, t)$, compute $x_t = \phi_t(x_0|x_1)$, regress $v_\theta(x_t, t)$ to $v_t(x_t|x_1)$ | No reflection; analytic flows |
| Sampling     | Numerically solve ODE with $v_\theta$; at each out-of-domain step, reflect at $\partial \Omega$ as needed per Algorithm 1           | Reflection at boundary; ODE solver integration |

During sampling, each ODE step that leads outside $\Omega$ 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:
$$
\tilde{v}_t(x|c) = w\,v_t(x|c) + (1-w)\,v_t(x|\emptyset),
$$
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 ($w=15$), FM oversaturates 51–64% of pixels; RFM maintains all values within $[0,255]$. RFM achieves FID ≈ 26.7, versus FM (31.8) and RDM (≫35 for low NFE) (Figure 4, 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 $v_t$ (Theorem 2.1, Pilipenko 2014). The error in the output distribution, measured in squared 2-Wasserstein distance, is bounded by the regression loss:
$$
W_2^2(p_1, p_{\theta,1}) \leq e^{1+2M} \mathcal{L}_{\rm RFM}(\theta)
$$
for convex $\Omega$ and M-Lipschitz $v_\theta$ (Theorem 2.3).

RFM retains the architectural flexibility of FM, supporting arbitrary priors and plug-and-play design of conditional flows constrained to $\Omega$. However, $\Omega$ must have a regular analytic boundary and admit a tractable family of differentiable conditional flows that remain entirely within $\Omega$ for all $t$.

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 [2405.16577].

Source: https://www.emergentmind.com/topics/reflected-flow-matching