# Posterior Collapse and
Latent Variable Non-identifiability

Yixin Wang
University of Michigan
<EMAIL_ADDRESS>
&David M. Blei
Columbia University
<EMAIL_ADDRESS>
&John P. Cunningham
Columbia University
<EMAIL_ADDRESS>

###### Abstract

Variational autoencoders model high-dimensional data by positing low-
dimensional latent variables that are mapped through a flexible distribution
parametrized by a neural network. Unfortunately, variational autoencoders
often suffer from posterior collapse: the posterior of the latent variables is
equal to its prior, rendering the variational autoencoder useless as a means
to produce meaningful representations. Existing approaches to posterior
collapse often attribute it to the use of neural networks or optimization
issues due to variational approximation. In this paper, we consider posterior
collapse as a problem of latent variable non-identifiability. We prove that
the posterior collapses if and only if the latent variables are non-
identifiable in the generative model. This fact implies that posterior
collapse is not a phenomenon specific to the use of flexible distributions or
approximate inference. Rather, it can occur in classical probabilistic models
even with exact inference, which we also demonstrate. Based on these results,
we propose a class of latent-identifiable variational autoencoders, deep
generative models which enforce identifiability without sacrificing
flexibility. This model class resolves the problem of latent variable non-
identifiability by leveraging bijective Brenier maps and parameterizing them
with input convex neural networks, without special variational inference
objectives or optimization tricks. Across synthetic and real datasets, latent-
identifiable variational autoencoders outperform existing methods in
mitigating posterior collapse and providing meaningful representations of the
data.

##  1 Introduction

Variational autoencoders (vae) are powerful generative models for high-
dimensional data [46, 26]. Their key idea is to combine the inference
principles of probabilistic modeling with the flexibility of neural networks.
In a vae, each datapoint is independently generated by a low-dimensional
latent variable drawn from a prior, then mapped to a flexible distribution
parametrized by a neural network.

Unfortunately, vae often suffer from posterior collapse, an important and
widely studied phenomenon where the posterior of the latent variables is equal
to prior [6, 8, 62, 38]. This phenomenon is also known as latent variable
collapse, KL vanishing, and over-pruning. Posterior collapse renders the vae
useless to produce meaningful representations, in so much as its per-datapoint
latent variables all have the exact same posterior.

Posterior collapse is commonly observed in the vae whose generative model is
highly flexible, leading to the common speculation that posterior collapse
occurs because vae involve flexible neural networks in the generative model
[11], or because it uses variational inference [59]. Based on these
hypotheses, many of the proposed strategies for mitigating posterior collapse
thus focus on modifying the variational inference objective (e.g. [44]),
designing special optimization schemes for variational inference in vae (e.g.
[32, 25, 18]), or limiting the capacity of the generative model (e.g. [60, 6,
16].)

In this paper, we consider posterior collapse as a problem of latent variable
non-identifiability. We prove that posterior collapse occurs if and only if
the latent variable is non-identifiable in the generative model, which loosely
means the likelihood function does not depend on the latent variable [42, 56,
40]. Below, we formally establish this equivalence by appealing to recent
results in Bayesian non-identifiability [49, 43, 42, 58, 40].

More broadly, the relationship between posterior collapse and latent variable
non-identifiability implies that posterior collapse is not a phenomenon
specific to the use of neural networks or variational inference. Rather, it
can also occur in classical probabilistic models fitted with exact inference
methods, such as Gaussian mixture models and probabilistic principal component
analysis (ppca). This relationship also leads to a new perspective on existing
methods for avoiding posterior collapse, such as the delta-VAE [44] or the
$\beta$-VAE [19]. These methods heuristically adjust the approximate inference
procedure embedded in the optimization of the model parameters. Though
originally motivated by the goal of patching the variational objective, the
results here suggest that these adjustments are useful because they help avoid
parameters at which the latent variable is non-identifiable and, consequently,
avoid posterior collapse.

The relationship between posterior collapse and non-identifiability points to
a direct solution to the problem: we must make the latent variable
identifiable. To this end, we propose latent-identifiable vae, a class of vae
that is as flexible as classical vae while also being identifiable. Latent-
identifiable vae resolves the latent variable non-identifiability by
leveraging Brenier maps [39, 37] and parameterizing them with input-convex
neural networks [2, 35]. Inference on identifiable vae uses the standard
variational inference objective, without special modifications or optimization
tricks. Across synthetic and real datasets, we show that identifiable vae
mitigates posterior collapse without sacrificing fidelity to the data.

Related work. Existing approaches to avoiding posterior collapse often modify
the variational inference objective, design new initialization or optimization
schemes for vae, or add neural network links between each data point and their
latent variables [6, 21, 52, 28, 8, 62, 61, 1, 15, 3, 32, 50, 63, 17, 12, 18,
25, 44, 51, 55, 38, 16, 34]. Several recent papers also attempt to provide
explanations for posterior collapse. Chen et al. [8] explains how the inexact
variational approximation can lead to inefficiency of coding in VAE, which
could lead to posterior collapse due to a form of information preference. Dai
et al. [11] argues that posterior collapse can be partially attributed to the
local optima in training vae with deep neural networks. Lucas et al. [33]
shows that posterior collapse is not specific to the variational inference
training objective; absent a variational approximation, the log marginal
likelihood of ppca has bad local optima that can lead to posterior collapse.
Yacoby et al. [59] discusses how variational approximation can select an
undesirable generative model when the generative model parameters are non-
identifiable. In contrast to these works, we consider posterior collapse
solely as a problem of latent variable non-identifiability, and not of
optimization, variational approximations, or neural networks per se. We use
this result to propose the identifiable vae as a way to directly avoid
posterior collapse.

Outside vae, latent variable identifiability in probabilistic models has long
been studied in the statistics literature [49, 43, 42, 58, 40, 56, 42]. More
recently, Betancourt [5] studies the effect of latent variable identifiability
on Bayesian computation for Gaussian mixtures. Khemakhem et al. [23, 24]
propose to resolve the non-identifiability in deep generative models by
appealing to auxiliary data. Kumar and Poole [29] study how the variational
family can help resolve the non-identifiability of vae. These works address
the identifiability issue for a different goal: they develop identifiability
conditions for different subsets of vae, aiming for recovering true causal
factors of the data and improving disentanglement or out-of-distribution
generalization. Related to these papers, we demonstrate posterior collapse as
an additional way that the concept of identifiability, though classical, can
be instrumental in modern probabilistic modeling. Considering identifiability
leads to new solutions to posterior collapse.

Contributions. We prove that posterior collapse occurs if and only if the
latent variable in the generative model is non-identifiable. We then propose
latent-identifiable vae, a class of vae that are as flexible as classical vae
but have latent variables that are provably identifiable. Across synthetic and
real datasets, we demonstrate that latent-identifiable vae mitigates posterior
collapse without modifying vae objectives or applying special optimization
tricks.

##  2 Posterior collapse and latent variable non-identifiability

Consider a dataset $\mathbf{\boldsymbol{x}}=(x_{1},\ldots,x_{n})$; each
datapoint is $m$-dimensional. Positing $n$ latent variables
$\mathbf{\boldsymbol{z}}=(z_{1},\ldots,z_{n})$, a variational autoencoder
(vae) assumes that each datapoint $x_{i}$ is generated by a $K$-dimensional
latent variable $z_{i}$:

$\displaystyle z_{i}$ $\displaystyle\sim p(z_{i}),\qquad x_{i}\,|\,z_{i}\sim
p(x_{i}\,|\,z_{i}\,;\,\theta)=\mathrm{EF}(x_{i}\,|\,f_{\theta}(z_{i})),$ (1)

where $x_{i}$ follows an exponential family distribution with parameters
$f_{\theta}(z_{i})$; $f_{\theta}$ parameterizes the conditional likelihood. In
a deep generative model $f_{\theta}$ is a parameterized neural network.
Classical probabilistic models like Gaussian mixture model [45] and
probabilistic PCA [54, 10, 48, 47] are also special cases of Eq. 1.

To fit the model, vae optimizes the parameters $\theta$ by maximizing a
variational approximation of the log marginal likelihood. After finding an
optimal $\hat{\theta}$, we can form a representation of the data using the
approximate posterior $q_{\hat{\phi}}(z\,|\,x)$ with variational parameters
$\hat{\phi}$ or its expectation
$\mathbb{E}_{q_{\hat{\phi}}(z\,|\,x)}\left[z\,|\,x\right]$.

Note that here we abstract away computational considerations and consider the
ideal case where the variational approximation is exact. This choice is
sensible: if the exact posterior suffers from posterior collapse then so will
the approximate posterior (a variational approximation cannot “uncollapse” a
collapsed posterior). That said we also note that there exist in practice
situations where variational inference alone can lead to posterior collapse. A
notable example is when the variational approximating family is overly
restrictive: it is then possible to have non-collapsing exact posteriors but
collapsing approximate posteriors.

###  2.1 Posterior collapse $\Leftrightarrow$ Latent variable non-
identifiability

We first define posterior collapse and latent variable non-identifiability,
then proving their connection.

######  Definition 1 (Posterior collapse [6, 8, 62, 38]).

Given a probability model
$p(\mathbf{\boldsymbol{x}},\mathbf{\boldsymbol{z}}\,;\,\theta)$, a parameter
value $\theta=\hat{\theta}$, and a dataset
$\mathbf{\boldsymbol{x}}=(x_{1},\ldots,x_{n})$, the posterior of the latent
variables $\mathbf{\boldsymbol{z}}$ collapses if

$\displaystyle
p(\mathbf{\boldsymbol{z}}\,|\,\mathbf{\boldsymbol{x}}\,;\,\hat{\theta})=p(\mathbf{\boldsymbol{z}}).$
(2)

The posterior collapse phenomenon can occur in a variety of probabilistic
models and with different latent variables. When the probability model is a
vae, it only has local latent variables
$\mathbf{\boldsymbol{z}}=(z_{1},\ldots,z_{n})$, and Eq. 2 is equivalent to the
common definition of posterior collapse
$p(z_{i}\,|\,x_{i}\,;\,\hat{\theta})=p(z_{i})$ for all $i$ [33, 44, 17, 12].
Posterior collapse has also been observed in Gaussian mixture models [5]; the
posterior of the latent mixture weights resembles their prior when the number
of mixture components in the model is larger than that of the data generating
process. Regardless of the model, when posterior collapse occurs, it prevents
the latent variable from providing meaningful summary of the dataset.

######  Definition 2 (Latent variable non-identifiability [42, 56]).

Given a likelihood function
$p(\mathbf{\boldsymbol{x}}\,|\,\mathbf{\boldsymbol{z}}\,;\,\theta)$, a
parameter value $\theta=\hat{\theta}$, and a dataset
$\mathbf{\boldsymbol{x}}=(x_{1},\ldots,x_{n})$, the latent variable
$\mathbf{\boldsymbol{z}}$ is non-identifiable if

$\displaystyle
p(\mathbf{\boldsymbol{x}}\,|\,\mathbf{\boldsymbol{z}}=\tilde{\mathbf{\boldsymbol{z}}}^{\prime}\,;\,\hat{\theta})=p(\mathbf{\boldsymbol{x}}\,|\,\mathbf{\boldsymbol{z}}=\tilde{\mathbf{\boldsymbol{z}}}\,;\,\hat{\theta})\qquad\forall\tilde{\mathbf{\boldsymbol{z}}}^{\prime},\tilde{\mathbf{\boldsymbol{z}}}\in\mathcal{Z},$
(3)

