Coupled Flow Matching

TL;DR

CPFM presents a controllable dimensionality-reduction framework that learns coupled continuous flows between high-dimensional data $x$ and a low-dimensional embedding $y$, enabling sampling of $p(y|x)$ and $p(x|y)$ while preserving residual information in the flow network. The method combines a kernelized generalized Gromov–Wasserstein OT coupling with a Dual Conditional Flow Matching network that shares a drift model to realize bidirectional conditional transports. Empirical results on MNIST, image datasets, and QM9 show improved embedding structure and higher reconstruction fidelity compared to baselines, illustrating effective semantic control and robust generation under severe compression (e.g., two-dimensional latent spaces). The work advances controllable dimensionality reduction and bidirectional generative modeling by integrating sophisticated OT priors with flow-based transport, offering practical pathways for interpretable embeddings and high-fidelity reconstructions. Despite computational demands, CPFM demonstrates strong potential for applications requiring semantic disentanglement and precise controllable generation.

Abstract

We introduce Coupled Flow Matching (CPFM), a framework that integrates controllable dimensionality reduction and high-fidelity reconstruction. CPFM learns coupled continuous flows for both the high-dimensional data x and the low-dimensional embedding y, which enables sampling p(y|x) via a latent-space flow and p(x|y) via a data-space flow. Unlike classical dimension-reduction methods, where information discarded during compression is often difficult to recover, CPFM preserves the knowledge of residual information within the weights of a flow network. This design provides bespoke controllability: users may decide which semantic factors to retain explicitly in the latent space, while the complementary information remains recoverable through the flow network. Coupled flow matching builds on two components: (i) an extended Gromov-Wasserstein optimal transport objective that establishes a probabilistic correspondence between data and embeddings, and (ii) a dual-conditional flow-matching network that extrapolates the correspondence to the underlying space. Experiments on multiple benchmarks show that CPFM yields semantically rich embeddings and reconstructs data with higher fidelity than existing baselines.

Figures (8)

• Figure 1: An overview of coupled flow matching.
• Figure 2: Based on the generalized GWOT transport plan $\pi$, each source sample $x$ is associated with a probability distribution over candidate embeddings ${y}$. A single embedding $y$ is randomly sampled according to its assigned weight $\pi_i$. The latent embedding distributions $\mu_{\mathcal{Y}}$ considered are: (Left) a Gaussian $\mathcal{N}(0, I_2)$; (Middle) a uniform distribution on the square $[-1,1]^2$; and (Right) a uniform distribution on the unit circle ${, y \in \mathbb{R}^2 : |y|=1 ,}$.
• Figure 3: The MNIST dataset compressed into two dimensions by DCFM, where different classes are distinguished by labels and then reconstructed.
• Figure 4: (Left) t-SNE visualization of the VAE latent space, where the labels do not show a clear separation. (Right) Embedding obtained from generalized GWOT, where the molecular embeddings vary continuously in label.
• Figure 5: The QM9 dataset compressed into two dimensions by DCFM and then reconstructed.

Theorems & Definitions (12)

• Proposition 4.1
• Theorem 4.2
• Proposition 4.3
• Theorem 4.4
• Theorem B.1
• proof
• Proposition B.2
• proof
• Proposition B.3
• proof