Simplifying Flow Matching Transformations with Low-Rank Mixture Models

2026-06-28

AI Summary

This work addresses the topological mismatch between standard normal latent variables and complex data distributions, which hinders the training efficiency and generative performance of normalizing flows. To mitigate this issue, the paper introduces, for the first time, a mixture of probabilistic principal component analyzers (MPPCA) as a learnable low-rank latent prior within the normalizing flow framework. This formulation effectively alleviates topological obstructions, simplifies the flow transformation architecture, and enables efficient initialization. The model is trained end-to-end by integrating the expectation-maximization (EM) algorithm with KL divergence minimization. Empirical evaluations on both tabular and image datasets demonstrate that the proposed approach significantly outperforms baseline methods, achieving faster convergence and superior sample quality.

Abstract

Normalizing flows are powerful generative models that learn an invertible mapping between complex data distributions and simple latent distributions, typically a standard normal density. However, this choice of latent density can impose unnecessary complexity on the learned flow transformation due to the topological mismatch between the latent and data densities, leading to slower training and suboptimal performance. In this work, we propose using mixtures of probabilistic principal component analyzers (MPPCA) as the latent density for normalizing flows. We simplify the learned flow transformation by learning a latent distribution that more closely aligns with the data distribution in terms of KL divergence, thus enabling faster convergence and improved generative performance. Critically, MPPCA models can be fit quickly and cheaply using the expectation-maximization algorithm, making them a practical choice for initializing latent distributions even in high-dimensional generative tasks. We validate our method on both tabular and image datasets, demonstrating consistent gains in training efficiency and generation quality compared to baselines.

Problem

• normalizing flows
• latent distribution
• topological mismatch
• generative models
• KL divergence

Innovation

• normalizing flows
• MPPCA
• latent distribution
• KL divergence
• expectation-maximization