PaperlayerView exact revision on arXiv

Source-linked AI summary

Flows for simultaneous manifold learning and density estimation

Johann Brehmer, Kyle Cranmer

arXiv:2003.13913v3stat.MLcs.LG

TL;DR

The paper addresses generative modeling when data lie near a lower-dimensional manifold, where standard flows represent the ambient space and may misrepresent the manifold. It introduces M-flows, which learn the manifold and a tractable density on it with separate manifold and density updates. Across experiments, M-flows learn higher-quality manifolds and generally outperform ambient-flow and PIE models on downstream inference and many generative metrics.

  • Problem

    Standard flows cannot exactly represent data restricted to a lower-dimensional manifold, motivating models that estimate both the manifold and its density.

  • Method

    M-flows learn an injective, invertible map from a lower-dimensional latent space to the data manifold, retain a tractable density there, and train with separate manifold and density updates.

  • Results

    M-flows learned higher-quality manifolds than PIE baselines and performed better than ambient-flow and PIE models on most downstream inference tasks and many evaluated generative metrics.

  • Takeaways & Limitations

    M-flows provide a generative model that combines tractable manifold density estimation with projection-based dimensionality reduction, denoising, and possible out-of-distribution detection.

  • Takeaways & Limitations

    The discussion assumes the manifold is connected, equivalent to R^n, describable by a single chart, and has known dimensionality n.

Abstract

from arXiv · show

We introduce manifold-learning flows (M-flows), a new class of generative models that simultaneously learn the data manifold as well as a tractable probability density on that manifold. Combining aspects of normalizing flows, GANs, autoencoders, and energy-based models, they have the potential to represent datasets with a manifold structure more faithfully and provide handles on dimensionality reduction, denoising, and out-of-distribution detection. We argue why such models should not be trained by maximum likelihood alone and present a new training algorithm that separates manifold and density updates. In a range of experiments we demonstrate how M-flows learn the data manifold and allow for better inference than standard flows in the ambient data space.

1. Introduction

M-flows address the mismatch between lower-dimensional data manifolds and ambient-space normalizing flows by learning both the manifold and a tractable density on it. They combine a learned manifold chart with projection-based off-manifold handling, while using separate training updates because maximum likelihood alone can fail.

  • Motivation: Standard normalizing flows cannot exactly represent data restricted to a lower-dimensional manifold and instead learn smeared-out ambient densities.Their full-dimensional latent spaces also increase memory and computational costs.
  • M-flow model: M-flows use an injective, invertible map from a lower-dimensional latent space to learn the data manifold and a tractable density over it.The model provides a bijective chart while learning the manifold from data rather than requiring a prescribed chart.
  • M-flow model: For off-manifold inputs, M-flows project onto the manifold and return both the distance to it and the density after projection.This supplies a defined treatment for points that do not lie on the learned manifold.
  • Advantages: The lower-dimensional latent space can reduce model complexity, support more expressive transformations, and scale to higher-dimensional data within a fixed computational budget.The architecture also naturally supports conditional densities on fixed manifolds, which the paper states should improve data efficiency in such settings.
  • Advantages: Projection onto the learned manifold provides dimensionality reduction and denoising, while distance from the manifold may help detect out-of-distribution samples.These uses follow from the model's explicit projection and distance outputs.
  • Training: Maximum likelihood alone may fail to learn both a good manifold and a good density, motivating separate manifold and density updates for efficient training.The paper identifies expensive density evaluation as a second training challenge addressed by its new strategies.

B. Flows on a prescribed manifold.

Flows on a prescribed manifold use a known chart to define a density restricted to the manifold, enabling tractable likelihood evaluation and sampling when the chart is available. The surrounding comparison highlights why learned-manifold approaches are needed when the manifold is unknown or when off-manifold evaluation matters.

  • Flow on a manifold (FOM): A flow on a prescribed manifold uses a known chart or atlas to construct a flow restricted to that manifold.The latent coordinates are modeled with a regular normalizing flow, while the chart maps those coordinates into the data manifold.
  • Flow on a manifold (FOM): Sampling draws from the base density, applies the latent flow, and then maps the result through the prescribed chart.The resulting model is by construction limited to the true manifold.
  • Comparison with other generative models: GANs and VAEs learn lower-dimensional data representations but generally lack tractable likelihoods, while standard flows cannot exactly represent lower-dimensional support.These differences motivate models that combine learned manifold structure with tractable density evaluation.
  • Learned-manifold alternatives: PIE-style models can sample from a learned manifold by fixing off-manifold latent variables, but that sampling density is not the same as their tractable ambient density.The alignment of latent variables with manifold coordinates is an inductive-bias issue, and maximum likelihood may not resolve it reliably.
  • Manifold-learning flow (M-flow): M-flows replace the prescribed chart with a learnable diffeomorphic map while retaining a tractable density defined on the learned manifold.They additionally project arbitrary inputs onto the manifold, returning both reconstruction distance and projected-manifold likelihood.
  • Manifold-learning flow (M-flow): M-flow reconstruction error is zero exactly on the learned manifold, providing a distance signal separate from the density evaluated after projection.This separation supports denoising, anomaly detection, and out-of-distribution detection as distinct uses of the model outputs.

