Reflected Flow Matching

TL;DR

Reflected Flow Matching (RFM) improves continuous normalizing flows by incorporating boundary constraints, enabling more natural samples on constrained domains through a simulation-free training approach that outperforms existing models.

Contribution

The paper introduces Reflected Flow Matching, a novel method that incorporates boundary constraints into flow matching for CNFs, improving sample quality on constrained domains.

Findings

RFM achieves comparable or better results on image benchmarks.

RFM produces high-quality class-conditioned samples with high guidance.

Analytical velocity fields in RFM avoid biased approximations.

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.