where $\mathcal{Z}$ denotes the domain of $\mathbf{\boldsymbol{z}}$, and
$\tilde{\mathbf{\boldsymbol{z}}}^{\prime},\tilde{\mathbf{\boldsymbol{z}}}$
refer to two arbitrary values the latent variable $\mathbf{\boldsymbol{z}}$
can take. As a consequence, for any prior $p(z)$ on $z$, we have the
conditional likelihood equal to the marginal
$p(\mathbf{\boldsymbol{x}}\,|\,\mathbf{\boldsymbol{z}}=\tilde{\mathbf{\boldsymbol{z}}}\,;\,\hat{\theta})=\int
p(\mathbf{\boldsymbol{x}}\,|\,\mathbf{\boldsymbol{z}}\,;\,\hat{\theta})p(\mathbf{\boldsymbol{z}})\mathop{}\\!\mathrm{d}\mathbf{\boldsymbol{z}}=p(\mathbf{\boldsymbol{x}}\,;\,\hat{\theta})\quad\forall\tilde{\mathbf{\boldsymbol{z}}}\in\mathcal{Z}.$

Definition 2 says a latent variable $\mathbf{\boldsymbol{z}}$ is non-
identifiable when the likelihood of the dataset $\mathbf{\boldsymbol{x}}$ does
not depend on $\mathbf{\boldsymbol{z}}$. It is also known as practical non-
identifiability [42, 56] and is closely related to the definition of
$\mathbf{\boldsymbol{z}}$ being conditionally non-identifiable (or
conditionally uninformative) given $\hat{\theta}$ [49, 43, 42, 58, 40]. To
enforce latent variable identifiability, it is sufficient to ensure that the
likelihood $p(\mathbf{\boldsymbol{x}}\,|\,\mathbf{\boldsymbol{z}},\theta)$ is
an injective (a.k.a. one-to-one) function of $\mathbf{\boldsymbol{z}}$ for all
$\theta$. If this condition holds then

$\displaystyle\tilde{\mathbf{\boldsymbol{z}}}^{\prime}\neq\tilde{\mathbf{\boldsymbol{z}}}\qquad\Rightarrow\qquad
p(\mathbf{\boldsymbol{x}}\,|\,\mathbf{\boldsymbol{z}}=\tilde{\mathbf{\boldsymbol{z}}}^{\prime}\,;\,\hat{\theta})\neq
p(\mathbf{\boldsymbol{x}}\,|\,\mathbf{\boldsymbol{z}}=\tilde{\mathbf{\boldsymbol{z}}}\,;\,\hat{\theta}).$
(4)

Note that latent variable non-identifiability only requires Eq. 3 be true for
a given dataset $\mathbf{\boldsymbol{x}}$ and parameter value $\hat{\theta}$.
Thus a latent variable may be identifiable in a model given one dataset but
not another, and at one $\theta$ but not another. See examples in Appendix A.

Latent variable identifiability (Definition 2) [42, 56] differs from model
identifiability [41], a related notion that has also been cited as a
contributing factor to posterior collapse [59]. Latent variable
identifiability is a weaker requirement: it only requires the latent variable
$\mathbf{\boldsymbol{z}}$ be identifiable at a particular parameter value
$\theta=\hat{\theta}$, while model identifiability requires both
$\mathbf{\boldsymbol{z}}$ and $\theta$ be identifiable.

We now establish the equivalence between posterior collapse and latent
variable non-identifiability.

######  Theorem 1 (Latent variable non-identifiability $\Leftrightarrow$
Posterior collapse).

Consider a probability model
$p(\mathbf{\boldsymbol{x}},\mathbf{\boldsymbol{z}}\,;\,\theta)$, a dataset
$\mathbf{\boldsymbol{x}}$, and a parameter value $\theta=\hat{\theta}$. The
local latent variables $\mathbf{\boldsymbol{z}}$ are non-identifiable at
$\hat{\theta}$ if and only if the posterior of the latent variable
$\mathbf{\boldsymbol{z}}$ collapses,
$p(\mathbf{\boldsymbol{z}}\,|\,\mathbf{\boldsymbol{x}})=p(\mathbf{\boldsymbol{z}})$.

###### Proof.

To prove that non-identifiability implies posterior collapse, note that, by
Bayes rule,

$\displaystyle
p(\mathbf{\boldsymbol{z}}\,|\,\mathbf{\boldsymbol{x}}\,;\,\hat{\theta})\propto
p(\mathbf{\boldsymbol{z}})p(\mathbf{\boldsymbol{x}}\,|\,\mathbf{\boldsymbol{z}}\,;\,\hat{\theta})=p(\mathbf{\boldsymbol{z}})p(\mathbf{\boldsymbol{x}}\,;\,\hat{\theta})\propto
p(\mathbf{\boldsymbol{z}}),$ (5)

where the middle equality is due to the definition of latent variable non-
identifiability. It implies
$p(\mathbf{\boldsymbol{z}}\,|\,\mathbf{\boldsymbol{x}}\,;\,\hat{\theta})=p(\mathbf{\boldsymbol{z}})$
as both are densities. To prove that posterior collapse implies latent
variable non-identifiability, we again invoke Bayes rule. Posterior collapse
implies that
$p(\mathbf{\boldsymbol{z}})=p(\mathbf{\boldsymbol{z}}\,|\,\mathbf{\boldsymbol{x}}\,;\,\hat{\theta})\propto
p(\mathbf{\boldsymbol{z}})\cdot
p(\mathbf{\boldsymbol{x}}\,|\,\mathbf{\boldsymbol{z}}\,;\,\hat{\theta})$,
which further implies that
$p(\mathbf{\boldsymbol{x}}\,|\,\mathbf{\boldsymbol{z}}\,;\,\hat{\theta})$ is
constant in $\mathbf{\boldsymbol{z}}$. If
$p(\mathbf{\boldsymbol{x}}\,|\,\mathbf{\boldsymbol{z}}\,;\,\hat{\theta})$
nontrivially depends on $\mathbf{\boldsymbol{z}}$, then
$p(\mathbf{\boldsymbol{z}})$ must be different from
$p(\mathbf{\boldsymbol{z}})p(\mathbf{\boldsymbol{x}}\,|\,\mathbf{\boldsymbol{z}}\,;\,\hat{\theta})$
as a function of $\mathbf{\boldsymbol{z}}$. ∎

The proof of Theorem 1 is straightforward, but Theorem 1 has an important
implication. It shows that the problem of posterior collapse mainly arises
from the model and the data, rather than from inference or optimization. If
the maximum likelihood parameters $\hat{\theta}$ of the vae renders the latent
variable $z$ non-identifiable, then we will observe posterior collapse.
Theorem 1 also clarifies why posteriors may change from non-collapsed to
collapsed (and back) while fitting a VAE. When fitting a VAE, Some parameter
iterates may lead to posterior collapse; others may not.

Theorem 1 points to why existing approaches can help mitigate posterior
collapse. Consider the $\beta$-vae [19], the vae lagging encoder [18], and the
semi-amortized vae [25]. Though motivated by other perspectives, these methods
modify the optimization objectives or algorithms of vae to avoid parameter
values $\theta$ at which the latent variable is non-identifiable. The
resulting posterior may not collapse, though the optimal parameters for these
algorithms no longer approximates the maximum likelihood estimate.

Theorem 1 can also help us understand posterior collapse observed in practice,
which manifests as the phenomenon that the posterior is approximately (as
opposed to exactly) equal to the prior,
$p(\mathbf{\boldsymbol{z}}\,|\,\mathbf{\boldsymbol{x}}\,;\,\hat{\theta})\approx
p(\mathbf{\boldsymbol{z}})$. In several empirical studies of vae (e.g. [18,
12, 25]), we observe that the Kullback-Leibler (kl) divergence between the
prior and posterior is close to zero but not exactly zero, a property that
stems from the likelihood
$p(\mathbf{\boldsymbol{x}}\,|\,\mathbf{\boldsymbol{z}})$ being nearly constant
in the latents $\mathbf{\boldsymbol{z}}$. In these cases, Theorem 1 provides
the intuition that the latent variable is nearly non-identifiable ,
$p(\mathbf{\boldsymbol{x}}\,|\,\tilde{\mathbf{\boldsymbol{z}}}^{\prime})\approx
p(\mathbf{\boldsymbol{x}}\,|\,\tilde{\mathbf{\boldsymbol{z}}}),\forall\tilde{\mathbf{\boldsymbol{z}}},\tilde{\mathbf{\boldsymbol{z}}}^{\prime}$
and so Eq. 2 holds approximately.

###  2.2 Examples of latent variable non-identifiability and posterior
collapse

We illustrate Theorem 1 with three examples. Here we discuss the example of
Gaussian mixture VAE (gmvae). See Appendix A for probabilistic principal
component analysis (ppca) and Gaussian mixture model (gmm).

The gmvae [51, 13] is the following model:

$\displaystyle p(z_{i})=\mathrm{Categorical}(1/K),\quad
p(w_{i}\,|\,z_{i}\,;\,\mu,\Sigma)=\mathcal{N}(\mu_{z_{i}},\Sigma_{z_{i}}),\quad
p(x_{i}\,|\,w_{i}\,;\,f,\sigma)=\mathcal{N}(f(w_{i}),\sigma^{2}\cdot I_{m}),$

where $\mu_{k}$’s are $d$-dimensional, $\Sigma_{k}$ are $d\times
d$-dimensional, and the parameters are $\theta=(\mu,\Sigma,f,\sigma^{2})$.
Suppose the function $f$ is fully flexible; thus $f(w_{i})$ can capture any
distribution of the data. The latent variable of interest is the categorical
$\mathbf{\boldsymbol{z}}=(z_{1},\ldots,z_{n})$. If its posterior collapses,
then $p(z_{i}=k\,|\,\mathbf{\boldsymbol{x}})=1/K$ for all $k=1,\ldots,K$.

Consider fitting a gmvae model with $K=2$ to a dataset of 5,000 samples. This
dataset is drawn from a gmvae also with $K=2$ well-separated clusters; there
is no model misspecification. A gmvae is typically fit by optimizing the
maximum log marginal likelihood
$\hat{\theta}=\operatorname*{arg\,max}_{\theta}\log
p(\mathbf{\boldsymbol{x}}\,|\,\theta)$. Note there may be multiple values of
$\theta$ that achieve the global optimum of this function.

We focus on two likelihood maximizers. One provides latent variable
identifiability and the posterior of $z_{i}$ does not collapse. The other does
not provide identifiablity; the posterior collapses.

  1. 1.

The first likelihood-maximizing parameter $\hat{\theta}_{1}$ is the truth; the
distribution of the $K$ fitted clusters correspond to the $K$ data-generating
clusters. Given this parameter, the latent variable $z_{i}$ is identifiable
because the $K$ data-generating clusters are different; different cluster
memberships $z_{i}$ must result in different likelihoods
$p(x_{i}\,|\,z_{i}\,;\,\hat{\theta}_{1})$. The posterior of $z_{i}$ does not
collapse.

  2. 2.