D. Manifolds with unknown dimensionality or nontrivial topology.

The framework initially assumes known manifold dimensionality and a single-chart topology, but it outlines extensions for unknown dimensions and more complex manifolds.

  • The simplifying assumptions are known manifold dimensionality n and topology equivalent to Rn, including representability by a single chart.
  • The presented algorithms can be extended when the assumptions about dimensionality and topology are relaxed.
  • When n is unknown, models can be trained across candidate dimensions, using reconstruction error and downstream performance rather than incomparable likelihoods.A performance drop is expected when the learned manifold is smaller than the true manifold dimension.
  • For PIE, trainable base-density variances with regularization toward 0 or 1 can allow the model to learn manifold dimensionality from data.
  • Multiple disjoint manifold pieces, potentially with different dimensionalities, may be handled with a mixture model using separate latent-to-data transformations.The text leaves the trainability of such a model unresolved.

3. Efficient training and evaluation

The paper frames M-flow training and evaluation as requiring strategies beyond ordinary maximum-likelihood procedures, especially because implicit generative models lack directly available likelihoods.

  • Implicit generative models use alternatives such as adversarial or optimal-transport objectives for GANs and lower-bound training for VAEs when likelihood evaluation is unavailable.

A. Maximum likelihood is not enough.

Naive maximum likelihood is problematic because the learned manifold changes the projected data variable and can favor pathological configurations. The paper therefore separates manifold and density updates, while noting substantial density-evaluation costs.

  • Comparing naive likelihoods across manifold transformations is meaningless because the likelihood’s data variable depends on the manifold-defining parameters.
  • A manifold perpendicular to the true data manifold can achieve high naive likelihood by projecting points into a high-density region.
  • In the toy example, reconstruction error is minimized at the true manifold and density configuration.
  • A combined reconstruction-error and negative-likelihood loss retains a lower unbounded pathological minimum, so outcomes depend on initialization and hyperparameters.
  • Separating manifold parameters from density parameters avoids gradient flow toward the pathological configuration and can recover both in the idealized infinite-capacity limit.
  • M-flow density evaluation requires full Jacobian entries rather than only layerwise log determinants, making it more expensive than standard-flow evaluation.
  • Because evaluation cost grows with ambient dimension, maximizing the full M-flow likelihood may become prohibitively expensive during high-dimensional training.
  • During manifold updates the likelihood is unnecessary, and during density updates the slow Jacobian term does not affect gradients and can be omitted.

B. Training strategies.

The paper separates manifold and density updates because naive maximum-likelihood training can fail, and compares sequential, alternating, adversarial, and regularized strategies. Its preferred M/D strategy trains the manifold and density in distinct phases, while other approaches introduce computational or optimization limitations.

  • Motivation: Naive maximum-likelihood training may fail to learn both a good manifold and a good density.The paper identifies this as a central training challenge.
  • Alternative strategies: Simultaneous training minimizes a combined loss containing negative naive log likelihood and reconstruction error.The same objective updates both the manifold-defining and density-defining transformations.
  • Separate manifold and density training: M/D training separates manifold updates from density updates to improve training speed and robustness.The manifold phase updates f using reconstruction error, while the density phase updates h.
  • Scheduling: Sequential training learns the manifold first and then trains the density, whereas alternating training switches between the two phases after fixed gradient updates.Algorithm 1 implements the alternating version with separate manifold and density update loops.
  • Alternative strategies: Adversarial and optimal-transport variants train generated samples against data, while OT/D alternates adversarial phases with likelihood-based density phases.The Sinkhorn divergence interpolates between Wasserstein and energy distances as its regularization changes.
  • Regularization: Geometric implicit regularization may yield high-quality manifolds but requires expensive Jacobian evaluations and is left for future work.The authors expect this strategy to perform poorly because pathological low-loss solutions can have poor manifold and density quality.

C. Likelihood evaluation.

