Gaussian mixture models with Wasserstein distance

1 Introduction

• VAEs offer a stable lower bound on log likelihood and an encoding distribution, whereas GANs support sharp deterministic samples but have less stable training.The paper notes blurry VAE samples associated with the objective and tractability requiring simple priors.Source: §1 · p. 1
• Optimal Transport offers a general framework for generative-model training, though the approach was still nascent despite promising early results.Source: §1 · p. 1
• The paper seeks to improve generative modeling with discrete and continuous latents and show cases where OT training is more effective than traditional VAE training.Source: §1 · p. 1
• Discrete latent-variable models matter because discreteness is common in the natural world and its datasets, but they are harder to train than continuous counterparts.Prior approaches include mitigating high-variance samples, parameterizing discrete distributions through continuous ones, and designing models to leverage conjugacy.Source: §1 · p. 1
• Earlier work found an unsupervised MNIST Gaussian-mixture model difficult to train with variational inference without substantial changes to the VAE objective [5].The discrete latent variational distribution can collapse, while labeled samples allowed learning in a semi-supervised version [16].Source: §1 · pp. 1–2
• Optimal Transport training enables Gaussian-mixture latent-variable models to train unsupervised on MNIST, supporting easier convergence than the VAE approach.The authors motivate this with the weaker topology induced by Wasserstein distance.Source: §1 · p. 2

2 Gaussian-Mixture Wasserstein Autoencoders

• The GM-LVM places discrete k above continuous z and is called GM-VAE under VAE training or GM-WAE under WAE training.Source: §2 · p. 2

2.1 The difficulty of training GM-VAEs

• In GM-VAEs, the discrete variational distribution can quickly form meaningless, unbalanced classes to accommodate an untrained generator, which then learns around them.Source: §2.1 · p. 2
• Within 20 steps, reconstruction loss drops as qD shuts off four k values; later continuous-latent learning improves reconstructions by roughly 1,000 steps.Source: §2.1 · pp. 2–3, Figure 1 · p. 3
• The shut-off k values do not return as the generator learns, because it would incur reconstruction penalties for z values associated with them.Source: §2.1 · p. 3
• The convergence difficulty is demonstrated for a specific implementation, and its root cause is not established generally; the authors expect a partial link to per-data-point ELBO KL effects.Source: §2.1 · p. 3

2.2 Optimal Transport facilitates training of GM-LVMs

• The constrained Wasserstein objective is relaxed with MMD and recovers the original distance as λ→∞.Source: §2.2 · p. 4
• The WAE uses deterministic gθ(z); its parameterized encoder uses diagonal-Gaussian qC for reparameterization and sums over discrete qD exactly for small MNIST K.Source: §2.2 · p. 4
• Unlike the per-data-point ELBO, the Wasserstein formulation aggregates q(z,k|x) under pdata, which the authors expect to let qD retain data-specific information while matching the prior overall.Source: §2.2 · p. 5

3 Results

• The MNIST model uses ten Gaussian prior mixtures for ten digits, a uniform mixture prior, and continuous latent dimension 10.The prior modes are chosen to be mutually equidistant with approximately 5% overlap.Source: §3 · p. 5

3.1 Reconstructions and samples

• The model reconstructs held-out MNIST digits and generates samples from each discrete latent value using its uniform prior.Figure 2 shows odd rows as test data and the rows beneath as reconstructions; its samples are ordered by discrete latent value.Source: §3.1 · pp. 5–6, Figure 2 · p. 6
• Interpolations show a smooth latent manifold, with nearly every prior mode representing a different digit and enabling control across modes.The paper reports that the learned manifold matches the prior modes; interpolations between data points are similar to other WAE and VAE approaches on MNIST.Source: §3.1 · p. 6, Figure 3 · p. 6
• After a few hundred VAE training steps, reconstructions deteriorate as the continuous KL term outweighs reconstruction loss.The authors attribute the deterioration to optimization sacrificing reconstruction to reduce the continuous KL term.Source: §3.1 · p. 7
• GM-WAE is qualitatively competitive with other approaches while using a low-complexity implementation and providing more control over generation and inference.The discrete-continuous latent structure lets each Gaussian encoder component capture data within one mode.Source: §3 · p. 5, §3.1 · p. 7

3.2 Latent variable fidelity

• The discrete encoder approximately assigns latent values to digit classes, while continuous-latent clusters overlap without fully aligning with prior clusters.The paper notes overlap between the 4 and 9 clusters and corresponding overlap in discrete assignments.Source: §3.2 · pp. 7–8
• 76% held-out digit-class assignment accuracy is achieved by GM-WAE, compared with 50–60% for basic K-means and 90% for.The result uses 16 discrete classes and a substantially different model and training procedure.Source: §3.2 · p. 8

• The Gaussian-mixture latent structure is suited to data with discrete classes and continuous variation within each class.Source: §4 · p. 8
• Optimal Transport, particularly Wasserstein Autoencoders, provides a natural framework for training these models, which are difficult to train with Variational Autoencoders.Source: §4 · p. 8
• MNIST experiments produced promising results and demonstrated additional control over generation.Source: §4 · p. 8