In the second likelihood-maximizing parameter $\hat{\theta}_{2}$, however, all
$K$ fitted clusters share the same distribution, each of which is equal to the
marginal distribution of the data. Specifically,
$(\mu_{k}^{*},\Sigma_{k}^{*})=(0,I_{d})$ for all $k$, and each fitted cluster
is a mixture of the $K$ original data generating clusters, i.e., the marginal.
At this parameter value, the model is still able to fully capture the mixture
distribution of the data. However, all the $K$ mixture components are the
same, and thus the latent variable $z_{i}$ is non-identifiable; different
cluster membership $z_{i}$ do not result in different likelihoods
$p(x_{i}\,|\,z_{i}\,;\,\hat{\theta}_{2})$, and hence the posterior of $z_{i}$
collapses. Figure 1(a) illustrates a fit of this (non-identifiable) gmvae to
the pinwheel data [22]. In Section 3, we construct an latent-identifiable VAE
(lidvae) that avoids this collapse.

Latent variable identifiability is a function of the both the model and the
true data-generating distribution. Consider fitting the same gmvae with $K=2$
but to a different dataset of 5,000 samples, this one drawn from a gmvae with
only one cluster. (There is model misspecification.) One maximizing parameter
value $\hat{\theta}_{3}$ is where both of the fitted clusters correspond to
the true data generating cluster. While this parameter value resembles that of
the first maximizer $\hat{\theta}_{1}$ above—both correspond to the true data
generating cluster—this dataset leads to a different situation for latent
variable identifiability. The two fitted clusters are the same and so
different cluster memberships do not result in different likelihoods of
$p(x_{i}\,|\,z_{i}\,;\,\hat{\theta}_{3})$. The latent variable $z_{i}$ is not
identifiable and its posterior collapses.

Takeaways. The gmvae example in this section (and the ppca and gmm examples in
Appendix A) illustrate different ways that a latent variable can be non-
identifiable in a model and suffer from posterior collapse. They show that
even the true posterior—without variational inference—can collapse in non-
identifiable models. They also illustrate that whether a latent variable is
identifiable can depend on both the model and the data. Posterior collapse is
an intrinsic problem of the model and the data, rather than specific to the
use of neural networks or variational inference.

The equivalence between posterior collapse and latent variable non-
identifiability in Theorem 1 also implies that, to mitigate posterior
collapse, we should try to resolve latent variable non-identifiability. In the
next section, we develop such a class of latent-identifiable vae.

##  3 Latent-identifiable vae via Brenier maps

We now construct latent-identifiable vae, a class of vae whose latent
variables are guaranteed to be identifiable, and thus the posteriors cannot
collapse.

###  3.1 The latent-identifiable vae

To construct the latent-identifiable vae, we rely on a key observation that,
to guarantee latent variable identifiability, it is sufficient to make the
likelihood function $P(x_{i}\,|\,z_{i}\,;\,\theta)$ injective for all values
of $\theta$. If the likelihood is injective, then, for any $\theta$, each
value of $z_{i}$ will lead to a different distribution
$P(x_{i}\,|\,z_{i}\,;\,\theta)$. In particular, this fact will be true for any
optimized $\hat{\theta}$ and so the latent $z_{i}$ must be identifiable,
regardless of the data. By Theorem 1, its posterior cannot collapse.

Constructing latent-identifiable vae thus amounts to constructing an injective
likelihood function for vae. The construction is based on a few building
blocks of linear and nonlinear injective functions, then composed into an
injective likelihood $p(x_{i}\,|\,z_{i}\,;\,\theta)$ mapping from
$\mathcal{Z}^{d}$ to $\mathcal{X}^{m}$, where $\mathcal{Z}$ and $\mathcal{X}$
indicate the set of values $z_{i}$ and $x_{i}$ can take. For example, if
$x_{i}$ is an m-dimensional binary vector, then $\mathcal{X}=\\{0,1\\}^{m}$;
if $z_{i}$ is a $K$-dimensional real-valued vector, then
$\mathcal{Z}=\mathbb{R}^{d}$.

The building blocks of lidvae: Injective functions. For linear mappings from
$\mathbb{R}^{d_{1}}$ to $\mathbb{R}^{d_{2}}$ $(d_{2}\geq d_{1})$, we consider
matrix multiplication by a $d_{1}\times d_{2}$-dimensional matrix $\beta$. For
a $d_{1}$-dimensional variable $z$, left multiplication by a matrix
$\beta^{\top}$ is injective when $\beta$ has full column rank [53]. For
example, a matrix with all ones in the diagonal and all other entries being
zero has full column rank.

For nonlinear injective functions, we focus on Brenier maps [4, 36]. A
$d$-dimensional Brenier map is is the gradient of a convex function from
$\mathbb{R}^{d}$ to $\mathbb{R}$. That is, a Brenier map satisfies $g=\nabla
T$ for some convex function $T:\mathbb{R}^{d}\rightarrow\mathbb{R}$. Brenier
maps are also known as a monotone transport map. They are guaranteed to be
bijective [4, 36] because their derivative is the Hessian of a convex $T$,
which must be positive semidefinite and has a nonnegative determinant [4].

To build a vae with Brenier maps, we require a neural network parametrization
of the Brenier map. As Brenier maps are gradients of convex functions, we
begin with the neural network parametrizaton of convex functions, namely the
input convex neural network (icnn) [2, 35]. This parameterization of convex
functions will enable Brenier maps to be paramterized as the gradient of icnn.

An $L$-layer icnn is a neural network mapping from $\mathbb{R}^{d}$ to
$\mathbb{R}$. Given an input $u\in\mathbb{R}^{d}$, its $l$th layer is

$\displaystyle\mathbf{\boldsymbol{z}}_{0}=\mathbf{u},\qquad\mathbf{\boldsymbol{z}}_{l+1}=h_{l}(\mathbf{W}_{l}\mathbf{\boldsymbol{z}}_{l}+\mathbf{A}_{l}\mathbf{u}+\mathbf{b}_{l}),\qquad(l=0,\ldots,L-1),$
(6)

where the last layer $\mathbf{\boldsymbol{z}}_{L}$ must be a scalar,
$\\{\mathbf{W}_{l}\\}$ are non-negative weight matrices with
$\mathbf{W}_{0}=\mathbf{\boldsymbol{0}}$. The functions
$\\{h_{l}:\mathbb{R}\rightarrow\mathbb{R}\\}$ are convex and non-decreasing
entry-wise activation functions for layer $l$; they are applied element-wise
to the vector
$(\mathbf{W}_{l}\mathbf{\boldsymbol{z}}_{l}+\mathbf{A}_{l}\mathbf{u}+\mathbf{b}_{l})$.
A common choice of $h_{0}:\mathbb{R}\rightarrow\mathbb{R}$ is the square of a
leaky RELU, $h_{0}(x)=(\max(\alpha\cdot x,x))^{2}$ with $\alpha=0.2$; the
remaining $h_{l}$’s are set to be a leaky RELU, $h_{l}(x)=\max(\alpha\cdot
x,x)$. This neural network is called “input convex” because it is guaranteed
to be a convex function.

Input convex neural networks can approximate any convex function on a compact
domain in sup norm (Theorem 1 of Chen et al. [9].) Given the neural network
parameterization of convex functions, we can parametrize the Brenier map
$g_{\theta}(\cdot)$ as its gradient with respect to the input
$g_{\theta}(u)=\partial z_{L}/\partial u.$ This neural network
parameterization of Brenier map is a universal approxiamtor of all Brenier
maps on a compact domain, because input convex neural networks are universal
approximators of convex functions [9].

The latent-identifiable VAE (lidvae). We construct injective likelihoods for
lidvae by composing two bijective Brenier maps with an injective matrix
multiplication. As the composition of injective and bijective mappings must be
injective, the resulting composition must be injective. Suppose
$g_{1,\theta}:\mathbb{R}^{K}\rightarrow\mathbb{R}^{K}$ and
$g_{2,\theta}:\mathbb{R}^{D}\rightarrow\mathbb{R}^{D}$ are two Brenier maps,
and $\beta$ is a $K\times D$-dimensional matrix $(D\geq K)$ with all the main
diagonal entries being one and all other entries being zero. The matrix
$\beta^{\top}$ has full column rank, so multiplication by $\beta^{\top}$ is
injective. Thus the composition
$g_{2,\theta}(\beta^{\top}\,g_{1,\theta}(\cdot))$ must be an injective
function from a low-dimensional space $\mathbb{R}^{K}$ to a high-dimensional
space $\mathbb{R}^{D}$.

######  Definition 3 (Latent-identifiable VAE (lidvae) via Brenier maps).

An lidvae via Brenier maps generates a $D$-dimensional datapoint
$x_{i},\in\\{1,\ldots,n\\}$ by:

$\displaystyle z_{i}\sim p(z_{i}),\qquad
x_{i}\,|\,z_{i}\sim\mathrm{EF}(x_{i}\,|\,g_{2,\theta}(\beta^{\top}\,g_{1,\theta}(z_{i}))),$
(7)

where $\mathrm{EF}$ stands for exponential family distributions; $z_{i}$ is a
$K$-dimensional latent variable, discrete or continuous. The parameters of the
model are $\theta\leavevmode\nobreak\ =\leavevmode\nobreak\
(g_{1,\theta},g_{2,\theta})$, where
$g_{1,\theta}:\mathbb{R}^{K}\rightarrow\mathbb{R}^{K}$ and
$g_{2,\theta}:\mathbb{R}^{D}\rightarrow\mathbb{R}^{D}$ are two continuous
Brenier maps. The matrix $\beta$ is a $K\times D$-dimensional matrix $(D\geq
K)$ with all the main diagonal entries being one and all other entries being
zero.

Contrasting lidvae (Eq. 7) with the classical vae (Eq. 1), the lidvae replaces
the function $f_{\theta}:\mathcal{Z}^{K}\rightarrow\mathcal{X}^{D}$ with the
injective mapping $g_{2,\theta}(\beta^{\top}\,g_{1,\theta}(\cdot))$, composed
by bijective Brenier maps $g_{1,\theta},g_{2,\theta}$ and a zero-one matrix
$\beta^{\top}$ with full column rank. As the likelihood functions of
exponential family are injective, the likelihood function
$p(x_{i}\,|\,z_{i}\,;\,\theta)=\mathrm{EF}(g_{2,\theta}(\beta^{\top}\,g_{1,\theta}(z_{i})))$
of lidvae must be injective. Therefore, replacing an arbitrary function
$f_{\theta}:\mathcal{Z}^{K}\rightarrow\mathcal{X}^{D}$ with the injective
mapping $g_{2,\theta}(\beta^{\top}\,g_{1,\theta}(\cdot))$ plays a crucial role
in enforcing identifiability for latent variable $z_{i}$ and avoiding
posterior collapse in lidvae. As the latent $z_{i}$ must be identifiable in
lidvae, its posterior does not collapse.

Despite its injective likelihood, lidvae are as flexible as vae; the use of
Brenier maps and icnn does not limit the capacity of the generative model.
Loosely, lidvae can model any distributions in $\mathbb{R}^{D}$ because
Brenier maps can map any given non-atomic distribution in $\mathbb{R}^{d}$ to
any other one in $\mathbb{R}^{d}$ [36]. Moreover, the icnn parametrization is
a universal approximator of Brenier maps [2]. We summarize the key properties
of lidvae in the following proposition.

######  Proposition 2.