M-flows provide tractable manifold likelihoods, but exact evaluation can be computationally expensive. The paper proposes conditional parameterization and approximate or amortized alternatives to reduce this cost.

  • Exact likelihood: Exact M-flow likelihood evaluation becomes more expensive with data dimension and transformation complexity.The paper reports that this is not limiting for low- to medium-dimensional inference problems.
  • Parameter inference: When the data manifold is independent of θ, only the density transformation h needs to depend on θ.This structure allows parameter inference without recomputing the expensive manifold Jacobian terms.
  • Parameter inference: Likelihood-ratio Jacobian terms cancel when comparing parameter points with a θ-independent manifold.The same cancellation applies to MCMC acceptance probabilities.
  • Approximate evaluation: Approximate likelihood methods replace full Jacobian calculations with randomly sampled matrix-vector products, trading precision against speed.The paper leaves empirical evaluation of this trade-off for future work.
  • Approximate evaluation: A stochastic estimator of a likelihood lower bound may be efficient, but its application to M-flows is left unexplored.The lower bound itself remains computationally expensive.
  • Amortization: Jacobian evaluation can potentially be amortized by regressing on j(u)=log det[J_g^T(u)J_g(u)].The proposed regression approach is presented as future work.

4. Experiments

The experiments progress from pedagogical toy examples to increasingly realistic settings, including the Lorenz attractor, particle physics data, and image datasets. Because projected likelihoods differ across M-flow variants, the paper evaluates models primarily through generated samples rather than direct test likelihood comparison.

  • Experimental scope: The experimental suite includes two pedagogical examples, the Lorenz attractor, a particle physics dataset, and image datasets.These experiments are organized across Sections 4.A–4.F.
  • Evaluation approach: Direct test log-likelihood comparison is not meaningful because M-flow variants evaluate likelihood after projecting onto different learned manifolds.The resulting likelihoods may use different data variables and units.
  • Evaluation approach: The experiments therefore assess performance through the quality of generated samples.This evaluation focuses on the generative mode rather than directly comparing projected likelihood values.

A. Gaussian on a circle.

The experiments show that M-flows can learn manifold-supported densities and outperform ambient-flow baselines on manifold quality and parameter inference. Their performance depends on data coverage and the training strategy.

  • Gaussian on a circle: On a noisy circle, ambient flows learn smeared-out densities, whereas M-flows restrict density support to the learned manifold.The FOM model uses the known manifold, while M-flow (M/D) learns the manifold from data.
  • Gaussian on a circle: M-flow (M/D) achieves good manifold quality and density estimation where training data are concentrated, but its learned manifold departs from truth in sparsely sampled regions.This demonstrates sensitivity to the amount of local training data.
  • Mixture model on a polynomial surface: On the polynomial-surface example, M-flow (M/D) most closely matches the true manifold, density, and parameter dependence across θ.The comparison includes θ=0 and changes at θ=±1.
  • Mixture model on a polynomial surface: M-flows exactly learn a two-dimensional manifold, while ambient-flow density remains supported off the manifold despite concentrating near it.The cross-section comparison visualizes this distinction through log likelihood and off-support regions.
  • Quantitative evaluation: Across the evaluated metrics, manifold-learning flows generally perform best, including better generated-sample and reconstruction quality than AF and PIE.They are competitive with PAE on these metrics despite retaining a tractable likelihood.
  • Quantitative evaluation: M-flow variants clearly outperform AF and PIE for inference on θ, but are weaker for out-of-distribution detection than AF log likelihood and PAE reconstruction error.The OT variants perform poorly across metrics, while M-flow (S) is slower and slightly worse.

C. Lorenz attractor.

The Lorenz experiment tests whether an M-flow can learn both a chaotic attractor’s lower-dimensional manifold and its invariant density. The learned model captures the attractor’s overall shape and a plausible density, while retaining imperfections associated with topology.

  • The Lorenz system has chaotic solutions that tend toward an attractor with approximate Hausdorff dimension 2.06 and an ergodic invariant probability measure.
  • The experiment trains an M-flow on samples from 100 simulated trajectories to learn a two-dimensional manifold and probability density.Trajectories are evolved from t = 0 to t = 1000, with samples taken after a warm-up period.
  • The learned model captures the attractor’s overall shape, including two disconnected manifold branches, and produces a plausible probability density on it.
  • The learned result is imperfect, with artifacts likely linked to the attractor’s nontrivial topology.

E. StyleGAN image manifolds.

On StyleGAN image manifolds, M-flow and Me-flow learn smooth manifolds that closely match the ground truth and support faithful projections. They also outperform AF and PIE baselines across the reported evaluation metrics, while PIE suffers mode collapse when sampling from its learned manifold.

  • Experimental setup: The experiments train models on 10 000 images for the two-dimensional manifold and 20 000 images for the 64-dimensional manifold using rational-quadratic neural spline flows.The architectures use multi-scale coupling-layer designs, with 20 layers for M-flow, Me-flow, and PIE and 28 for AF.
  • Generated samples: PIE suffers mode collapse when restricted to manifold sampling, especially for n = 2, where it repeatedly generates the same face.The learned PIE level set v = 0 is not aligned with the true image manifold.
  • Learned manifolds: M-flow and Me-flow learn manifolds covering virtually identical images to the ground truth, with smooth changes along manifold coordinates.The PIE models instead assign most image-distribution variance to off-manifold coordinates.
  • Learned manifolds: M-flow and Me-flow projections reconstruct test images faithfully, with only small differences visible in fine details such as hair.The PIE manifolds are not very helpful for projection.
  • Evaluation: M-flow and Me-flow clearly outperform AF and PIE on all reported FID and manifold-projection reconstruction metrics.Table 4 evaluates sample quality with FID and projection quality with mean reconstruction error across independently initialized runs.

