Correcting Source Mismatch in Flow Matching with Radial-Angular Transport

TL;DR

This paper introduces Radial--Angular Flow Matching (RAFM), a novel framework that corrects source mismatch in flow matching for heavy-tailed and anisotropic data, improving transport accuracy without complex modifications.

Contribution

RAFM explicitly matches the radial distribution of the source to data, reducing the transport problem to angular alignment and providing theoretical and empirical improvements.

Findings

RAFM improves flow matching performance on heavy-tailed data.

The framework offers explicit density and stability guarantees.

Empirically, RAFM outperforms standard Gaussian flow matching.

Abstract

Flow Matching is typically built from Gaussian sources and Euclidean probability paths. For heavy-tailed or anisotropic data, however, a Gaussian source induces a structural mismatch already at the level of the radial distribution. We introduce \textit{Radial--Angular Flow Matching (RAFM)}, a framework that explicitly corrects this source mismatch within the standard simulation-free Flow Matching template. RAFM uses a source whose radial law matches that of the data and whose conditional angular distribution is uniform on the sphere, thereby removing the Gaussian radial mismatch by construction. This reduces the remaining transport problem to angular alignment, which leads naturally to conditional paths on scaled spheres defined by spherical geodesic interpolation. The resulting framework yields explicit Flow Matching targets tailored to radial--angular transport without modifying the underlying deterministic training pipeline. We establish the exact density of the matched-radial source, prove a radial--angular KL decomposition that isolates the Gaussian radial penalty, characterize the induced target vector field, and derive a stability result linking Flow Matching error to generation error. We further analyze empirical estimation of the radial law, for which Wasserstein and CDF metrics provide natural guarantees. Empirically, RAFM substantially improves over standard Gaussian Flow Matching and remains competitive with recent non-Gaussian alternatives while preserving a lightweight deterministic training procedure. Overall, RAFM provides a principled source-and-path design for Flow Matching on heavy-tailed and extreme-event data.