The latent variable $z_{i}$ is identifiable in lidvae, i.e. for all
$i\in\\{1,\ldots,n\\}$, we have

$\displaystyle
p(x_{i}\,|\,z_{i}=\tilde{z}^{\prime}\,;\,\theta)=p(x_{i}\,|\,z_{i}=\tilde{z}\,;\,\theta)\qquad\Rightarrow\qquad\tilde{z}^{\prime}=\tilde{z},\qquad\forall\tilde{z}^{\prime},\tilde{z},\theta.$
(8)

Moreover, for any vae-generated data distribution, there exists an lidvae that
can generate the same distribution. (The proof is in Appendix B.)

###  3.2 Inference in lidvae

Performing inference in lidvae is identical to the classical vae, as the two
vae differ only in their parameter constraints. To fit an lidvae, we use the
classical amortized inference algorithm of vae; we maximize the evidence lower
bound (elbo) of the log marginal likelihood [26].

In general, lidvae are a drop-in replacement for vae. Both have the same
capacity (Proposition 2) and share the same inference algorithm, but lidvae is
identifiable and does not suffer from posterior collapse. The price we pay for
lidvae is computational: the generative model (i.e. decoder) is parametrized
using the gradient of a neural network; its optimization thus requires
calculating gradients of the gradient of a neural network, which increases the
computational complexity of vae inference and can sometimes challenge
optimization. While fitting classical vae using stochastic gradient descent
has $O(k\cdot p)$ computational complexity, where $k$ is the number of
iterations and $p$ is the number of parameters, fitting latent-identifiable
vae may require $O(k\cdot p^{2})$ computational complexity.

###  3.3 Extensions of lidvae

(a) Non-ID GMVAE

(b) IDGMVAE

(c) Accuracy

(d) Log-likelihood

Figure 1: (a)-(b): The posterior of the classical gmvae [51, 27, 13] collapses
when fit to the pinwheel dataset; the latents predict the same value for all
datapoints. The posteriors of lidgmvae, however, do not collapse and provide
meaningful representations. (c)-(d) The latent-identifiable gmvae produces
posteriors that are substantially more informative than gmvae when fit to
fashion MNIST. It also achieves higher test log likelihood.

The construction of lidvae reveals a general strategy to make the latent
variables of generative models identifiable: replacing nonlinear mappings with
injective nonlinear mappings. We can employ this strategy to make the latent
variables of many other vae variants identifiable. Below we give two examples,
mixture vae and sequential vae.

The mixture vae, with gmvae as a special case, models the data with an
exponential family mixture and mapped through a flexible neural network to
generate the data. We develop its latent-identifiable counterpart using
Brenier maps.

######  Example 1 (Latent-identifiable mixture VAE (lidmvae)).

An lidmvae generates a $D$-dimensional datapoint $x_{i},i\in\\{1,\ldots,n\\}$
by

$\displaystyle z_{i}\sim\mathrm{Categorical}(1/K),\quad
w_{i}\,|\,z_{i}\sim\mathrm{EF}(w_{i}\,|\,{\beta_{1}^{\top}\,z_{i}}),\quad
x_{i}\,|\,w_{i}$
$\displaystyle\sim\mathrm{EF}(x_{i}\,|\,g_{2,\theta}(\beta_{2}^{\top}\,g_{1,\theta}(w_{i}))),$
(9)

where $W_{i}$ is a $K$-dimensional one-hot vector that indicates the cluster
assignment. The parameters of the model are $\theta\leavevmode\nobreak\
=\leavevmode\nobreak\ (g_{1,\theta},g_{2,\theta})$, where the functions
$g_{1,\theta}:\mathbb{R}^{M}\rightarrow\mathbb{R}^{M}$ and
$g_{2,\theta}:\mathbb{R}^{D}\rightarrow\mathbb{R}^{D}$ are two continuous
Brenier maps. The matrices $\beta_{1}$ and $\beta_{2}$ are a $K\times
M$-dimensional matrix $(M\geq K)$ and a $M\times D$-dimensional matrix $(D\geq
M)$ respectively, both having all the main diagonal entries being one and all
other entries being zero.

The lidmvae differs from the classical mixture vae in $p(x_{i}\,|\,z_{i})$,
where we replace its neural network mapping with its injective counterpart,
i.e. a composition of two Brenier maps and a matrix multiplication
$g_{2,\theta}(\beta_{2}^{\top}\,g_{1,\theta}(\cdot))$. As a special case,
setting both exponential families in Example 1 as Gaussian gives us lidgmvae,
which we will use to model images in Section 4.

Next we derive the identifiable counterpart of sequential vae, which models
the data with an autoregressive model conditional on the latents.

######  Example 2 (Latent-identifiable sequential VAE (lidsvae)).

An lidsvae generates a $D$-dimensional datapoint $x_{i},i\in\\{1,\ldots,n\\}$
by

$\displaystyle z_{i}\sim p(z_{i}),\qquad
x_{i}\,|\,z_{i},x_{<i}\sim\mathrm{EF}(g_{2,\theta}(\beta_{2}^{\top}\,g_{1,\theta}([z_{i},f_{\theta}(x_{<i})]))),$

where $x_{<i}=(x_{1},\ldots,x_{i-1})$ represents the history of $x$ before the
$i$th dimension. The function
$f_{\theta}:\mathcal{X}_{<i}\rightarrow\mathbb{R}^{H}$ maps the history
$X_{<i}$ into an $H$-dimensional vector. Finally, $[z_{i},f_{\theta}(x_{<i})]$
is an $(K+H)\times 1$ vector that represents a row-stack of the vectors
$(z_{i})_{K\times 1}$ and $(f_{\theta}(x_{<i}))_{H\times 1}$.

Similar with mixture vae, the lidsvae also differs from sequential vae only in
its use of $g_{2,\theta}(\beta_{2}^{\top}\,g_{1,\theta}(\cdot))$ function in
$p(x_{i}\,|\,z_{i},x_{<i})$. We will use lidsvae to model text in Section 4.

##  4 Empirical studies

| Fashion-MNIST |  | Omniglot |
---|---|---|---|---
| AU | KL | MI | LL |  |  | AU | KL | MI | LL
vae [26] | 0.1 | 0.2 | 0.9 | -258.8 |  |  | 0.02 | 0.0 | 0.1 | -862.1
SA-vae [25] | 0.2 | 0.3 | 1.3 | -252.2 |  |  | 0.1 | 0.2 | 1.0 | -853.4
Lagging vae [18] | 0.4 | 0.6 | 1.6 | -248.5 |  |  | 0.5 | 1.0 | 3.6 | -849.4
$\beta$-vae [19] ($\beta$=0.2) | 0.6 | 1.2 | 2.4 | -245.3 |  |  | 0.7 | 1.4 | 5.9 | -842.6
lidgmvae (this work) | 1.0 | 1.6 | 2.6 | -242.3 |  |  | 1.0 | 1.7 | 7.5 | -820.3

| Synthetic | Yahoo | Yelp
---|---|---|---
| AU | KL | MI | LL | AU | KL | MI | LL | AU | KL | MI | LL
vae [26] | 0.0 | 0.0 | 0.0 | -46.5 | 0.0 | 0.0 | 0.0 | -519.7 | 0.0 | 0.0 | 0.0 | -635.9
SA-vae [25] | 0.4 | 0.1 | 0.1 | -40.2 | 0.2 | 1.0 | 0.2 | -520.2 | 0.1 | 1.9 | 0.2 | -631.5
Lagging vae [18] | 0.5 | 0.1 | 0.1 | -40.0 | 0.3 | 1.6 | 0.4 | -518.6 | 0.2 | 3.6 | 0.1 | -631.0
$\beta$-vae [19] ($\beta$=0.2) | 1.0 | 0.1 | 0.1 | -39.9 | 0.5 | 4.7 | 0.9 | -524.4 | 0.3 | 10.0 | 0.1 | -637.3
lidsvae | 1.0 | 0.5 | 0.6 | -40.3 | 0.8 | 7.2 | 1.1 | -519.5 | 0.7 | 9.1 | 0.9 | -634.2

Table 1: Across image and text datasets, lidvae outperforms existing vae
variants in preventing posterior collapse while achieving similar goodness-of-
fit to the data.

We study lidvae on images and text datasets, finding that lidvae do not suffer
from posterior collapse as we increase the capacity of the generative model,
while achieving similar fits to the data. We further study ppca, showing how
likelihood functions nearly constant in latent variables lead to collapsing
posterior even with Markov chain Monte Carlo (mcmc).

###  4.1 lidvae on images and text

We consider three metrics for evaluating posterior collapse: (1) kl divergence
between the posterior and the prior,
$\lx@glossaries@gls@link{acronym}{KL}{\leavevmode\textsc{kl}}(q(\mathbf{\boldsymbol{z}}\,|\,\mathbf{\boldsymbol{x}})||p(\mathbf{\boldsymbol{z}}))$;
(2) Percentange of active units
(au):$\lx@glossaries@gls@link{acronym}{AU}{\leavevmode\textsc{au}}=\sum_{d=1}^{D}\mathbb{1}\\{\mathrm{Cov}_{p(\mathbf{\boldsymbol{x}})}(\mathbb{E}_{q(\mathbf{\boldsymbol{z}}\,|\,\mathbf{\boldsymbol{x}})}\left[\mathbf{\boldsymbol{z}}_{d}\right])\geq\epsilon\\},$
where $\mathbf{\boldsymbol{z}}_{d}=(z_{1d},\ldots,z_{nd})$ is the $d$th
dimension of the latent variable $\mathbf{\boldsymbol{z}}$ for all the $n$
data points. In calculating au, we follow Burda et al. [7] to calculate the
posterior mean,
$(\mathbb{E}\left[z_{1d}\,|\,\mathbf{\boldsymbol{x}}_{1}\right],\ldots,\mathbb{E}\left[z_{nd}\,|\,\mathbf{\boldsymbol{x}}_{n}])\right]$
for all data points, and calculate the sample variance of
$\mathbb{E}\left[z_{id}\,|\,\mathbf{\boldsymbol{x}}_{i}\right]$ across $i$’s
from this vector. The threshold $\epsilon$ is chosen to be 0.01 [7]; the
theoretical maximum of
$\%\lx@glossaries@gls@link{acronym}{AU}{\leavevmode\textsc{au}}$ is one; (3)
Approximate Mutual information (mi) between $\mathbf{\boldsymbol{x}}_{i}$ and
$\mathbf{\boldsymbol{z}}_{i}$,
$I(\mathbf{\boldsymbol{x}},\mathbf{\boldsymbol{z}})=\mathbb{E}_{\mathbf{\boldsymbol{x}}}\left[\mathbb{E}_{q(\mathbf{\boldsymbol{z}}\,|\,\mathbf{\boldsymbol{x}})}\left[\log(q(\mathbf{\boldsymbol{z}}\,|\,\mathbf{\boldsymbol{x}}))\right]\right]-\mathbb{E}_{\mathbf{\boldsymbol{x}}}\left[\mathbb{E}_{q(\mathbf{\boldsymbol{z}}\,|\,\mathbf{\boldsymbol{x}})}\left[\log(q(\mathbf{\boldsymbol{z}}))\right]\right]$.
We also evaluate the model fit using the importance weighted estimate of log-
likelihood on a held-out test set [7]. For mixture vae, we also evaluate the
predictive accuracy of the categorical latents against ground truth labels to
quantify their informativeness.

