Manifold Learning by Mixture Models of VAEs for Inverse Problems
Giovanni S. Alberti, Johannes Hertrich, Matteo Santacesaria, Silvia Sciutto
TL;DR
This work proposes learning data manifolds with arbitrary topology by representing them as an atlas of charts, each modeled by a variational autoencoder enhanced with a normalizing flow to flexibly shape the latent prior. The decoders/encoders and flows yield analytic inverses for charts, enabling a Riemannian gradient descent on the learned manifold to solve inverse problems such as deblurring and electrical impedance tomography. The authors derive a tractable ELBO-based loss for training the mixture, analyze chart overlap, and introduce retractions for gradient steps that respect the manifold structure. Empirically, using multiple charts improves manifold coverage and reconstruction quality compared to a single generator, especially in ill-posed inverse problems, and the approach is extensible to Bayesian settings and diffusion-model representations.
Abstract
Representing a manifold of very high-dimensional data with generative models has been shown to be computationally efficient in practice. However, this requires that the data manifold admits a global parameterization. In order to represent manifolds of arbitrary topology, we propose to learn a mixture model of variational autoencoders. Here, every encoder-decoder pair represents one chart of a manifold. We propose a loss function for maximum likelihood estimation of the model weights and choose an architecture that provides us the analytical expression of the charts and of their inverses. Once the manifold is learned, we use it for solving inverse problems by minimizing a data fidelity term restricted to the learned manifold. To solve the arising minimization problem we propose a Riemannian gradient descent algorithm on the learned manifold. We demonstrate the performance of our method for low-dimensional toy examples as well as for deblurring and electrical impedance tomography on certain image manifolds.
