Table of Contents
Fetching ...

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.

Manifold Learning by Mixture Models of VAEs for Inverse Problems

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.
Paper Structure (38 sections, 6 theorems, 69 equations, 9 figures, 3 algorithms)

This paper contains 38 sections, 6 theorems, 69 equations, 9 figures, 3 algorithms.

Key Result

Lemma 6

Let $x\in\mathcal{M}$, $U_x\subseteq\mathbb{R}^n$ be a neighborhood of $x$ in $\mathbb{R}^n$, $\pi\colon U_x\to\mathcal{M}\cap U_x$ be a differentiable map such that $\pi\circ\pi=\pi$. Set $V_x=\{h\in T_x\mathcal{M}\subseteq\mathbb{R}^n: x+h\in U_x\}$. Then defines a retraction in $x$.

Figures (9)

  • Figure 1: Example of a one-dimensional manifold that admits no global parameterization.
  • Figure 2: Plot of the unnormalized latent density $q$.
  • Figure 3: Data sets used for the different manifolds.
  • Figure 4: Learned charts for the different manifolds. For the manifolds "two circles" and "ring", each color represents one chart. For the manifolds "sphere", "swiss roll" and "torus" we plot each chart in a separate figure.
  • Figure 5: Generated samples by the learned mixture of VAEs. The color of a point indicates from which generator the point was sampled.
  • ...and 4 more figures

Theorems & Definitions (15)

  • Remark 1: Lipschitz regularization
  • Remark 2: Number of charts $K$
  • Remark 3: Projection onto learned charts
  • Example 1
  • Remark 4
  • Definition 5
  • Lemma 6
  • Lemma 7
  • Remark 8: Descent algorithm
  • Remark 9: Dimension of the Manifold
  • ...and 5 more