Competing methods. We compare lidvae with the classical vae [26], the
$\beta$-vae ($\beta$=0.2) [19], the semi-amortized vae [25], and the lagging
vae [18]. Throughout the empirical studies, we use flexible variational
approximating families (RealNVPs [14] for image and LSTMs [20] for text).

Results: Images. We first study lidgmvae on four subsampled image datasets
drawn from pinwheel [22], MNIST [31], Fashion MNIST [57], and Omniglot [30].
Figures 1(a) and 1(b) illustrate a fit of the gmvae and the lidgmvae to the
pinwheel data [22]. The posterior of the gmvae latents collapse, attributing
all datapoints to the same latent cluster. In contrast, lidgmvae produces
categorical latents faithful to the clustering structure. Figure 1 examines
the lidgmvae as we increase the flexibility of the generative model. Figure
1(c) shows that the categorical latents of the lidgmvae are substantially more
predictive of the true labels than their classical counterparts. Moreover, its
performance does not degrade as the generative model becomes more flexible.
Figure 1(d) shows that the lidgmvae consistently achieve higher test log-
likelihood. Table 1 compares different variants of vae in a 9-layer generative
model. Across four datasets, lidgmvae mitigates posterior collapse. It
achieves higher au, kl and mi than other variants of vae. It also achieves a
higher test log-likelihood.

Results: Text. We apply lidsvae to three subsampled text datasets drawn from a
synthetic text dataset, the Yahoo dataset, and the Yelp dataset [60]. The
synthetic dataset is generated from a classical two-layer sequential vae with
a five-dimensional latent. Table 1 compares the lidsvae with the sequential
vae. Across the three text datasets, the lidsvae outperforms other variants of
vae in mitigating posterior collapse, generally achieving a higher au, kl, and
mi.

###  4.2 Latent variable non-identifiability and posterior collapse in PPCA

(a) $\sigma=0.2$

(b) $\sigma=0.5$

(c) $\sigma=1.0$

(d) $\sigma=1.5$

Figure 2: As the noise level increases in ppca, the latent variable becomes
closer to non-identifiable because the likelihood and more susceptible to
posterior collapse. Its likelihood surface becomes flatter and its posterior
becomes closer to the prior. Top panel: Likelihood surface of ppca as a
function of the two latents $z_{1},z_{2}$. When $\sigma$ increase, the
likelihood surface becomes flatter and the latent variables $z_{1},z_{2}$ are
closer to non-identifiable. Bottom panel: Posterior of $z_{1}$ under different
$\sigma$ values. When $\sigma$ increase, the posterior becomes closer to the
prior.

Here we show that the PPCA posterior becomes close to the prior when the
latent variable becomes close to be non-identifiable. We perform inference
using Hamiltonian Monte Carlo (hmc), avoiding the effect of variational
approximation on posterior collapse.

Consider a ppca with two latent dimensions,
$p(z_{i})=\mathcal{N}(z_{i}\,;\,0,I_{2}),\,p(x_{i}\,|\,z_{i}\,;\,\theta)=\mathcal{N}(x_{i}\,;\,z_{i}^{\top}w,\sigma^{2}\cdot
I_{5}),$ where the value of $\sigma^{2}$ is known, $z_{i}$’s are the latent
variables of interest, and $w$ is the only parameter of interest. When the
noise $\sigma^{2}$ is set to a large value, the latent variable $z_{i}$ may
become nearly non-identifiable. The reason is that the likelihood function
$p(x_{i}\,|\,z_{i})$ becomes slower-varying as $\sigma^{2}$ increases. For
example, Figure 2 shows that the likelihood surface becomes flatter as
$\sigma^{2}$ increases. Accordingly, the posterior becomes closer to the prior
as $\sigma^{2}$ increases. When $\sigma=1.5$, the posterior collapses. This
non-identifiability argument provides an explanation to the closely related
phenomenon described in Section 6.2 of [33].

##  5 Discussion

In this work, we show that the posterior collapse phenomenon is a problem of
latent variable non-identifiability. It is not specific to the use of neural
networks or particular inference algorithms in vae. Rather, it is an intrinsic
issue of the model and the dataset. To this end, we propose a class of lidvae
via Brenier maps to resolve latent variable non-identifiability and mitigate
posterior collapse. Across empirical studies, we find that lidvae outperforms
existing methods in mitigating posterior collapse.

The latent variables of lidvae are guaranteed to be identifiable. However, it
does not guarantee that the latent variables and the parameters of lidvae are
jointly identifiable. In other words, the lidvae model may not be identifiable
even though its latents are identifiable. This difference between latent
variable identifiability and model identifiability may appear minor. But the
tractability of resolving latent variable identifiability plays a key role in
making non-identifiability a fruitful one perspective of posterior collapse.
To enforce latent variable identifiability, it is sufficient to ensure that
the likelihood
$p(\mathbf{\boldsymbol{x}}\,|\,\mathbf{\boldsymbol{z}},\hat{\theta})$ is an
injective function of $\mathbf{\boldsymbol{z}}$. In contrast, resolving model
identifiability for the general class of vae remains a long standing open
problem, with some recent progress relying on auxiliary variables [23, 24].
The tractability of resolving latent variable identifiability is a key
catalyst of a principled solution to mitigating posterior collapse.

There are a few limitations of this work. One is that the theoretical argument
focuses on the collapse of the exact posterior. The rationale is that, if the
exact posterior collapses, then its variational approximation must also
collapse because variational approximation of posteriors cannot “uncollapse” a
posterior. That said, variational approximation may “collapse” a posterior,
i.e. the exact posterior does not collapse but the variational approximate
posterior collapses. The theoretical argument and algorithmic approaches
developed in this work does not apply to this setting, which remains an
interesting venue of future work.

A second limitation is that the latent-identifiable vae developed in this work
bear a higher computational cost than classical vae. While the latent-
identifiable vae ensures the identifiability of its latent variables and
mitigates posterior collapse, it does come with a price in computation because
its generative model (i.e. decoder) is parametrized using gradients of a
neural network. Fitting the latent-identifiable vae thus requires calculating
gradients of gradients of a neural network, leading to much higher
computational complexity than fitting the classifical vae. Developing
computationally efficient variants of the latent-identifiable vae is another
interesting direction for future work.

Acknowledgments.  We thank Taiga Abe and Gemma Moran for helpful discussions,
and anonymous reviewers for constructive feedback that improved the
manuscript. David Blei is supported by ONR N00014-17-1-2131, ONR
N00014-15-1-2209, NSF CCF-1740833, DARPA SD2 FA8750-18-C-0130, Amazon, and the
Simons Foundation. John Cunningham is supported by the Simons Foundation,
McKnight Foundation, Zuckerman Institute, Grossman Center, and Gatsby
Charitable Trust.

