Simplifying Flow Matching Transformations with Low-Rank Mixture Models

2026-06-29Machine Learning

Machine Learning
AI summary

The authors explain that normalizing flows are models that transform complicated data into simple forms and back, usually using a standard normal distribution as a starting point. They found that this usual choice can make the transformation harder and slower because the simple form doesn't match the real data well. To fix this, the authors propose using a mixture of probabilistic principal component analyzers (MPPCA) as the starting point, which better fits the data shape. This approach helps the model learn faster and generate better results. They tested their idea on different types of data and saw consistent improvements.

Normalizing flowsLatent distributionStandard normal distributionKL divergenceMixture modelsProbabilistic principal component analysisExpectation-maximization algorithmGenerative modelsHigh-dimensional data
Authors
Liam A. Kruse, Houjun Liu, Alexandros E. Tzikas, Mansur M. Arief, Mykel J. Kochenderfer
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.