Correcting Source Mismatch in Flow Matching with Radial-Angular Transport

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.

• The paper introduces RAFM which decouples the radial and angular components of transport to correct source mismatch, improving sample fidelity in high-dimensional and heavy-tailed settings.
• It replaces the standard Gaussian source with one that exactly matches the data’s radial law, eliminating the KL penalty from mismatched radial tails.
• Extensive benchmarks show that RAFM outperforms baseline methods by achieving higher quality samples with improved generation stability and faster computation.

Correcting Source Mismatch in Flow Matching with Radial–Angular Transport

Motivation and Background

Flow Matching (FM) is a robust framework for generative modeling, prescribing probability paths (via source distribution and interpolation maps) between a simple source and the target data distribution, and learning a neural vector field that realizes this transport. Canonically, FM initializes with an isotropic Gaussian source, leveraging analytical and computational convenience. However, this imposition can induce structural source mismatch—particularly along the radial dimension—when the data exhibit heavy-tailed or non-Gaussian norm statistics. This mismatch is not benign: the FM dynamics are then burdened by correcting the artificial radial discrepancy, confounding angular and radial learnability, and impeding sample fidelity, especially in high-dimensional or anisotropic regimes.

The paper introduces Radial–Angular Flow Matching (RAFM), which decomposes and corrects the source mismatch by decoupling the radial and angular components of the transport. RAFM constructs an explicit source distribution whose radial law matches the data and whose conditional angular distribution is uniform, thus isolating the source correction to the radial part. Subsequently, conditional paths are defined as geodesics on the appropriately scaled spheres, focusing the learned transport on angular alignment. This structural decoupling yields concrete theoretical and empirical advantages, notably enhanced sample fidelity for non-Gaussian, heavy-tailed, and extreme-event datasets.

Figure 1: Conceptual comparison of Gaussian FM (corrects both radius and angle over the flow) and RAFM (preserves radial law; rearranges mass predominantly via angular transport on spheres; right shows matched-radius spherical geodesic interpolation vs. linear Euclidean interpolation).

Formal Structure of RAFM

Radial Source Construction

Let $X \sim p_{\mathrm{data}} \subset \mathbb{R}^d$ with $R = \|X\|$ and $U = X/\|X\|$. RAFM replaces the canonical Gaussian source with $X_0 = R\,U_0$, where $R \sim p_R$ (the data radial law) and $U_0$ is uniform on $S^{d-1}$. The resulting source distribution $q_{\mathrm{rad}}$ is absolutely continuous off the origin and preserves the data's norm law exactly. Sampling is implemented practically via empirical distribution over the training set radii, with CDF and Wasserstein guarantees provided for the estimation error.

Theoretical analysis (via KL divergence decomposition) shows that the standard Gaussian source incurs an additional penalty $\mathrm{KL}(p_R\|p_{\chi_d})$—quantifying the cost of mismatched radial tails—that is eliminated in RAFM. After this correction, only the angular component remains to be modeled.

Spherical Geodesic Conditional Paths

Once the source is matched, transport aims to rearrange mass angularly at fixed radius. For each data pair, the interpolation follows the geodesic on $S^{d-1}$ of radius $R = \|X\|$0, connecting the randomly initialized direction $R = \|X\|$1 to the target direction $R = \|X\|$2. This produces flows that are globally radius-preserving, with the learned vector field tangent everywhere to the sphere for almost all points, and analytically computable via path derivatives.

The approach is compatible with standard, simulation-free FM training and does not unduly complicate architecture or optimization pipelines.

Theoretical Guarantees

The paper provides rigorous theoretical support for the RAFM construction:

• Norm Law Control: The marginal distribution of norm under the learned FM flow is governed entirely by the source under tangential vector fields. If the learned field is (or is projected to be) tangent to the sphere, the radial law is exactly preserved.
• Generation Stability: A strong stability result relates the population regression error (between the learned and ideal vector field over the reference paths) to the Wasserstein-2 distance between generated and data distributions. Specifically, if the FM regression loss is small, then generation is accurate, with an exponential amplification governed by the spatial Lipschitz constant of the vector field.

This decoupling elucidates the benefits of radial–angular separation: once the source is adapted, only the conditional angular transport remains as the locus of approximation.

Empirical Evaluation

Extensive benchmarks validate the framework’s claims, traversing synthetic (correlated Student-$R = \|X\|$3, anisotropic Gaussian) and real-world (planar PIV turbulence) non-Gaussian data. The evaluation decomposes the effect of source correction versus conditional path geometry and compares RAFM to baseline Gaussian FM, source-only correction, and the recent multiplicative (non-Gaussian noise) methods such as MSGM.

Key findings:

• Gaussian FM significantly degrades under radial tail mismatch, especially in higher dimensions.
• Simply matching the radial law at source (source-only) recovers a majority of the fidelity gap; the remaining gains from RAFM are manifest in the most severe (heavily-tailed, high-dimensional) regimes, showing that path geometry matters after source adaptation.
• Compared to MSGM, RAFM achieves competitive or superior quality with orders-of-magnitude less wall-clock time in high dimensions.

  Figure 2: Top—Comparison of radial laws: the empirical (RAFM) source closely matches data, while the Gaussian source has substantial mismatch. Bottom—Tails of generated sample norms: Gaussian FM fails to capture real tails, source-only corrects much of the deficit, RAFM provides highly accurate tail matching and sample quality; MSGM is competitive but slower.

Analysis of Limitations and Ablations

• Tangential Projection: Projecting the predicted velocity onto the tangent space of the sphere is critical for stability and fidelity in challenging (heavy-tailed, high-$R = \|X\|$4) settings, though less impactful in easier scenarios.
• Empirical Source Quality: Empirical and oracle (exact) source distributions yield similar results, justifying the approach on practical datasets.
• Low-Dimensional Fragility: Very low-dimensional (e.g., $R = \|X\|$5) radial–angular settings reveal numerical fragility near the origin due to angular degeneracy; this exposes a limitation of the geometric decoupling for such cases.

Implications and Future Directions

RAFM underscores that FM source distribution is not a neutral implementation detail but a structural component of the generative geometry—particularly for heavy-tailed, extreme-event, or fundamentally non-Gaussian data. By matching the radial law and focusing learning resources on angular alignment, RAFM achieves significantly better-aligned generative flows, yielding theoretical and practical improvements.

The framework opens several avenues:

• Angular transport generalization: More robust angular interpolations and non-spherical geometries could further improve low-dimensional and multi-modal data scenarios.
• Non-Gaussian conditional path extensions: The geometric decoupling concept can inform source-and-path design in other simulation-free or continuous-flow frameworks.
• Applications: RAFM is particularly suited for modeling scientific, financial, and physical data with extreme-value behavior, where Gaussianity as a baseline is fundamentally inadequate.

Conclusion

RAFM provides a principled geometric correction to Flow Matching generative models via explicit radial source adaptation and matched-radius angular transport. The framework addresses the critical challenge of artificial radial correction, improves learnability and fidelity on non-Gaussian datasets, and is competitive with, or surpasses, computationally more demanding alternatives. The results demonstrate that source-and-path co-design, with explicit accommodation of data geometry, is central for robust high-dimensional generative modeling in non-Gaussian settings.

Reference:

"Correcting Source Mismatch in Flow Matching with Radial-Angular Transport" (2604.04291)