## References

  * [1] A. A. Alemi, B. Poole, I. Fischer, J. V. Dillon, R. A. Saurous, and K. Murphy (2017) Fixing a broken ELBO.  arXiv preprint arXiv:1711.00464.  Cited by: §1.
  * [2] B. Amos, L. Xu, and J. Z. Kolter (2017) Input convex neural networks.  In Proceedings of the 34th International Conference on Machine Learning-Volume 70,  pp. 146–155.  Cited by: §1, §3.1, §3.1.
  * [3] A. Asperti (2019) Variational autoencoders and the variable collapse phenomenon.  Sensors & Transducers 234 (6),  pp. 1–8.  Cited by: §1.
  * [4] K. Ball (2004) An elementary introduction to monotone transportation.  In Geometric aspects of functional analysis,  pp. 41–52.  Cited by: §3.1.
  * [5] M. Betancourt (2017) Identifying Bayesian mixture models.  Note: https://mc-stan.org/users/documentation/case-studies/identifying_mixture_modelsAccessed: 2021-05-04 Cited by: §1, §2.1.
  * [6] S. Bowman, L. Vilnis, O. Vinyals, A. Dai, R. Jozefowicz, and S. Bengio (2016) Generating sentences from a continuous space.  In Proceedings of The 20th SIGNLL Conference on Computational Natural Language Learning,  pp. 10–21.  Cited by: §1, §1, §1, Definition 1.
  * [7] Y. Burda, R. Grosse, and R. Salakhutdinov (2015) Importance weighted autoencoders.  arXiv preprint arXiv:1509.00519.  Cited by: §4.1.
  * [8] X. Chen, D. P. Kingma, T. Salimans, Y. Duan, P. Dhariwal, J. Schulman, I. Sutskever, and P. Abbeel (2016) Variational lossy autoencoder.  arXiv preprint arXiv:1611.02731.  Cited by: §1, §1, Definition 1.
  * [9] Y. Chen, Y. Shi, and B. Zhang (2018) Optimal control via neural networks: a convex approach.  arXiv preprint arXiv:1805.11835.  Cited by: §3.1.
  * [10] M. Collins, S. Dasgupta, and R. E. Schapire (2001) A generalization of principal components analysis to the exponential family..  In Nips,  Vol. 13,  pp. 23.  Cited by: §2.
  * [11] B. Dai, Z. Wang, and D. Wipf (2019) The usual suspects? Reassessing blame for VAE posterior collapse.  arXiv preprint arXiv:1912.10702.  Cited by: §1, §1.
  * [12] A. B. Dieng, Y. Kim, A. M. Rush, and D. M. Blei (2018) Avoiding latent variable collapse with generative skip models.  arXiv preprint arXiv:1807.04863.  Cited by: §1, §2.1, §2.1.
  * [13] N. Dilokthanakul, P. A. Mediano, M. Garnelo, M. C. Lee, H. Salimbeni, K. Arulkumaran, and M. Shanahan (2016) Deep unsupervised clustering with Gaussian mixture variational autoencoders.  arXiv preprint arXiv:1611.02648.  Cited by: §2.2, Figure 1, Figure 1.
  * [14] L. Dinh, J. Sohl-Dickstein, and S. Bengio (2016) Density estimation using real NVP.  arXiv preprint arXiv:1605.08803.  Cited by: §4.1.
  * [15] H. Fu, C. Li, X. Liu, J. Gao, A. Celikyilmaz, and L. Carin (2019) Cyclical annealing schedule: a simple approach to mitigating KL vanishing.  arXiv preprint arXiv:1903.10145.  Cited by: §1.
  * [16] I. Gulrajani, K. Kumar, F. Ahmed, A. A. Taiga, F. Visin, D. Vazquez, and A. Courville (2016) Pixelvae: a latent variable model for natural images.  arXiv preprint arXiv:1611.05013.  Cited by: §1, §1.
  * [17] S. Havrylov and I. Titov (2020) Preventing posterior collapse with Levenshtein variational autoencoder.  arXiv preprint arXiv:2004.14758.  Cited by: §1, §2.1.
  * [18] J. He, D. Spokoyny, G. Neubig, and T. Berg-Kirkpatrick (2019) Lagging inference networks and posterior collapse in variational autoencoders.  arXiv preprint arXiv:1901.05534.  Cited by: §1, §1, §2.1, §2.1, §4.1, Table 1, Table 1.
  * [19] I. Higgins, L. Matthey, A. Pal, C. Burgess, X. Glorot, M. Botvinick, S. Mohamed, and A. Lerchner (2016) $\beta$-VAE: learning basic visual concepts with a constrained variational framework.  Cited by: §1, §2.1, §4.1, Table 1, Table 1.
  * [20] S. Hochreiter and J. Schmidhuber (1997) Long short-term memory.  Neural computation 9 (8),  pp. 1735–1780.  Cited by: §4.1.
  * [21] M. D. Hoffman and M. J. Johnson (2016) ELBO surgery: Yet another way to carve up the variational evidence lower bound.  Cited by: §1.
  * [22] M. J. Johnson, D. K. Duvenaud, A. Wiltschko, R. P. Adams, and S. R. Datta (2016) Composing graphical models with neural networks for structured representations and fast inference.  In Advances in neural information processing systems,  pp. 2946–2954.  Cited by: item 2, §4.1.
  * [23] I. Khemakhem, D. P. Kingma, and A. Hyvärinen (2019) Variational autoencoders and nonlinear ICA: a unifying framework.  arXiv preprint arXiv:1907.04809.  Cited by: §1, §5.
  * [24] I. Khemakhem, R. P. Monti, D. P. Kingma, and A. Hyvarinen (2020) ICE-BeeM: Identifiable conditional energy-based deep models based on nonlinear ICA.  External Links: 2002.11537 Cited by: §1, §5.
  * [25] Y. Kim, S. Wiseman, A. Miller, D. Sontag, and A. Rush (2018) Semi-amortized variational autoencoders.  In International Conference on Machine Learning,  pp. 2678–2687.  Cited by: §1, §1, §2.1, §2.1, §4.1, Table 1, Table 1.
  * [26] D. P. Kingma and M. Welling (2014) Auto-encoding variational Bayes.  In Proceedings of the International Conference on Learning Representations (ICLR),  Vol. 1.  Cited by: §1, §3.2, §4.1, Table 1, Table 1.
  * [27] D. P. Kingma, S. Mohamed, D. J. Rezende, and M. Welling (2014) Semi-supervised learning with deep generative models.  In Advances in neural information processing systems,  pp. 3581–3589.  Cited by: Figure 1, Figure 1.
  * [28] D. P. Kingma, T. Salimans, R. Jozefowicz, X. Chen, I. Sutskever, and M. Welling (2016) Improved variational inference with inverse autoregressive flow.  In Advances in neural information processing systems,  pp. 4743–4751.  Cited by: §1.
  * [29] A. Kumar and B. Poole (2020) On implicit regularization in $\beta$-vaes.  In International Conference on Machine Learning,  pp. 5480–5490.  Cited by: §1.
  * [30] B. M. Lake, R. Salakhutdinov, and J. B. Tenenbaum (2015) Human-level concept learning through probabilistic program induction.  Science 350 (6266),  pp. 1332–1338.  Cited by: §4.1.
  * [31] Y. LeCun, C. Cortes, and C. Burges (2010) MNIST handwritten digit database.  ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist 2.  Cited by: §4.1.
  * [32] B. Li, J. He, G. Neubig, T. Berg-Kirkpatrick, and Y. Yang (2019) A surprisingly effective fix for deep latent variable modeling of text.  arXiv preprint arXiv:1909.00868.  Cited by: §1, §1.
  * [33] J. Lucas, G. Tucker, R. B. Grosse, and M. Norouzi (2019) Don’t blame the ELBO! A linear VAE perspective on posterior collapse.  In Advances in Neural Information Processing Systems,  pp. 9403–9413.  Cited by: §1, §2.1, §4.2.
  * [34] L. Maaløe, M. Fraccaro, V. Liévin, and O. Winther (2019) BIVA: A very deep hierarchy of latent variables for generative modeling.  In Advances in Neural Information Processing Systems, H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (Eds.),  Vol. 32,  pp. .  External Links: Link Cited by: §1.
  * [35] A. V. Makkuva, A. Taghvaei, S. Oh, and J. D. Lee (2019) Optimal transport mapping via input convex neural networks.  arXiv preprint arXiv:1908.10962.  Cited by: §1, §3.1.
  * [36] R. J. McCann and N. Guillen (2011) Five lectures on optimal transportation: geometry, regularity and applications.  Analysis and geometry of metric measure spaces: Lecture notes of the séminaire de Mathématiques Supérieure (SMS) Montréal,  pp. 145–180.  Cited by: §3.1, §3.1.
  * [37] R. J. McCann et al. (1995) Existence and uniqueness of monotone measure-preserving maps.  Duke Mathematical Journal 80 (2),  pp. 309–324.  Cited by: §1.
  * [38] A. v. d. Oord, O. Vinyals, and K. Kavukcuoglu (2017) Neural discrete representation learning.  arXiv preprint arXiv:1711.00937.  Cited by: §1, §1, Definition 1.
  * [39] G. Peyré, M. Cuturi, et al. (2019) Computational optimal transport.  Foundations and Trends® in Machine Learning 11 (5-6),  pp. 355–607.  Cited by: §1.
  * [40] D. J. Poirier (1998) Revising beliefs in nonidentified models.  Econometric Theory 14 (4),  pp. 483–509.  Cited by: §1, §1, §2.1.
  * [41] B.L.S.P. Rao and R.B.L.S. Prakasa (1992) Identifiability in stochastic models: characterization of probability distributions.  Probability and mathematical statistics, Academic Press.  External Links: ISBN 9780125640152, LCCN lc91040038, Link Cited by: §2.1.
  * [42] A. Raue, C. Kreutz, T. Maiwald, J. Bachmann, M. Schilling, U. Klingmüller, and J. Timmer (2009) Structural and practical identifiability analysis of partially observed dynamical models by exploiting the profile likelihood.  Bioinformatics 25 (15),  pp. 1923–1929.  Cited by: §1, §1, §2.1, §2.1, Definition 2.
  * [43] A. Raue, C. Kreutz, F. J. Theis, and J. Timmer (2013) Joining forces of Bayesian and frequentist methodology: a study for inference in the presence of non-identifiability.  Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences 371 (1984),  pp. 20110544.  Cited by: §1, §1, §2.1.
  * [44] A. Razavi, A. v. d. Oord, B. Poole, and O. Vinyals (2019) Preventing posterior collapse with delta-VAEs.  arXiv preprint arXiv:1901.03416.  Cited by: §1, §1, §1, §2.1.
  * [45] D. A. Reynolds (2009) Gaussian mixture models..  Encyclopedia of biometrics 741,  pp. 659–663.  Cited by: §2.
  * [46] D. J. Rezende, S. Mohamed, and D. Wierstra (2014) Stochastic backpropagation and approximate inference in deep generative models.  arXiv preprint arXiv:1401.4082.  Cited by: §1.
  * [47] S. Roweis and Z. Ghahramani (1999) A unifying review of linear gaussian models.  Neural computation 11 (2),  pp. 305–345.  Cited by: §2.
  * [48] S. Roweis (1998) EM algorithms for pca and spca.  Advances in neural information processing systems,  pp. 626–632.  Cited by: §2.
  * [49] E. San Martın and J. González (2010) Bayesian identifiability: contributions to an inconclusive debate.  Chilean Journal of Statistics 1 (2),  pp. 69–91.  Cited by: §1, §1, §2.1.
  * [50] B. Seybold, E. Fertig, A. Alemi, and I. Fischer (2019) Dueling decoders: regularizing variational autoencoder latent spaces.  arXiv preprint arXiv:1905.07478.  Cited by: §1.
  * [51] R. Shu (2016-12) Gaussian mixture VAE: lessons in variational inference, generative models, and deep nets.  External Links: Link Cited by: §1, §2.2, Figure 1, Figure 1.
  * [52] C. K. Sønderby, T. Raiko, L. Maaløe, S. K. Sønderby, and O. Winther (2016) How to train deep variational autoencoders and probabilistic ladder networks.  In 33rd International Conference on Machine Learning (ICML 2016),  Cited by: §1.
  * [53] G. Strang, G. Strang, G. Strang, and G. Strang (1993) Introduction to linear algebra.  Vol. 3, Wellesley-Cambridge Press Wellesley, MA.  Cited by: §3.1.
  * [54] M. E. Tipping and C. M. Bishop (1999) Probabilistic principal component analysis.  Journal of the Royal Statistical Society: Series B (Statistical Methodology) 61 (3),  pp. 611–622.  Cited by: §2.
  * [55] J. M. Tomczak and M. Welling (2017) VAE with a VampPrior.  arXiv preprint arXiv:1705.07120.  Cited by: §1.
  * [56] F. Wieland, A. L. Hauber, M. Rosenblatt, C. Tönsing, and J. Timmer (2021) On structural and practical identifiability.  Current Opinion in Systems Biology.  Cited by: §1, §1, §2.1, §2.1, Definition 2.
  * [57] H. Xiao, K. Rasul, and R. Vollgraf (2017) Fashion-MNIST: a novel image dataset for benchmarking machine learning algorithms.  CoRR abs/1708.07747.  External Links: Link, 1708.07747 Cited by: §4.1.
  * [58] Y. Xie and B. P. Carlin (2006) Measures of Bayesian learning and identifiability in hierarchical models.  Journal of Statistical Planning and Inference 136 (10),  pp. 3458–3477.  Cited by: §1, §1, §2.1.
  * [59] Y. Yacoby, W. Pan, and F. Doshi-Velez (2020) Characterizing and avoiding problematic global optima of variational autoencoders.  External Links: 2003.07756 Cited by: §1, §1, §2.1.
  * [60] Z. Yang, Z. Hu, R. Salakhutdinov, and T. Berg-Kirkpatrick (2017) Improved variational autoencoders for text modeling using dilated convolutions.  In Proceedings of the 34th International Conference on Machine Learning-Volume 70,  pp. 3881–3890.  Cited by: §1, §4.1.
  * [61] S. Yeung, A. Kannan, Y. Dauphin, and L. Fei-Fei (2017) Tackling over-pruning in variational autoencoders.  arXiv preprint arXiv:1706.03643.  Cited by: §1.
  * [62] T. Zhao, K. Lee, and M. Eskenazi (2018) Unsupervised discrete sentence representation learning for interpretable neural dialog generation.  In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers),  pp. 1098–1107.  Cited by: §1, §1, Definition 1.
  * [63] Y. Zhao, P. Yu, S. Mahapatra, Q. Su, and C. Chen (2020) Discretized bottleneck in VAE: posterior-collapse-free sequence-to-sequence learning.  arXiv preprint arXiv:2004.10603.  Cited by: §1.

Supplementary Materials

Posterior Collapse and Latent Variable Non-identifiability

##  Appendix A Examples of posterior collapse continued

We present two additional examples of posterior collapse, probabilistic
principal component analysis and Gaussian mixture model.

###  A.1 Probabilistic principal component analysis

We consider classical probabilistic principal component analysis (ppca) and
show that its local latent variables can suffer from posterior collapse at
maximum likelihood parameter values (i.e. global maxima of log marginal
likelihood). This example refines the perspective of Lucas et al. [7], which
demonstrated that posterior collapse can occur in ppca absent any variational
approximation but due to local maxima in the log marginal likelihood. Here we
show that posterior collapse can occur even with global maxima, absent
optimization issues due to local maxima.

Consider a ppca with two latent dimensions,

$\displaystyle p(z_{i})$ $\displaystyle=\mathcal{N}(z_{i}\,|\,0,I_{2}),$
$\displaystyle p(x_{i}\,|\,z_{i}\,;\,\theta)$
$\displaystyle=\mathcal{N}(x_{i}\,|\,z_{i}^{\top}w,\sigma^{2}\cdot I_{5}),$

