Structured Coupling for Flow Matching

Abstract: Standard flow matching scales well but typically relies on an unstructured source distribution, limiting its ability to learn interpretable latent structure. Latent-variable models, by contrast, capture structure but often sacrifice generative quality. We bridge this gap by proposing Structured Coupling for Flow Matching (SCFM), a cooperative framework that augments flow matching with structured latent representation learning. By introducing structured latent variables and exogenous noise into the source, SCFM jointly learns a structured prior (via latent variable modeling) and a continuous transport map (via flow matching). It uses a shared time-dependent recognition network for both latent variable model variational inference and intermediate-time flow velocity estimation. This yields a structurally informed yet unconditional, simulation-free flow model, where the latent variable model can also assist flow sampling. Empirically, SCFM facilitates unsupervised latent representation learning for clustering, disentanglement and downstream tasks, while remaining competitive with flow matching in sample quality, showing that meaningful structure can be learned without sacrificing generative fidelity.

• The paper introduces a unified framework (SCFM) that incorporates structured latent priors into flow matching to enhance generative quality and representation interpretability.
• The methodology redesigns source coupling to use a tuple (z, ε) and a shared recognition network for consistent variational inference and flow velocity estimation.
• Empirical evaluations on datasets like MNIST, CIFAR-10, and ImageNet show improved clustering, disentanglement, and sample quality with modest computational overhead.

Structured Coupling for Flow Matching: A Unified Framework for Latent-Variable Representation Learning and High-Fidelity Generation

Introduction

The study introduces Structured Coupling for Flow Matching (SCFM), a methodological advance designed to unify the advantages of flow-based generative models with structured latent-variable modeling. Classical flow matching approaches, despite their strong sample quality and scalable training, rely on simple, unstructured source distributions (commonly isotropic Gaussians), which impedes interpretable representation learning. Conversely, latent-variable models such as VAEs encode semantically rich latent spaces, benefitting clustering and disentanglement, but typically fall short in generative fidelity due to decoder constraints and variational inference approximation errors. SCFM addresses this dichotomy by embedding a learnable, structured latent prior within the flow matching paradigm, thereby supporting unsupervised latent representation learning without sacrificing sample quality.

SCFM: Methodology

Structured Source Coupling

SCFM reformulates the flow matching source endpoint from an unstructured noise variable $x_0$ to a tuple $(z, \epsilon)$, where $z$ is a learnable structural latent and $\epsilon$ is exogenous noise. The structural prior $p_\psi(z)$ is parameterized (e.g., as a GMM), and $p(\epsilon)$ is standard Gaussian. The encoder-induced coupling constructs joint samples $(z, \epsilon, x_1)$ through $$x_1 \sim p_{data}(x_1), z \sim q_\phi(z|x_1), \epsilon \sim p(\epsilon)$$. This coupling both encodes latent structure during training and regularizes the transport map to meaningfully separate semantic factors from residual variability.

Time-Dependent Shared Recognition Network

A single shared neural network family, parameterized as $\mu_\theta(x, t)$, facilitates both variational inference at the endpoint ($t=1$) and posterior mean estimation for flow velocity at intermediate times ($(z, \epsilon)$0). At $(z, \epsilon)$1, $(z, \epsilon)$2 yields VAE encoder outputs for $(z, \epsilon)$3. For $(z, \epsilon)$4, $(z, \epsilon)$5 predicts the posterior mean for the regression-based flow-matching loss. This parameter sharing aligns prior learning with flow transport, ensuring the latent structure directly informs the generative process.

Joint Training Objective

Training combines:

• Variational Flow Matching (VFM) loss for $(z, \epsilon)$6, where regression targets the expected source endpoint under the induced posterior;
• VAE objective at $(z, \epsilon)$7, enforcing KL alignment of the aggregated posterior with the prior and reconstruction through the decoder;
• Regularization of exogenous variables to retain Gaussian structure in $(z, \epsilon)$8.

This cooperative objective ensures that the aggregated posterior over $(z, \epsilon)$9 matches $z$0, and that the flow samples from the structured source are transported to data likelihood-indistinguishable outputs.

