Decomposing Out-of-Distribution Error in Conditional Flow Matching via Wasserstein Geometry

Conditional flow matching has emerged as a powerful generative modeling framework that learns a vector field to transport an initial distribution toward a target data distribution. However, theoretical understanding of its out-of-distribution (OOD) performance under unseen conditions remains limited. In this work, we establish a rigorous geometric formulation to decompose the source of generalization error. We treat the conditional task as a map from the condition space to the Wasserstein space and derive a generalization bound under a coarse embedding assumption. The resulting decomposition separates OOD error into three tractable components: Interpolation Sparsity, Geometric Distortion, and In-Distribution Fit. Our empirical evaluation confirms that this framework demonstrates three key functions: (1) it acts as a diagnostic tool that tracks the dynamics of generalization during training; (2) it identifies dataset-specific failure modes (e.g., topological gaps, geometric instability); and (3) it enables mathematically motivated interventions that yield predictable gains by minimizing specific terms.

Modern AI systems can generate images, sounds, and other content on demand, but they are often asked to create examples they have never seen before, which can lead to unexpected failures. This paper introduces a mathematical framework to diagnose why these failures occur, showing that they stem from one of three sources: insufficient diversity in the training data, a mismatch between the model’s internal representation and the task it faces, or incomplete training. Like a medical scan, our framework pinpoints the root cause of failure rather than treating the model as a black box. Experiments on both synthetic and real image datasets show that once the cause is identified, targeted improvements consistently lead to better performance.