where $z_{i}$’s are the latent variables of interest and others
$\theta=(w,\sigma^{2})$ are parameters of the model.

Consider fitting this model to two datasets, each with 500 samples, focusing
on maximum likelihood parameter values. Depending on the true distribution of
the dataset, ppca may or may not suffer from posterior collapse.

(a) Likelihood (1D PPCA)

(b) Posterior (1D PPCA)

(c) Likelihood (2D PPCA)

(d) Posterior (2D PPCA)

Figure 3: Fitting ppca with more latent dimensions than enough leads to non-
identifiable local latent variables and collapsed posteriors. (a)-(b) Fit a
two-dimensional ppca to data drawn from a one-dimensional ppca. The likelihood
surface is constant in one dimension of the latent variable, i.e. this latent
variable is non-identifiable. Hence its corresponding posterior collapses.
(c)-(d) Fit a two-dimensional ppca to data from a two-dimensional ppca does
not suffer from posterior collapse; its likelihood surface varies in all
dimensions.

  1. 1.

Sample the data from a one-dimensional ppca,

$\displaystyle
x_{i}\sim\mathcal{N}(x_{i}\,|\,\mathcal{N}(0,I_{1})\cdot\bar{w}_{1},\bar{\sigma}_{1}\cdot
I_{5}).$ (10)

(The model remains two dimensional.) The latent variables $z_{i}$’s are not
(fully) identifiable in this case. The reason is that one set of maximum
likelihood parameters is
$\hat{\theta}=(\hat{w},\hat{\sigma})=([\mathbf{\boldsymbol{0}},\bar{w}_{1}],\bar{\sigma}_{1})$,
i.e. setting one latent dimension as zero and the other equal to the true data
generating direction. Under this $\hat{\theta}$, the likelihood function is
constant in the first dimension of the latent variable, i.e. $z_{i1}$; see
Figure 3(a). The posterior of $z_{i1}$ thus collapses, matching the prior,
while the posterior of $z_{i2}$ stays peaked (Figure 3(b)).

  2. 2.

Sample the data from from a two-dimensional ppca,

$\displaystyle
x_{i}\sim\mathcal{N}(x_{i}\,|\,\mathcal{N}(0,I_{2})\cdot\bar{w}_{2},\bar{\sigma}_{2}\cdot
I_{5}).$ (11)

The latent variables $z_{i}$ are identifiable. The likelihood function varies
against both $z_{i1}$ and $z_{i2}$; the posteriors of both $z_{i1}$ and
$z_{i2}$ are peaked (Figures 3(c) and 3(d)).

###  A.2 Gaussian mixture model

(a) Likelihood function

(b) Posterior histogram

Figure 4: When a latent variable is non-identifiable (non-ID) in a model, its
likelihood function is a constant function and its posterior is equal to the
prior, i.e. its posterior collapses. Consider a Gaussian mixture model with
two clusters
$x\sim\alpha\cdot\mathcal{N}(\mu_{1},\sigma_{1}^{2})+(1-\alpha)\cdot\mathcal{N}(\mu_{2},\sigma_{2}^{2})$,
treating the mixture weight $\alpha$ as the latent variable and others as
parameters. Fit the model to datasets generated respectively by one Gaussian
cluster ($\alpha$ non-identifiable), two overlapping Gaussian clusters
($\alpha$ nearly non-identifiable), and two non-overlapping Gaussian clusters
($\alpha$ identifiable). Under optimal parameters, the likelihood function
$p(x\,|\,\alpha)$ is (close to) a constant when the latent variable $\alpha$
is (close to) non-identifiable; its posterior is also (close to) the prior.
Otherwise, the likelihood function is non-constant and the posterior is
peaked.

Though we have focused on the posterior collapse of local latent variables, a
model can also suffer from posterior collapse of its global latent variables.
Consider a simple Gaussian mixture model (gmm) with two clusters,

$\displaystyle p(\alpha)$ $\displaystyle=\mathrm{Beta}(\alpha\,|\,5,5),$
$\displaystyle p(x_{i}\,|\,\alpha\,;\,\theta)$
$\displaystyle=\alpha\cdot\mathcal{N}(x_{i}\,|\,\mu_{1},\sigma_{1}^{2})+(1-\alpha)\cdot\mathcal{N}(x_{i}\,|\,\mu_{2},\sigma_{2}^{2}).$

Here $\alpha$ is a global latent variable and
$\theta=(\mu_{1},\mu_{2},\sigma_{1},\sigma_{2})$ are the parameters of the
model. Fit this model to three datasets, each with $10^{5}$ samples.

  1. 1.

Sample the data from two non-overlapping clusters,

$\displaystyle x_{i}\sim
0.15\cdot\mathcal{N}(-10,1)+0.85\cdot\mathcal{N}(10,1).$ (12)

The latent variable $\alpha$ is identifiable. The two data generating clusters
are substantially different, so the likelihood function varies across
$\alpha\in[0,1]$ under the maximum likelihood (ml) parameters (Figure 4(a)).
The posterior of $\alpha$ is also peaked (Figure 4(b)) and differs much from
the prior.

  2. 2.

Sample the data from two overlapping clusters,

$\displaystyle x_{i}\sim
0.15\cdot\mathcal{N}(-0.5,1)+0.85\cdot\mathcal{N}(0.5,1).$ (13)

The latent variable $\alpha$ is identifiable. However, it is nearly non-
identifiable. While the two data generating clusters are different, they are
very similar to each other because they overlap. Therefore, the likelihood
function $p(x_{i}\,|\,\alpha\,;\,\theta^{*})$ is slowly varying under ml
parameters
$\theta^{*}=(\mu_{1}^{*},\mu_{2}^{*},\sigma_{1}^{*},\sigma_{2}^{*})=(-0.5,0.5,1,1)$;
see Figure 4(a). Consequently, the posterior of $\alpha$ remains very close to
the prior; see Figure 4(b).

  3. 3.

Sample the data from a single Gaussian distribution,
$x_{i}\sim\mathcal{N}(-1,1)$. The latent variable $\alpha$ is non-
identifiable. The reason is that one set of ml parameters is
$\theta^{*}=(\mu_{1}^{*},\mu_{2}^{*},\sigma_{1}^{*},\sigma_{2}^{*})=(-1,-1,1,1)$,
i.e. setting both of the two mixture components equal to the true data
generating Gaussian distribution.

Under this $\theta^{*}$, the latent variable $\alpha$ is non-identifiable and
its likelihood function $p(\\{x_{i}\\}_{i=1}^{n}\,|\,\alpha\,;\,\theta^{*})$
is constant in $\alpha$ because the two mixture components are equal; Figure
4(a) illustrates this fact. Moreover, the posterior of $\alpha$ collapses,
$p(\alpha\,|\,\\{x_{i}\\}_{i=1}^{n}\,;\,\theta^{*})=p(\alpha).$ Figure 4(b)
illustrates this fact: The hmc samples of the $\alpha$ posterior closely match
those drawn from the prior. (Exact inference is intractable in this case, so
we use hmc as a close approximation to exact inference.) This example
demonstrates the connection between non-identifiability and posterior
collapse; it also shows that posterior collapse is not specific to variational
inference but is an issue of the model and the data.

As for PPCA, these GMM examples demonstrate that whether a latent variable is
identifiable in a probabilistic model not only depends on the model but also
the data. While all three examples were fitted with the same gmm model, their
identifiability situation differs as the samples are generated in different
ways.

##  Appendix B Proof of Proposition 2

We prove a general version of Proposition 2 by establishing the latent
variable identifiability and flexibility of the most general form of the
lidvae. The lidvae, lidmvae, and lidsvae (Definitions 4, 1 and 2) will all be
its special cases. Then Proposition 2 will also be a special case of the more
general result stated below (Proposition 3).

We first define the most general form of lidvae.

######  Definition 4 (General lidvae via Brenier maps).

A general lidvae via Brenier maps generates an $D$-dimensional data-point
$x_{i},\in\\{1,\ldots,n\\}$ by:

$\displaystyle(z_{i})_{K\times 1}$ $\displaystyle\sim p(z_{i}),$ (14)
$\displaystyle(w_{i})_{M\times 1}\,|\,z_{i}$
$\displaystyle\sim\mathrm{EF}(w_{i}\,|\,{\beta_{1}^{\top}\,z_{i}}),$ (15)
$\displaystyle(x_{i})_{D\times 1}\,|\,w_{i},x_{<i}$
$\displaystyle\sim\mathrm{EF}(x_{i}\,|\,h\circ
g_{2,\theta}(\beta_{2}^{\top}\,g_{1,\theta}([w_{i},f_{\theta}(x_{<i})]))),$
(16)

where $\mathrm{EF}$ stands for exponential family distributions; $z_{i}$ is a
$K$-dimensional latent variable, discrete or continuous. The parameters of the
model are $\theta\leavevmode\nobreak\ =\leavevmode\nobreak\
(g_{1,\theta},g_{2,\theta},f_{\theta})$, where
$f_{\theta}:\mathcal{X}_{<i}\rightarrow\mathbb{R}^{H}$ is a function that maps
all previous data points $x_{<i}$ to an $H$-dimensional vector,
$g_{1,\theta}:\mathbb{R}^{M+H}\rightarrow\mathbb{R}^{M+H}$ and
$g_{2,\theta}:\mathbb{R}^{D}\rightarrow\mathbb{R}^{D}$ are two continuous
monotone transport maps. The function $h(\cdot)$ is a bijective link function
for the exponential family, e.g. the sigmoid function. The matrix $\beta_{1}$
is a $K\times M$-dimensional matrix $(M\geq K)$ all the main diagonal entries
being one and all other entries being zero, and thus with full row rank.
Similarly, $\beta_{2}$ is a $(M+H)\times D$-dimensional matrix $(D\geq M+H)$
with all the main diagonal entries being one and all other entries being zero,
also with full row rank. Finally, $[w_{i},f_{\theta}(x_{<i})]$ is an
$(M+H)\times 1$ vector that represents a row-stack of the vectors
$(w_{i})_{M\times 1}$ and $(f_{\theta}(x_{<i}))_{H\times 1}$.

The general lidvae differs from the classical vae whose general form is

$\displaystyle(z_{i})_{K\times 1}$ $\displaystyle\sim p(z_{i}),$ (17)
$\displaystyle(w_{i})_{M\times 1}\,|\,z_{i}$
$\displaystyle\sim\mathrm{EF}(w_{i}\,|\,{\beta_{1}^{\top}\,z_{i}}),$ (18)
$\displaystyle(x_{i})_{D\times 1}\,|\,w_{i},x_{<i}$
$\displaystyle\sim\mathrm{EF}(x_{i}\,|\,h\circ
g_{\theta}([w_{i},f_{\theta}(x_{<i})])),$ (19)

The key difference is in Eq. 19, where the classical vae uses an arbitrary
function $g:\mathbb{R}^{M+H}\rightarrow\mathbb{R}^{D}$ in Eq. 19. In contrast,
lidvae uses a composition
$g_{2,\theta}(\beta_{2}^{\top}\,g_{1,\theta}(\cdot))$ with additional
constraints in Eq. 16.