Sampling and Decoder-Initialized Refinement

Sampling in SCFM supports:

• ODE-based sampling: Samples $z$1 are mapped to data via ODE integration with the learned velocity field.
• Conditional reconstruction: Given an observation $z$2, $z$3 allows for data-conditioned reconstructions.
• Decoder-initialized refinement: Decoding $z$4 provides a high-level sample, which is subsequently refined by partial ODE integration. This yields rapid, high-fidelity samples at lower compute cost, contingent on decoder quality.

Experimental Evaluation

Representation Learning and Disentanglement

SCFM demonstrates high efficacy in unsupervised clustering and disentanglement benchmarks:

• On MNIST, SCFM outperforms VaDE and MFCVAE on both NMI and clustering accuracy, improving NMI by $z$58 points and ACC by $z$613 points over VaDE.
• On Cars3D and Shapes3D, SCFM achieves FactorVAE and DCI disentanglement scores on par with or exceeding diffusion or VAE-based models. For instance, in Cars3D, SCFM with a $z$7-TCVAE endpoint yields the highest FactorVAE score (0.977), indicating well-separated, controllable latent factors.
• Qualitative factor swaps validate that the latent codes correspond to semantic generative factors (e.g., shape, color, pose).

Latent Probing in Large-Scale Settings

• On CIFAR-10, SCFM achieves the highest linear probe accuracy and competitive nonlinear probe accuracy, outperforming VAE, AAE, and BiGAN.
• For ImageNet-128, latent representations trained by SCFM permit higher Top-1/Top-5 accuracy for both linear and nonlinear probes than strong VAE baselines.
• Visualization of mixture-component samples demonstrates that the structured prior organizes data into clusters with coherent appearance statistics, even without label supervision.

Generative Quality

• On CIFAR-10, SCFM essentially matches the sample quality of flow-matching baselines (FID 2.117 vs. 2.137).
• On unconditional ImageNet-128, SCFM (FID 17.180) outperforms unconditional and label-conditional SiT-XL/2 baselines, showing that structured coupling scales to high-complexity datasets.
• Decoder-initialized refinement achieves a trade-off between generation quality and computational efficiency, rapidly approaching full-flow FID with a fraction of the FLOPs.

Complexity and Cost

SCFM introduces moderate compute and parameter overhead compared to plain flow matching, attributed to the additional decoder and structured latent-branch (e.g., 101M vs. 73.6M parameters on CIFAR-10). The benefits in representation quality, controllability, and interpretability, however, are substantial.

Theoretical and Practical Implications

SCFM extends the expressiveness of simulation-free flow models, showing that explicit latent-structure learning is compatible with high-fidelity generation in continuous flow-based architectures. It avoids introducing conditionality at sampling, contrasting with classifier-free guidance which requires external signals.

The theoretically motivated endpoint alignment ensures aggregated posterior and prior consistency, and the method is compatible with established advances in stochastic interpolant flows and structured priors. SCFM brings flow matching into the regime of representation learning, disentanglement, and clustering, where previous flow-based approaches were largely outperformed by VAEs.

From a practical standpoint, SCFM's ability to learn controllable, semantically meaningful latents enables downstream tasks such as unsupervised classification, factor manipulation, and semantic image editing, all with generative quality competitive with state-of-the-art sample-based flows and diffusions.

Future Directions

Potential avenues for extension include:

• Reduction of training and inference overhead through architectural/algorithmic optimizations.
• Improvements in decoder expressivity to avoid posterior collapse and stabilize the endpoint variational model.
• Application to multimodal, temporally-extended, or conditional generative settings, leveraging the implicit controllability and interpretability of the structured latent space.

Conclusion

SCFM constitutes a significant advance in the integration of structured latent-variable models with transport-based generative flows. It demonstrates empirically and theoretically that a learnable, semantically organized source prior enables clustering, disentanglement, and interpretable representation learning without compromising generative sample quality or simulation-free training. This positions flow matching as a viable candidate for unified generative and representation learning in high-dimensional, complex data regimes (2605.07676).