F. Real-world images.

On CelebA-HQ, M-flow and Me-flow learn higher-quality manifolds than PIE, but AF produces the most realistic images. The authors attribute this gap partly to the proof-of-principle choice of manifold dimension and note that performance appears to stabilize for n ≳100.

  • Experimental setup: CelebA-HQ experiments use 64 × 64 images and set the M-flow and Me-flow manifold dimensionality to n = 512.The existence and dimensionality of the real-world image manifold are left undetermined for future work.
  • CelebA-HQ results: M-flow and Me-flow learn a higher-quality manifold than PIE on CelebA-HQ, while AF produces the most realistic images.The comparison uses generated-sample FID and reconstruction errors from projection to the learned manifold.
  • Interpretation: The authors suggest that AF’s stronger image realism may reflect a suboptimal manifold dimension in this proof-of-principle study.The CelebA-HQ manifold dimension was not systematically determined.
  • Manifold dimension: Illustrative CelebA results hint that M-flow performance stabilizes with n ≳100.Figure 14 reports one run per dimension after limited training, so the result is explicitly illustrative.
  • Contributions: The paper introduces M-flows and Me-flows, separates manifold and density updates with M/D training, and demonstrates them across experiments.The broader contribution includes a discussion of their relation to several generative-model families and a conditional PIE version.

A. Broader impact

The paper frames manifold-learning flows as potentially useful for simulation-based inference and other scientific settings with manifold-structured data. It also emphasizes risks from limited diversity, inherited bias, and possible misuse for generating fake data.

  • Potential benefits: Manifold-learning flows may help scientists extract knowledge from large-scale experiments by modeling manifold-structured data in simulation-based inference settings.Normalizing-flow surrogates can support inference when scientific simulations lack tractable likelihoods.
  • Risks and safeguards: Manifold-learning flows could be abused to generate fake data masquerading as real documents.The authors note that current models remain far from generating realistic high-resolution images, videos, or audio, while urging long-term attention to this risk.
  • Risks and safeguards: Face models trained in the paper lack diversity and reproduce or reinforce biases present in their training data.The authors state that such biases must be understood, measured, and mitigated before real-world use.

C. Particle physics.

The particle-physics section evaluates M-flows and comparison models on synthetic datasets, using flow-based architectures and sequential manifold/density training. Evaluation includes likelihood-based inference with MCMC and kernel-density estimation, alongside image-generation metrics on manifold datasets.

  • Particle-physics data: The particle-physics data use 40-dimensional summary-statistic vectors generated with MadGraph5_aMC, Pythia8, and Delphes 3, with two effective-field-theory parameters drawn from a unit Gaussian prior.The statistics include energy, momentum, pseudorapidity, invariant mass, and related reconstructed quantities.
  • Architectures: AF, PIE, M-flow, and Me-flow models use rational-quadratic neural spline flows with coupling layers and invertible linear transformations.The particle-physics setup uses 35 coupling layers for AF and 20 layers for f and g plus 15 layers for h in the other flow models.
  • Training: The M-flow and Me-flow models are trained with a sequential M/D algorithm, while likelihood phases use negative mean log likelihood in bits per dimension.The supplied training description also specifies AdamW optimization, cosine annealing, and weight decay of 10^-5.
  • Inference evaluation: Inference uses Metropolis-Hastings MCMC on three synthetic observed datasets, each containing 15 independent samples generated at three parameter settings.For each model and observed dataset, four MCMC chains of length 750 are generated.
  • Image-manifold data: The image-manifold experiments use StyleGAN2-generated datasets with intrinsic dimensions n = 2 and n = 64 embedded in 64×64×3-dimensional image space.The images are downsampled to 64 × 64, and the n = 2 dataset contains 10^4 training images while the n = 64 dataset contains 2·10^4 training images.
  • Image evaluation: Generated images are evaluated with Fréchet Inception Distance, while the 64-dimensional dataset also supports parameter inference using an MCMC chain and kernel-density estimation.The image-manifold models use rational-quadratic neural spline flows, and the CelebA-HQ experiment uses manifold dimension n = 512.
Loading 2003.13913v3…