General lidvae can handle both i.i.d. and sequential data. For i.i.d data
(e.g. images), we can set $f_{\theta}(\cdot)$ to be a zero function, which
implies $P(x_{i}\,|\,w_{i},x_{<i})=P(x_{i}\,|\,w_{i})$. For sequential data
(e.g. text), we can set $f_{\theta}(\cdot)$ to be an LSTM that embeds the
history $x_{<i}$ into an $H$-dimensional vector.

General lidvae emulate many existing vae. Letting $z_{i}$ be categorical (one-
hot) vectors, the distribution $\mathrm{EF}(z_{i}^{\top}\beta_{\theta})$ is an
exponential family mixture. The identifiable vae then maps this mixture model
through a flexible function $g_{\theta}$. When $z_{i}$ is real-valued, it
mimics classical vae by mapping an exponential family PCA through flexible
functions.

lidgmvae is a special case of the general lidvae when we set $z_{i}$ be
categorical (one-hot) vectors, set the exponential family distribution
$\mathrm{EF}$ to be Gaussian in Eqs. 15 and 16. In this case,
$w_{i}\sim\mathrm{Gaussian}(z_{i}^{\top}\beta_{\theta},\gamma_{\theta})$ is a
Gaussian mixture. Then, we set $f_{\theta}(\cdot)$ to be a zero function,
which implies $P(x_{i}\,|\,w_{i},x_{<i})=P(x_{i}\,|\,w_{i})$, and finally set
$h$ as the identity function.

This general lidvae also subsumes the Bernoulli mixture model, which is a
common variant of lidgmvae for the MNIST data. Specifically, we can set
$z_{i}$ be categorical (one-hot) vectors, and then set the exponential family
distribution $\mathrm{EF}$ to be Gaussian in Eq. 15, making
$w_{i}\sim\mathrm{Gaussian}(z_{i}^{\top}\beta_{\theta},\gamma_{\theta})$ to be
a Gaussian mixture. Next we set $f_{\theta}(\cdot)$ to be a zero function,
which implies $P(x_{i}\,|\,w_{i},x_{<i})=P(x_{i}\,|\,w_{i})$, then set $h$ to
be the sigmoid function, and finally set the $\mathrm{EF}$ to be Bernoulli in
Eq. 16.

lidsvae is another special case of the general lidvae when we set the
$\mathrm{EF}$ to be a point mass and $\beta_{1,\theta}$ to be identity matrix
in Eq. 15, which implies $w_{i}=z_{i}$. Then setting the $\mathrm{EF}$ to be a
categorical distribution and $h$ to be identity in Eq. 16 leads to a
configuration that is the same as Example 2.

lidvae can be made deeper with more layers by introducing additional full row-
rank matrices $\beta_{k}$ (e.g. ones with all the main diagonal entries being
one and all other entries being zero) and additional Brenier maps
$g_{k,\theta}$. For example, we can expand Eq. 16 with an additional layer by
setting

$(x_{i})_{D\times
1}\,|\,w_{i},x_{<i}\sim\mathrm{EF}(g_{3,\theta}(\beta_{3}^{\top}g_{2,\theta}(\beta_{2}^{\top}\,g_{1,\theta}([w_{i},f_{\theta}(x_{<i})])))).$

Next we establish the latent variable identifiability and flexibility of this
general class of lidvae, which will imply the identifiability and flexibility
of all the special cases above.

######  Proposition 3.

The latent variable $z_{i}$ is identifiable in lidvae, i.e. for all
$i\in\\{1,\ldots,n\\}$, we have

$\displaystyle
p(x_{i}\,|\,z_{i}=\tilde{z}^{\prime},x_{<i}\,;\,\theta)=p(x_{i}\,|\,z_{i}=\tilde{z},x_{<i}\,;\,\theta)\qquad\Rightarrow\qquad\tilde{z}^{\prime}=\tilde{z},\qquad\forall\tilde{z}^{\prime},\tilde{z},\theta.$
(20)

Moreover, for any data distribution generated by the classical vae (Eqs. 17,
18 and 19), there exists an lidvae that can generate the same distribution.

###### Proof.

We first establish the latent variable identifiability. To show that the
latent variable $z_{i}$ is identifiable, it is sufficient to show that the
mapping from $z_{i}$ to $p(x_{i}\,|\,z_{i}\,;\,\theta)$ is injective for all
$\theta$. The injectivity holds because all the transformations
$(\beta_{1},\beta_{2},g_{1,\theta},g_{2,\theta})$ involved in the mapping is
injective, and their composition must be injective: the linear transformations
$(\beta_{1},\beta_{2})$ have full row rank and hence are injective; the
nonlinear transformations $(g_{1,\theta},g_{2,\theta})$ are monotone transport
maps and are guaranteed to be bijective [1, 8]; finally, the exponential
family likelihood is injective.

We next establish the flexibility of the lidvae, by proving that any vae-
generated $p(\mathbf{\boldsymbol{x}})$ can be generated by an lidvae. The
proof proceeds in two steps: (1) we show any vae-generated
$p(\mathbf{\boldsymbol{x}})$ can be generated by a vae with injective
likelihood $p(x_{i}\,|\,z_{i}\,;\,\theta)$; (2) we show any
$p(\mathbf{\boldsymbol{x}})$ generated by an injective vae can be generated by
an lidvae.

To prove (1), suppose $\beta_{1}$ does not have full row rank and $g_{\theta}$
is not injective. Then there exists some $Z^{\prime}\in\mathbb{R}^{d}$, $d<K$,
and injective $\beta^{\prime}_{1,\theta},g^{\prime}_{\theta}$ such that the
new vae can represent the same $p(\mathbf{\boldsymbol{x}})$. The reason is
that we can always turn an non-injective function into an injective one by
considering its quotient space. In particular, we consider the quotient space
with the equivalence relation between $z,z^{\prime}$ defined as
$p(x\,|\,z\,;\,\theta)=p(x\,|\,z^{\prime}\,;\,\theta)\\}$, which induces a
bijection into $\mathbb{R}^{d}$. When $p(z^{\prime})$ is no longer standard
Gaussian, there must exist a bijective Brenier map
$\tilde{z}=f_{t}(z^{\prime})$ such that $p(\tilde{x})$ is standard Gaussian
(Theorem 6 of McCann and others [9]).

To prove (2), we show that any vae with injective mapping can be
reparameterized as a lidvae. To prove this claim, it is sufficient to show
that any injective function
$l_{\theta}:\mathbb{R}^{M+H}\rightarrow\mathbb{R}^{D}$ can be reparametrized
as $g_{2,\theta}(\beta_{2}^{\top}\,g_{1,\theta}(\cdot))$. Below we provide
such a reparametrization by solving for $g_{1},g_{2}$ and $\beta$ in
$l_{\theta}(z)=g_{2,\theta}(\beta_{2}^{\top}g_{1,\theta}(z))$. We set
$g_{1,\theta}$ as an identity map, $\beta_{2}$ as an $(M+H)\times D$ matrix
with all the main diagonal entries being one and all other entries being zero,
and $g_{2,\theta}$ as an invertible $\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$
mapping which coincides with $l_{\theta}$ on the $(M+H)$-dimensional subspace
of $z$.

Finally, we note that the same argument applies to the variant of vae where
$w_{i}=z_{i}$. It coincides with the classical vae in Kingma and Welling [6].
Applying the same argument as above establishes Proposition 2.

∎

##  Appendix C Experiment details

For image experiments, all hidden layers of the neural networks have 512
units. We choose the number of continuous latent variables as 64 and the
dimensionality of categorical variables as the number of ground truth labels.
Then we use two-layer RealNVP ([2]) as an approximating family to tease out
the effect of variational inference.

For text experiments, all hidden layers of the neural networks have 1024
units. We choose the dimensionality of the embedding as 1024. Then we use two-
layer LSTM as an approximating family following common practice of fitting
sequential vae.

##  Appendix D Additional experimental results

Table 2 includes additional experimental results of lidvae on image datasets
(Pinwheel and MNIST).

| Pinwheel |  |  | MNIST
---|---|---|---|---
| AU | KL | MI | LL |  |  | AU | KL | MI | LL
vae [6] | 0.2 | 1.4e-6 | 2.0e-3 | -6.2 (5e-2) |  |  | 0.1 | 0.1 | 0.2 | -108.2 (5e-1)
SA-vae [5] | 0.2 | 1.6e-5 | 2.0e-2 | -6.5 (5e-2) |  |  | 0.4 | 0.4 | 0.6 | -106.3 (7e-1)
Lagging vae [3] | 0.6 | 0.7e-3 | 1.5e0 | -6.5 (4e-2) |  |  | 0.5 | 0.8 | 1.7 | -105.2 (5e-1)
$\beta$-vae [4] ($\beta$=0.2) | 1.0 | 1.2e-3 | 2.3e0 | -6.6 (6e-2) |  |  | 0.8 | 1.5 | 2.8 | -100.4 (6e-1)
lidgmvae (this work) | 1.0 | 1.2e-3 | 2.2e0 | -6.5 (5e-2) |  |  | 1.0 | 1.8 | 3.9 | -95.4 (7e-1)
Table 2:  lidgmvae do not suffer from posterior collapse and achieves better
fit than its classical counterpart in a 9-layer generative model. The reported
number is mean (sd) over ten different random seeds. (Higher is better.)

## References

  * [1] K. Ball (2004) An elementary introduction to monotone transportation.  In Geometric aspects of functional analysis,  pp. 41–52.  Cited by: Appendix B.
  * [2] L. Dinh, J. Sohl-Dickstein, and S. Bengio (2016) Density estimation using real NVP.  arXiv preprint arXiv:1605.08803.  Cited by: Appendix C.
  * [3] J. He, D. Spokoyny, G. Neubig, and T. Berg-Kirkpatrick (2019) Lagging inference networks and posterior collapse in variational autoencoders.  arXiv preprint arXiv:1901.05534.  Cited by: Table 2.
  * [4] I. Higgins, L. Matthey, A. Pal, C. Burgess, X. Glorot, M. Botvinick, S. Mohamed, and A. Lerchner (2016) $\beta$-VAE: learning basic visual concepts with a constrained variational framework.  Cited by: Table 2.
  * [5] Y. Kim, S. Wiseman, A. Miller, D. Sontag, and A. Rush (2018) Semi-amortized variational autoencoders.  In International Conference on Machine Learning,  pp. 2678–2687.  Cited by: Table 2.
  * [6] D. P. Kingma and M. Welling (2014) Auto-encoding variational Bayes.  In Proceedings of the International Conference on Learning Representations (ICLR),  Vol. 1.  Cited by: Appendix B, Table 2.
  * [7] J. Lucas, G. Tucker, R. B. Grosse, and M. Norouzi (2019) Don’t blame the ELBO! A linear VAE perspective on posterior collapse.  In Advances in Neural Information Processing Systems,  pp. 9403–9413.  Cited by: §A.1.
  * [8] R. J. McCann and N. Guillen (2011) Five lectures on optimal transportation: geometry, regularity and applications.  Analysis and geometry of metric measure spaces: Lecture notes of the séminaire de Mathématiques Supérieure (SMS) Montréal,  pp. 145–180.  Cited by: Appendix B.
  * [9] R. J. McCann et al. (1995) Existence and uniqueness of monotone measure-preserving maps.  Duke Mathematical Journal 80 (2),  pp. 309–324.  Cited by: Appendix B.