V ARIATIONAL AUTOENCODERS IN THE PRESENCE OF LOW - DIMENSIONAL DATA : LANDSCAPE AND IMPLICIT BIAS
Anonymous authors
Paper under double-blind review
A BSTRACT
Variational Autoencoders (VAEs) are one of the most commonly used generative models, particularly for image data.
A prominent difﬁculty in training VAEs is data that is supported on a lower dimensional manifold.
Recent work by Dai and
Wipf (2019) suggests that on low-dimensional data, the generator will converge to a solution with 0 variance which is correctly supported on the ground truth manifold.
In this paper, via a combination of theoretical and empirical results, we show that the story is more subtle.
Precisely, we show that for linear encoders/decoders, the story is mostly true and VAE training does recover a generator with support equal to the ground truth manifold, but this is due to the implicit bias of gradient descent rather than merely the VAE loss itself.
In the nonlinear case, we show that the VAE training frequently learns a higher-dimensional manifold which is a superset of the ground truth manifold.
I NTRODUCTION
Variational autoencoders (VAEs) have recently enjoyed a revived interest, both due to architectural choices that have led to improvements in sample quality (Oord et al., 2017; Razavi et al., 2019b;
Vahdat & Kautz, 2020) and
due to algorithmic insights (Dai et al., 2017; Dai & Wipf, 2019).
Nevertheless, ﬁne-grained understanding of the behavior of VAEs is lacking, both on the theoretical and empirical level.
In our paper, we study a common setting of interest where the data is supported on a low-dimensional manifold — often argued to be the setting relevant to real-world image and text data due to the manifold hypothesis (see e.g. Goodfellow et al.
In this setting, Dai & Wipf (2019) proposed a two-stage training process for VAEs, based on a combination of empirical and theoretical arguments suggesting that for standard VAE training with such data distributions: (1) the generator’s covariance will converge to
, (2) the generator will learn the correct manifold, but not the correct density on the manifold (3) the number of approximately 0 eigenvalues in the encoder covariance will equal the intrinsic dimensionality of the manifold (see also Dai et al.
In this paper we revisit this setting and explore the behaviour of both the VAE loss and the training dynamics.
Through a combination of theory and experiments we show that:
• In the case of the data manifold being linear (i.e. the data is Gaussian, supported on a linear subspace—equivalently, it is produced as the pushforward of a Gaussian through a linear map), and the encoder/decoder being parametrized as linear maps, we show that: a) the set of optima includes parameters for which the generator’s support is a strict superset of the data manifold; b) the gradient descent dynamics are such that they converge to generators with support equal to the support of the data manifold.
This provides a full proof of the conjecture in Dai & Wipf (2019), albeit we show the phenomenon is a combination of both the location of the minima of the loss as well as an implicit bias of the training dynamics .
• In the case of the data manifold being nonlinear (i.e. the data distribution is the pushforward of the Gaussian through a nonlinear map f : R r
→ R d , r ≤ d ), the gradient descent dynamics from a random start often converges to generators G whose support strictly contains the support of the underlying data distribution, while driving reconstruction error to 0 and driving the VAE loss to
−∞ .
This shows that the conjecture in Dai & Wipf (2019) does not hold for general nonlinear data manifolds and architectures for the generator/encoder.
Organization: We will provide an informal overview of our ﬁndings in Section 3.
The rigorous discussion on the VAE landscape are in Section 4 and on the implicit bias of gradient descent in
Section 5.
S ETUP
We will study the behavior of VAE learning when data lies on a low-dimensional manifold—more precisely, we study when the generator can recover the support of the underlying data distribution.
In order to have a well-deﬁned “ground truth”, both for our theoretical and empirical results, we will consider synthetic dataset that are generated by a “ground truth” generator as follows.
Data distribution: To generate a sample point x for the data distribution, we will sample z ∼
N (0 , I r ) , and output x = f ( z ) , for a suitably chosen f .
In the linear case, f ( z ) =
Az , for some matrix A ∈ R d × r .
In the nonlinear case, f ( z ) will be a nonlinear function f : R r → R d .
We will consider several choices for f .
Parameterization of the trained model: For the model we are training, the generator will sample z ∼ N (0 , I r ) and output x ∼ N ( f ( z ) , (cid:15)I ) , for trainable f, (cid:15) ; the encoder given input x will output z ∼ N ( g ( x ) , D ) , where D ∈ R r × r is a diagonal matrix, and g, D are trainable.
In the linear case, f, g will be parameterized as matrices ˜ A, ˜ B ; in the nonlinear case, several different parameterizations will be considered.
In either case the VAE Loss will be denoted by L ( · ) , see (3).
O UR R ESULTS
Linear VAEs: the correct distribution is not recovered.
Recall in the linear case, we train a linear encoder and decoder to learn a Gaussian distribution consisting of data points x ∼ N (0 , Σ) — equivalently, the data distribution is the pushforward of a standard Gaussian z ∼ N (0 , I r ) through a linear generator x = Az with AA T = Σ ; see also Section 2 above.
In Theorem 1 of Lucas et al.
(2019), the authors proved that in a certain probabilistic PCA setting where Σ is full-rank, the landscape of the VAE loss has no spurious local minima: at any global minima of the loss, the VAE decoder exactly matches the ground truth distribution, i.e. ˜ A ˜ A T + (cid:15) 2 I = Σ .
We revisit this problem in the setting where Σ has rank less than d so that the data lies on the lowerdimensional manifold/subspace spanned by the columns of A or equivalently Σ , denoted span ( A ) .
We show empirically (i.e. via simulations) in Section 6 that when Σ is rank-degenerate the VAE actually fails to recover the correct distribution .
More precisely, the recovered ˜ A has the correct column span but fails to recover the correct density — conﬁrming predictions made in Dai & Wipf (2019).
We then explain theoretically why this happens, where it turns out we ﬁnd some surprises.
Landscape Analysis: Linear and Nonlinear VAE.
Dai & Wipf (2019) made their predictions on the basis of the following observation about the loss landscape : there can exist sequences of VAE solutions whose objective value approaches −∞ (i.e. are asymptotic global minima), for which the generator has the correct column span, but does not recover the correct density on the subspace.
They also informally argued that these are all of the asymptotic global minima of loss landscape (Pg 7 and Appendix I in Dai & Wipf (2019)), but did not give a formal theorem or proof of this claim.
We settle the question by showing this is not the case : 1 namely, there exist many convergent sequences of VAE solutions which still go to objective value −∞ but also do not recover the correct column span — instead, the span of such ˜ A is a strictly larger subspace.
More precisely, we obtain a tight characterization of all asymptotic global minima of the loss landscape:
Theorem 1 (Optima of Linear VAE Loss, Informal Version of Theorem 3) .
Suppose that ˜ A, ˜ B are ﬁxed matrices such that A =
˜
A ˜ BA and suppose that # { i : ˜
A i = 0 } > r − d , i.e. the number of zero columns of ˜ A is strictly larger than r − d .
Then there exists ˜ (cid:15) t → 0 and positive diagonal matrices
˜ D t such that lim t →∞ L ( ˜ A, ˜ B, ˜ D t , ˜ (cid:15) t ) =
−∞ .
Also, these are all of the asymptotic global minima: any convergent sequence of points ( ˜ A t , ˜ B t , ˜ D t , ˜ (cid:15) t ) along which the loss L goes to −∞ satisﬁes ˜
A t → ˜ A, ˜ B t → ˜ B with A = ˜
A ˜ BA such that # { i : ˜
A i = 0 } > r − d .
To interpret the constraint # { i : ˜
A i = 0 } > r − d , observe that if the data lies on a lowerdimensional subspace of dimension r ∗ < d (i.e. r ∗ is the rank of Σ ), then there exists a generator which generates the distribution with r − r ∗ > r − d zero columns by taking an arbitrary lowrank factorization LL T = Σ with L : d × r ∗ and deﬁning A : d × r by A =
[ L 0 d × r − r ∗ ] .
The larger the gap is between the manifold/intrinsic dimension r ∗ and the ambient dimension d , the more ﬂexibility we have in constructing asymptotic global minima of the landscape.
Also, we note there is no constraint in the Theorem that r − d ≥ 0 : the assumption is automatically satisﬁed if r < d .
To summarize, the asymptotic global minima satisfy A =
˜
A ˜ BA , so the column span of ˜ A contains that of A , but in general it can be a higher dimensional space.
For example
, if d, r ≥ r ∗ + 2 and and the ground truth generator is A = (cid:20) I r ∗   (cid:21) , then ˜ A = (cid:20) I r ∗ +    (cid:21) and ˜ B =
(cid:20) I r ∗ +    (cid:21) is a perfectly valid asymptotic global optima of the landscape,
even though decoder ˜ A generates a different higher-dimensional Gaussian distribution N (cid:18) 0 , (cid:20) I r ∗ +    (cid:21)(cid:19) than the ground truth.
In the above result we showed that there are asymptotic global minima with higher dimensional spans even with the common restriction that the encoder variance is diagonal; if we considered the case where the encoder variance is unconstrained, as done in Dai & Wipf (2019), and/or can depend on its input x , this can only increase the number of ways to drive the objective value to −∞ .
We also consider the analogous question in the nonlinear VAE setting where the data lies on a lowdimensional manifold.
We prove in Theorem 6 that even in a very simple example where we ﬁt a
VAE to generate data produced by a 1-layer ground truth generator, there exists a bad solution with strictly larger manifold dimension which drives the reconstruction error to zero (and VAE loss to −∞ ).
The proof of this result does not depend strongly on the details of this setup and it can be adapted to construct bad solutions for other nonlinear VAE settings.
We note that the nature both of these result is asymptotic: that is, they consider sequences of solutions whose loss converges to −∞ — but not the rate at which they do so.
In the next section, we will consider the trajectories the optimization algorithm takes, when the loss is minimized through gradient descent.
Linear VAE: implicit regularization of gradient ﬂow.
The above theorem indicates that studying the minima of the loss landscape alone cannot explain the empirical phenomenon of linear VAE training recovering the support of the ground truth manifold in experiments; the only prediction that can be made is that the VAE will recover a possibly larger manifold containing the data.
We resolve this tension by proving that gradient ﬂow , the continuous time version of gradient descent, has an implicit bias towards the low-rank global optima.
Precisely, we measure the effective rank quantitatively in terms of the singular values: namely, if σ k ( ˜ A ) denotes the k -th largest singular value of matrix ˜ A , we show that all but the largest dim(span A ) singular values of ˜ A decay at an exponential rate, as long as: (1) the gradient ﬂow continues to exist 2 , and (2) the gradient ﬂow does not go off to inﬁnity, i.e. neither ˜ A or ˜ (cid:15) go to inﬁnity (in simulations, the decoder ˜
A converges to a bounded point and ˜ (cid:15) → 0 so the latter assumption is true).
To simplify the proof, we work with a slightly modiﬁed loss which “eliminates” the encoder variance by setting it to its optimal value:
L 1 ( ˜ A, ˜ B, ˜ (cid:15) ) :
= min ˜ D L ( ˜ A, ˜ B, ˜ (cid:15), ˜ D ) ; this loss has a simpler closed form, and we believe the theorems should hold for the standard loss as well.
(Generally, gradient descent on the original loss L will try to optimize ˜ D in terms of the other parameters, and if it succeeds to keep ˜ D well-tuned in terms of ˜ A, ˜ B, ˜ (cid:15) then L will behave like the simpliﬁed loss L 1 .)
Theorem 2 (Implicit Bias of Gradient Flow, Informal version of Theorem 5) .
Let A : d × r be arbitrary and deﬁne W to be the span of the rows of A , let ˜Θ(0) =
( ˜ A (0) , ˜ B (0) , ˜ (cid:15) (0)) be an arbitrary initialization and deﬁne the gradient ﬂow ˜Θ( t ) =
( ˜ A ( t ) , ˜ B ( t ) , ˜ (cid:15) ( t )) by the ordinary differential equation (ODE)
[Equation]
with initial condition Θ 0 .
If the solution to this equation exists on the time interval [0 , T ] and satisﬁes max t ∈ [0 ,T ] max j [ (cid:107) ( ˜ A t ) j (cid:107)
+ ˜ (cid:15) 2 t ] ≤ K , then for all t ∈ [0 , T ] we have
[Equation]
where C ( A, ˜ A ) :
= (cid:107) P W ⊥
˜
A T (0) (cid:107) 2 F and P W ⊥ is the orthogonal projection onto the orthogonal complement of W .
Together, our Theorem 1 and Theorem 2 show that if gradient descent converges to a point while driving the loss to −∞ , then it successfully recovers the ground truth subspace/manifold span A .
This shows that, in the linear case, the conjecture of Dai & Wipf (2019) can indeed be validated provided we incorporate training dynamics into the picture.
The prediction of theirs we do not prove is that the number of zero entries of the encoder covariance D converges to the intrinsic dimension; as shown in Table 1, in a few experimental runs this does not occur — in contrast, Theorem 2 implies that ˜
A should have the right number of nonzero singular values and our experiments agree with this.
Nonlinear VAE: Dynamics and Simulations.
In the linear case, we showed that the implicit bias of gradient descent leads the VAE training to converge to a distribution with the correct support.
In the nonlinear case, we show that this does not happen—even in simple cases.
Precisely, in the setup of the one-layer ground truth generator, where we proved (Theorem 6) there exist bad global optima of the landscape, we verify experimentally (see Figure 1) that gradient descent from a random start does indeed converge to such bad asymptotic minima.
In particular, this shows that whether or not gradient descent has a favorable implicit bias strongly depends on the data distribution and VAE architecture.
More generally, by performing experiments with synthetic data of known manifold dimension, we make the following conclusions: (1) gradient descent training recovers manifolds approximately containing the data, (2) these manifolds are generally not the same dimension as the ground truth manifold, but larger (this is in contrast to the conjecture in Dai & Wipf (2019) that they should be equal) even when the decoder and encoder are large enough to represent the ground truth and the reconstruction error is driven to 0 (VAE loss is driven to −∞ ), and (3) of all manifolds containing the data, gradient descent still seems to favor those with relatively low (but not always minimal) dimension.
Further investigating the precise role of VAE architecture and optimization algorithm, as well as the interplay with the data distribution is an exciting direction for future work.
3. R ELATED WORK
Implicit regularization.
Interestingly, the implicit bias towards low-rank solutions in the VAE which we discover is consistent with theoretical and experimental results in other settings, such as deep linear networks/matrix factorization (e.g. Gunasekar et al.
; Li et al.
Arora et al.
; Li et al.
Jacot et al.
, although it seems to be for a different mathematical reason — unlike those settings, initialization scale does not play a major role.
Similar to the setting of implicit margin maximization (see e.g. Ji & Telgarsky (2018); Schapire & Freund (2013);
Soudry et al.
(2018)), in our VAE setting the optima are asymptotic (though approaching a ﬁnite point, not off at inﬁnity) and the loss goes to −∞ .
Kumar & Poole (2020); Tang & Yang (2021) also explore some implicit regularization effects tied to the Jacobian of the generator and the covariance of the
Gaussian noise.
Architectural and Algorithmic Advances for VAEs.
There has been a recent surge in activity with the goal of understanding VAE training and improving its performance in practice.
Much of the work has been motivated by improving posterior modeling to avoid problems such as “posterior collapse”, see e.g. (Dai et al., 2020;
Razavi et al., 2019a; Pervez & Gavves, 2021; Lucas et al., 2019;
He et al., 2019; Oord et al., 2017; Razavi et al., 2019b; Vahdat & Kautz, 2020).
Most relevant to the current work are probably the works Dai & Wipf (2019) and Lucas et al.
(2019) discussed earlier.
A relevant previous work to these is Dai et al.
(2017); one connection to the current work is that they also performed experiments with a ground truth manifold, in their case given as the pushforward of a Gaussian through a ReLU network.
In their case, they found that for a certain decoder and encoder architectures that they could recover the intrinsic dimension using a heuristic related to the encoder covariance eigenvalues from Dai & Wipf (2019); our results are complementary in that they show that this phenomena is not universal and does not hold for other natural datasets (e.g. manifold data on a sphere ﬁt with a standard VAE architecture).
VAE L ANDSCAPE A NALYSIS
In this section, we analyze the landscape of a VAE, both in the linear and non-linear case.
Preliminaries and notation.
We use a VAE to model a datapoint x ∈ R d as the pushforward of z ∼ N (0 , I r ) .
We have the following standard VAE architecture:
[Equation]
where (cid:15) 2 > 0 is the decoder variance, D is a diagonal matrix with nonnegative entries, and f, g, D, (cid:15) are all trainable parameters.
(For simplicity, our D does not depend on x ; this is the most common setup in the linear VAE case we will primarily focus on.)
The VAE objective (see Appendix A for explicit derivation) is to minimize :
[Equation]
We also state a general fact about VAEs for the case that the objective value can be driven to −∞ , which was observed in (Dai & Wipf, 2019): they must satisfy (cid:15) → 0 and achieve perfect limiting reconstruction error.
The ﬁrst claim in this Lemma is established in the proof of Theorem 4 and the second claim is Theorem 5 in Dai & Wipf (2019).
For completeness, we include a self-contained proof in Appendix B. [Equation]
Lemma 1 (Theorems 4 and 5 of Dai & Wipf (2019)) .
Suppose f t , g t , D t , (cid:15) t for t ≥ 1 are a sequence such that lim t →∞ L ( f t , g t , D t , (cid:15) t ) =
−∞ .
Then: 1) lim t →∞ (cid:15) t = 0 and 2)
[Equation]
In fact [Equation] the reconstruction error and (cid:15) are closely linked in a simple way:
Lemma 2.
If f, g, D are ﬁxed, then the optimal value of (cid:15) to minimize L is given by (cid:15)
=
(cid:113)
[Equation]
4. L INEAR VAE
Setup:
In the linear VAE case, we assume the data is generated from the model x = Az with
A ∈ R d × r ∗ and z ∼ N (0 , I r ∗ ) .
We will denote the training parameters by ˜ A ∈ R d × r , ˜ B ∈ R r × d , ˜ D ∈ R r × r , and ˜ (cid:15) > 0 , where r ≥ 1 is a ﬁxed hyperparameter which corresponds to the latent dimension in the trained generator, and we assume ˜ D is a diagonal matrix.
With this notation in place, the implied VAE has generator/decoder ˜ x ∼ N ( ˜ Az, ˜ (cid:15) 2 I d ) and encoder ˜
z ∼ N ( ˜ Bx, ˜ D ) .
The
VAE objective as a function of parameters ˜Θ =
( ˜ A, ˜ B, ˜ D, ˜ (cid:15) ) is (see Appendix A):
[Equation]
Our analysis makes use of a simpliﬁed objective L 1 , which “eliminates” D out of the objective by plugging in the optimal value of D for a choice of the other variables.
We use this as a technical tool when analyzing the landscape of the original loss L .
Lemma 3 (Deriving the simpliﬁed loss L 1 ) .
Suppose that ˜ A, ˜ B, ˜ (cid:15) are ﬁxed.
Then the objective
L is minimized by choosing for all i that ˜ D ii =
[Equation]
2 where ˜ A i is column i of ˜ A , and for
L 1 ( ˜ A, ˜ B, ˜ (cid:15) ) :
= min ˜ D L ( ˜ A, ˜ B, ˜ D, ˜ (cid:15) )
it holds that
[Equation]
Taking advantage of this simpliﬁed formula, we can then identify (for the original loss L ) simple sufﬁcient conditions on ˜ A, ˜ B which ensure they can be used to approach the population loss minimum by picking suitable ˜ (cid:15) t , ˜ D t and prove matching necessary conditions.
Theorem 3.
First, suppose that
˜ A : d × r, ˜ B : r × d are ﬁxed matrices such that A
= ˜
A ˜ BA and suppose that # { i : ˜
A i = 0 } > r − d , i.e. the number of zero columns of ˜ A is strictly larger than r − d .
Then for any sequence of positive ˜ (cid:15) t → 0 there exist a sequence of positive diagonal matrices ˜
D t such that:
For every i such that ˜
A i (cid:54) = 0 , i.e. column i of ˜
A is nonzero, we have ( ˜ D t ) ii → 0 . 2.
lim t →∞ L ( ˜ A, ˜ B, ˜ D t , ˜
(cid:15) t ) =
−∞ .
Conversely, suppose that that ˜ A t , ˜ B t , ˜ D t , ˜ (cid:15) t is an arbitrary sequence such that lim t →∞ L ( ˜ A t , ˜ B t , ˜ D t , ˜ (cid:15) t ) =
−∞ .
Then as t → ∞ , we must have that:
˜ (cid:15) t → 0 and (cid:107) A − ˜
A t ˜ B t A (cid:107) 2 F → 0 . 2.
max i ( ˜ D t ) ii (cid:107) ( ˜ A t ) i (cid:107) 2 F → 0 where ( ˜ A t ) i denotes the i -th column of ˜ A t . 3.
For any δ > 0 , lim inf t →∞ # { i : (cid:107) ( ˜ A t ) i (cid:107) 22 < δ }
> r − d , i
e. asymptotically ˜
A t has strictly more than r − d columns arbitrarily close to zero
In particular, if ( ˜ A t , ˜ B t , ˜ D t , ˜ (cid:15) t ) converge to a point ( ˜ A, ˜ B, ˜ D, ˜ (cid:15) ) then ˜ (cid:15) = 0 , A
= ˜
A ˜ BA , ˜ D ii =
for every i such that ˜
A i (cid:54) = 0 , and # { i : ˜
A i = 0 } > r − d .
For the sufﬁciency direction, we observe that in (5) the ﬁrst term is zero if A =
˜
A ˜ BA and the sum of the last two terms goes to −∞
if ˜ (cid:15) → 0 and enough columns of ˜ A are zero.
Based upon similar reasoning combined with Lemma 1, we show necessity.
The full proof is in the Appendix.
4. N ONLINEAR VAE
In this section, we give a simple example of a nonlinear VAE architecture which can represent the ground truth distribution perfectly, but has another asymptotic global minimum where it outputs data lying on a manifold of a larger dimension ( r ∗ + s instead of r ∗ for any s ≥ 1 ).
The ground truth model is a one-layer network (“sigmoid dataset” in Section 6) and the bad decoder we construct outputs a standard Gaussian in r ∗ + s dimensions padded with zeros.
Theorem 4 (Theorem 6 in Appendix) .
Let s ≥ 1 be arbitrary and consider the sigmoid setup from Section 6.
There exists ˜
A 1 , ˜
A 2 , ˜ B s.t.
for ˜ (cid:15) t → 0 there exists
˜ D
t s.t.
(1) the VAE loss
L ( ˜ A 1 , ˜
A 2 , ˜ B, ˜ D t , ˜ (cid:15) t ) →
−∞ ; (2) The output of the decoder ˜ x = ˜
A 1 ˜ z
+ σ ( ˜ A 2 ˜ z ) , ˜ z ∼ N (0 , I r ) is a standard Gaussian in the ﬁrst r ∗ + s coordinates and zero in the remaining ones.
Thus, the generator constructed has as support a manifold of larger dimension ( r ∗ + s ).
Moreover, in Section 6, we show that this is not merely a theoretical possibility: we show through simulations that gradient descent from a random initialization often converges to similar minima.
I MPLICIT BIAS OF GRADIENT DESCENT IN L INEAR VAE
In this section, we prove that even though the landscape of the VAE loss contains generators with strictly larger support than the ground truth, the gradient ﬂow is implicitly biased towards low-rank solutions .
We prove this for the simpliﬁed loss L 1 ( ˜ A, ˜ B, ˜ (cid:15) )
=
min ˜ D L 1 ( ˜ A, ˜ B, ˜ (cid:15), ˜ D ) , which makes the calculations more tractable, though we believe our results should hold for the original loss L as well.
The main result we prove is as follows:
Theorem 5 (Implicit bias of gradient descent) .
Let A : d × r be arbitrary and deﬁne W to be the span of the rows of A , let ˜Θ(0) =
( ˜ A (0) , ˜ B (0) , ˜ (cid:15) (0)) be an arbitrary initialization and deﬁne the gradient ﬂow ˜Θ( t ) =
( ˜ A ( t ) , ˜ B ( t ) , ˜ (cid:15) ( t )) by the differential equation (1) .
with initial condition ˜Θ 0 .
If the solution to this equation exists on the time interval [0 , T ] and satisﬁes max t ∈ [0 ,T ] max j [ (cid:107) ( ˜ A t ) j (cid:107)
+ ˜ (cid:15) 2 t ] ≤ K , then for all t ∈ [0 , T ] we have
[Equation]
(6) where P W ⊥ is the orthogonal projection onto the orthogonal complement of W .
Towards showing the above result, we ﬁrst show how to reduce to matrices where A has d − dim ( rowspan ( A )) rows that are all-zero.
To do this, we observe that the linear VAE objective is invariant to arbitrary rotations in the output space (i.e. x -space), so the gradient descent/ﬂow trajectories transform naturally under rotations.
Thus, we can “rotate” the ground truth parameters as well as the training parameters.
This is formally captured as Lemma 6 in the Appendix.
Recall that by the singular value decomposition A = USV T for some orthogonal matrices U, V and diagonal matrix S , and rotation invariance in the x -space lets us reduce to analyzing the case where U =
I , i.e. A = SV T .
This matrix has a zero row for every zero singular value.
Analysis when A has zero rows.
Having reduced our analysis to the case where A has zero rows, the following key lemma shows that for every i such that row i of A (denoted A ( i ) ) is zero, the gradient descent step −∇ L or −∇ L 1 will be negatively correlated with the corresponding row
˜ A ( i ) .
Lemma 4 (Gradient correlation) .
If row i of A is zero, then
[Equation]
The way we use it is to notice that since the negative gradient points towards zero, gradient descent will shrink the size of ˜ A ( i ) .
Since the size of the matrix ˜
A stays bounded, this should mean that for small step sizes the norm of row i of ˜
A shrinks by a constant factor at every step of gradient descent on loss L 1 .
We formalize this in continuous time for the gradient ﬂow, i.e. the limit of gradient descent as step size goes to zero: for the special case of Theorem 2 in the zero row setting, the corresponding rows of ˜ A decay exponentially fast.
Lemma 5 (Exponential decay of extra rows) .
Let A be arbitrary, and let ˜Θ(0) =
( ˜ A (0) , ˜ B (0) , ˜ (cid:15) (0)) be an arbitrary initialization and deﬁne the gradient ﬂow ˜Θ( t ) =
( ˜ A ( t ) , ˜ B ( t ) , ˜ (cid:15) ( t )) to be a solution of the differential equation (1) with initial condition ˜Θ(0) .
If the solution exists on the time interval [0 , T ] and satisﬁes max t ∈ [0 ,T ] max j [ (cid:107) ( ˜ A ( t ))
j (cid:107)
+ ˜ (cid:15) ( t ) 2 ] ≤ K for some K > 0 , then for all i such that row i of A is zero we have (cid:107) ˜
A ( i ) ( t ) (cid:107) 2 ≤ e − t/K (cid:107) ˜
A ( i ) (0) (cid:107) 2 for all t ∈ [0 , T ] .
S IMULATIONS
In this section, we provide extensive empirical support for the questions we addressed theoretically.
In particular we investigate the kinds of minima VAEs converge to when optimized via gradient descent over the course of training.
Linear VAEs: First, we investigate whether linear VAEs are able to ﬁnd the correct support for a distribution supported over a linear subspace.
The setup is as follows.
We choose a ground truth linear transformation matrix A by concatenating an r ∗ × r ∗ matrix consisting of iid standard
Gaussian entries with a zero matrix of dimension ( d − r ∗ ) × r ∗ ; the data is generated as Az, z ∼
N (0 , I r ∗ ) .
Thus the data lies in a r ∗ -dimensional subspace embedded in a d -dimensional space.
[Table] We ran the experiment with various choices for r ∗ , d as well as the latent dimension of the trained decoder (Table 1).
Every result is the mean over three experiments run with the same dimensionality and setup but a different random seed.
Results: From Table 1 we can see that the optima found by gradient descent capture the support of the manifold accurately across all choices of d, r , with the correct number of nonzero decoder rows.
We also almost always see the correct number of zero dimensions in the diagonal matrix corresponding to the encoder variance.
However, gradient descent is unable to recover the density of the data on the learned manifold in the linear setting — in sharp contrast to the full rank case (Lucas et al., 2019).
We conclude this by comparing the eigenvalues of the data covariance matrix and the learned generator covariance matrix.
In order to understand whether the distribution on the linear subspace has the right density, we compute the eigenvalue error by forming matrices X, ˆ X with n rows, for which each row is sampled from the ground truth and learned generator distribution respectively.
We then compute the vector of eigenvalues λ, ˆ λ for the ground truth covariance matrix AA T and empirical covariance matrix (1 /n )
ˆ X T
ˆ X respectively and compute the normalized eigenvalue error
|| ˆ λ −
λ ||
/ ||
λ
|| .
In no case does the density of the learned distribution come close to the ground truth.
Nonlinear Dataset
In this section, we investigate whether VAEs are able to ﬁnd the correct support in nonlinear settings.
Unlike the linear setting, there is no “canonical” data distribution suited for a nonlinear VAE, so we explore two setups:
• Sphere dataset: The data are generated from the unit sphere concatenated with zero padding at the end.
This can be interpreted as a unit sphere embedded in a higher dimensional space.
We usedlayers of 200 hidden units to parameterize our encoder and decoder networks.
To measure how well the VAE has learnt the support of the distribution
we evaluate the average of ( (cid:107) ˜ x
:( r +1) (cid:107) 2 − 1) 2 , where ˜ x are generated by the learnt generator.
We will call this quantity manifold error .
We have also evaluated the padding error , which is deﬁned as (cid:107) ˜ x
r +2: (cid:107) 22 .
• Sigmoid Dataset: Let z ∼ N (0 , I r ) , the sigmoid dataset concatenates z with σ ( (cid:104) a ∗ , z (cid:105) ) where a ∗ ∈ R r is generated according to N (0 , I r ) .
We add additional zero paddings to embed the generated data in a higher dimensional ambient space.
The decoder is parameterized by a nonlinear function f ( z ) =
˜ Az + σ ( ˜ Cz ) and the encoder is parameterized by a linear function g ( x ) = ˜
Bx .
The intrinsic dimension of the dataset is r .
To measure how well the VAE has learnt the support of the distribution, we evaluate the average of ( σ ( (cid:104) a ∗ , ˜ x : r (cid:105) )
− ˜ x
r +1 ) 2 , where ˜ x are generated by the learnt decoder.
We will call this quantity manifold error .
The padding error is deﬁned as similarly as the sphere dataset.
Results: Due to space constraints, results for the sphere dataset (Table 3 and Figure 3) are in Appendix D.
In both of the nonlinear dataset experiments, we see that the number of zero entries in the diagonal encoder variance is less reﬂective of the intrinsic dimension of the manifold than the linear dataset.
It is, however, at least as large as the intrinsic dimension (Table 3, 2).
We consider a coordinate to be 0 if it’s less than 0.1.
We found that the magnitude of each coordinate to be well separated, i.e. the smaller coordinates tend to be smaller than 0.1 and the larger tend to be bigger than 0.5.
Thus the threshold selection is not crucial.
We did not include padding error in the tables because it reaches zero in all experiments
[Table]
Dimensions
Ambient Dimensions
VAE Latent Dimensions
[Table]
Error 0. Mean #0’s in Encoder Variance 3.
We show the progression of manifold error, decoder variance and VAE loss during training for the sphere data in Figure 3 and for the sigmoid data in Figure 2, which are included in Appendix D due to space constraints.
Datasets of all conﬁgurations of dimensions reached close to zero decoder variances, meaning the VAE loss is approaching −∞ .
To demonstrate Theorem 6, we took examples from both datasets to visualize their output.
For the sphere dataset, we visualize the data generated from the model, with 8 latent dimensions, trained on unit sphere with 2 intrinsic dimensions and 16 ambient dimensions (Column 2 in Table 3).
Its training progression is shown as the orange curve in Figure 3 .
We create a histogram of the norm of its ﬁrst 3 dimensions (Figure 1 (a)) and found that more than half of the generated data falls inside of the unit sphere.
The generated data has one intrinsic dimension higher than its training data, despite its decoder variance approaching zero, which is equivalent to its reconstruction error approaching zero by Lemma 2.
In the sigmoid dataset, the featured model has 24 latent dimension, and is trained on a 7-dimensional manifold embedded in a 28-dimensional ambient space.
We produced a scatter plot givengenerated data points ˜ x
r +1 from the decoder.
The x -axis in the Figure 1(b) is (cid:104) a ∗ , ˜ x : r (cid:105) and the y -axis is ˜ x
r +1 .
In contrast to the groundtruth data, whose scatter points roughly form a sigmoid function, the scatter points of the generate data resemble a gaussian distribution.
This closely resembles the example provided in Theorem 6.
Hence, despite its decoder variance and reconstruction error both approaching zero and loss consistently decreasing, the generated data do not recover the training data distribution and the data distribution recovered has higher intrinsic dimensions than the training data.
We also investigated the effect of lower bounding the decoder variance as a possible way to improve the VAE performance (details are given in Appendix E).
This enabled the VAE to recover the correct manifold dimension in the sigmoid example, but not the sphere example; methods of improvements to the VAE’s manifold recovery is an important direction for future work.
A D ERIVATIONS OF VAE LOSSES
We have (for some constants C 1 , C 2 , C 3 ):
[Equation]
where the ﬁrst line uses that log √ det (cid:15)
I = log √ (cid:15)
d = d log( (cid:15) ) .
The VAE objective is to maximize the expectation of log p ( x | z ) + log p ( z ) − log q ( z | x ) for x from the data p ∗ and z ∼ q ( z | x ) .
This means that explicitly the objective is to maximize
[Equation]
which simpliﬁes to (up to additive constant)
[Equation]
i
[Equation]
and converting this to minimization form gives the VAE Loss (3)
[Equation]
Linear VAE derivation.
Plugging in the linear VAE parameters into the loss function, we get
[Equation]
We can write out the expectation as:
[Equation]
where we used that z, z (cid:48) are independent and the identity E z ∼ N (0 ,I ) (cid:107)
Mz (cid:107) 2 = (cid:104) MM T , I (cid:105) =
(cid:107) M (cid:107) 2 F .
Next, we can observe that
[Equation]
where ˜ A i is the i -th column of the matrix ˜ A .
Therefore we recover (4).
B D EFERRED P ROOFS FROM S ECTIONB. G ENERAL F ACTS
Proof of Lemma 1.
For completeness, we include the proof of these claims; they are similar to the proofs of Theorems 4 and 5 in Dai & Wipf (2019).
First, consider the objective for ﬁxed f, g, D, (cid:15) and omit the subscript t .
We have
[Equation]
and
[Equation]
from the inequality x − log x ≥ 1 for x ≥ 0 .
Since these terms are both bounded above, the only way the objective goes to negative inﬁnity is if d log (cid:15) →
−∞ which means (cid:15) → 0 .
Now that we know (cid:15) t → 0 , we claim that lim t
[Equation]
E
[Equation]
Suppose otherwise: then this for inﬁnitely many t this quantity is lower bounded by some constant c > 0 , hence the objective for those t is lower bounded by c/(cid:15)
+ d log( (cid:15) )
+ r/
and this goes to + ∞ as (cid:15) → 0 , instead of −∞ .
Proof of Lemma 2.
Taking the partial derivative of (3) with respect to (cid:15) and setting it to zero gives
[Equation]
[Equation]
and solving for (cid:15) gives the result.
B. L INEAR VAE
Proof of Lemma 3.
Taking the partial derivative with respect to ˜ D ii gives 0 = (cid:107) ˜
A i (cid:107) 2 / ˜ (cid:15)
2 +1 − 1 / ˜ D ii which means
[Equation]
hence
[Equation]
It follows that the objective value at the optimal D is
[Equation]
i
[Equation] [Equation]
Proof of Theorem
First we prove the sufﬁciency direction,
i
[Equation]
e. that
if A =
˜
A ˜ BA and there exists i such that
˜
A i =
then we show how to drive the loss to −∞ .
By Lemma 3, if we make the optimal choice of D (which clearly satisﬁes the conditions on D described in the Lemma)
the objective simpliﬁes to
[Equation]
we have (1 / 2) log( (cid:107) ˜
A i (cid:107) 2 + ˜ (cid:15) 2 ) = log ˜
(cid:15) so the objective will go to −∞ provided ( d − r + # { i : ˜
A i = 0 } ) log ˜ (cid:15) →
−∞ .
Since ˜ (cid:15) → 0 this is equivalent to asking d − r + # { i : ˜
A i = 0 } > 0 , which is exactly the assumption of the Theorem.
Next we prove the converse direction, i.e. the necessary conditions.
Note: we split the ﬁrst item in the lemma into two conclusions in the proof below (so there are four conclusions instead of three).
The ﬁrst conclusion follows from the ﬁrst conclusion of Lemma 1.
The second conclusion of Lemma 1 tells us that
[Equation]
which gives us the second and third conclusions above.
For the fourth conclusion, since
L 1 ( ˜ A t , ˜ B t , ˜ D t ) ≤ L ( ˜ A t , ˜ B t , ˜ D t , ˜ (cid:15) t ) we know that lim t →∞ L 1 ( ˜ A t , ˜ B t , ˜ D t ) =
−∞ and recalling
L 1 ( ˜ A, ˜ B, ˜ (cid:15) )
= 1 2˜ (cid:15) 2 (cid:107) A −
˜
A ˜ BA (cid:107) 2 F + 1
(cid:107) ˜
BA
[Equation] [Equation]
+ ( d − r ) log ˜ (cid:15) + 1 2 (cid:88)
[Table]
[Equation]
[Table]
we see that, because the ﬁrst two terms are nonnegative, this is possible only if the sum of the last two terms goes to −∞ .
Based on similar reasoning to the sufﬁciency case, this is only possible if strictly more than r − d of the columns of ( ˜ A t ) become arbitrarily close to zero; precisely, if there exists δ such that at most r − d of the columns of ˜ A t have norm less than δ , then
[Equation]
[Table] [Equation] [Table] [Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
which does not go to −∞ as ˜ (cid:15) → 0 (and the other terms of L 1 are nonnegative).
B. N ONLINEAR VAE
We give the full details of the construction of the bad nonlinear VAE optimum and prove that it is an asymptotic global minimum.
(Note: in the notation of Section 6 we are considering a ∗ with 0/ entries, but the proof generalizes straightforwardly for arbitrary a ∗ with the correct support.)
Setup: Suppose s ≥ 1 is arbitrary and the ground truth x ∈ R d with d > r ∗ + s is generated in the following way: ( x 1 , . . .
, x r ∗ ) ∼ N (0 , I r ∗ ) , x r ∗ +1 =
σ ( x 1 + · · · + x r ∗ ) for an arbitrary nonlinearity σ , and x r ∗ +2 = · · · = x d = 0 .
Furthermore, suppose the architecture for the decoder with latent dimension r > r ∗ + 1 is
[Equation]
where σ ( · ) is applied as an entrywise nonlinearity, and the encoder is linear, g ( x ) :
=
˜ Bx .
Observe that the ground truth decoder can be expressed by taking ˜ A 2 to have a single nonzero row in position r + 1 with entries (1 , . .
[Equation]
where ˜ B is a ground truth encoder which achieves perfect reconstruction.
On the other hand, the following VAE different from the ground truth achieves perfect reconstruction:
[Equation]
The output of this decoder is a Gaussian N (cid:18) 0 , (cid:20) I r ∗ + s   (cid:21)(cid:19) , which means it is strictly higherdimensional than the ground truth dimension r ∗ .
(This also means that if we drew the corresponding plot of to Figure 1 (b) for this model, we would get something that looks just like the experimentally obtained result.)
We prove in the Appendix that it this is an asymptotic global optima:
Theorem 6.
Let s ≥ 1 be arbitrary and the ground truth and VAE architecture is as deﬁned as above.
For any sequence ˜ (cid:15) t → 0 , there exist diagonal matrices ˜ D t such that:
the VAE loss L ( ˜ A 1 , ˜
A 2 , ˜ B, ˜ D t , ˜ (cid:15) t ) →
−∞ where
˜ A 1 , ˜ A 2 , ˜ B are deﬁned by (9) 2.
The number of coordinates of ˜ D t which go to zero equals r ∗ + s .
Proof.
We show how to pick ˜ D t as a function of ˜ (cid:15) t and that if ˜ (cid:15) t → 0 , the loss goes to −∞ .
From now on we drop the subscripts.
With these parameters, the VAE loss is
[Equation]
i
[Equation]
=
[Equation]
1 / 2˜ (cid:15)
[Equation] [Equation]
Taking the partial derivative with respect to ˜ D ii for i ≤ r ∗ + s and optimizing gives 0 = (1 / ˜ (cid:15) 2 ) + 1 − 1 / ˜ D ii
i
[Equation]
e.
[Equation]
and plugging this into the objective gives
[Equation]
Setting the remaining ˜ D ii to
[Equation]
we see that using d > r ∗ + s that the loss goes to −∞ provided ˜ (cid:15) → 0 , proving the result.
C D EFERRED P ROOFS FROM S ECTION
First, we formalize the rotation invariance of the objective.
Lemma 6 (Rotational Invariance of Gradient Descent on Linear VAE) .
Let L A ( ˜ A, ˜ B, ˜ D, ˜ (cid:15) ) denote the VAE population loss objective (4) .
Then for an arbitrary orthogonal matrix U , we have
[Equation]
A, ˜ B, ˜ D, ˜ (cid:15) )
=
L UA ( U ˜ A, ˜ BU T , ˜ D, ˜ (cid:15) ) .
Furthermore,
U ∇ ˜
A
[Equation]
A ( ˜ A, ˜ B, ˜ D, ˜
[Equation]
U
[Equation]
A L UA
[Equation]
U
[Equation]
, ˜ BU T ,
[Equation]
) and ( ∇ ˜ B L A
[Equation]
˜
[Equation] [Equation]
[Equation]
D,
[Equation] [Equation]
[Equation]
As a consequence, if for any η ≥ 0 we deﬁne ( ˜ A 1 , ˜ B 1 , ˜ D 1 , ˜ (cid:15) 1 ) =
( ˜ A, ˜ B, ˜ D, ˜ (cid:15) )
− η ∇ L A ( ˜ A, ˜ B, ˜ D, ˜ (cid:15) )
then
[Equation]
i.e. gradient descent preserves rotations by U .
The same result holds for the gradient ﬂow (i.e. continuous time gradient descent), or replacing everywhere the loss L by the simpliﬁed loss L 1 .
Proof of Lemma 6.
We give the proof for L as stated, but it is exactly the same for the simpliﬁed loss L 1 .
From the objective function (4) and U T =
U − 1 observe that
[Equation]
Then from the above and the multivariate chain rule have
[Equation]
so multiplying both sides on the left by U and using U T
=
U − 1 gives the second claim
[Equation]
and similarly
[Equation]
gives the third claim.
Then the gradient descent claim follows immediately.
Proof of Lemma 4.
First we prove the conclusion for the original loss L .
Since ( ˜ A ˜ BA ) i(cid:96)
= (cid:80) j,k ˜
A ij
˜ B jk
A k(cid:96)
we have that
[Equation]
˜ B jk A k(cid:96) (cid:33) and if we know the corresponding row i in A is zero then this simpliﬁes
to
[Equation]
BA
[Equation]


[Equation]
k(cid:96) (cid:33) which means that
[Equation]
where the notation A ( i ) denotes row i of matrix A .
Thus, for this term the gradient with respect to row ˜
A ( i ) has nonnegative dot product with row ˜
A ( i ) .
Also,
[Equation]
and so
[Equation]
which gives the ﬁrst result.
For the second result with the simpliﬁed loss L 1 , observe that
[Equation]
so
[Equation]
and the other terms in the loss behave the same in the case of L .
Including the factor of 1 / 2 from the loss function gives the result.
Proof of Lemma 5.
From Lemma 4 we have that for any such row i ,
[Equation]
which by Gronwall’s inequality implies (cid:107)
[Equation]
desired.
Proof of Theorem 5.
Before proceeding, we observe that the ﬁrst inequality in (6) is a consequence of the general min-max characterization of singular values, see e.g. Horn & Johnson (2012).
We now prove the rest of the statement.
As explained at the beginning of the section, we start by taking the Singular Value Decomposition
A = USV T where S is the diagonal matrix of singular values and U, V are orthogonal.
We assume the diagonal matrix S is sorted so that its top-left entry is the largest singular value and its bottomright is the smallest.
Note that this means the ﬁrst dim( W ) rows of U are an orthonormal basis for W .
Note that for any time t , (cid:107) P W ⊥
˜
A T ( t ) (cid:107) 2 F = (cid:80) di =dim( W )+1 (cid:107) ( U ˜ A ( t ) T ) i (cid:107) 2 because the rows ( U dim( W )+1 , . .
, U d ) are an orthonormal basis for W ⊥ .
Therefore we have that
[Equation]
, proving the result, provided we justify the middle inequality.
Deﬁne A ∗ := U T A = SV T , which has a zero row for every zero singular value of A , and apply Lemma 5 (using that the deﬁnition of K is invariant to left-multiplication of ˜ A by an orthogonal matrix) and Lemma 6 to conclude that the rows of U T ˜
A ( t ) , i.e. the columns of U ˜
A ( t ) T , corresponding to zero rows of A ∗ shrink by a factor of e − t/K .
This directly gives the desired inequality, completing the proof.
D D EFERRED F IGURES AND P LOTS FROM S ECTION
Eigenvalues of Linear Data.
As we’ve discussed, in our linear setting the VAE does not recover the ground truth data density.
Since our generative process for ground-truth data is x = Az for a matrix A and z normally distributed, we can characterize the density function by the eigenvalues of the true or estimated covariance matrix.
We give ﬁgures for the normalized error of these eigenvalues between the learned generator and the ground truth in Table 1.
A concrete example of eigenvalue mismatch for a problem with 6 nonzero dimensions is a ground-truth set of covariance eigenvalues
[Equation] [Equation] [Equation]
4] while the trained linear VAE distribution has covariance eigenvalues [Equation] [Equation] [Equation]
Here, the VAE was easily able to learn the support of the data but clearly is very off when it comes to the structure of the covariances.
E E XPERIMENTS WITH D ECODER V ARIANCE C LIPPING
As was suggested by an anonymous ICLR 2022 reviewer, one potential way to evade the results in our paper is to restrict the decoder variance from converging to 0.
In this section, we examine (empirically) the impact of clipping the decoder variance during training.
We caveat though, that our paper does not analyze the landscape of the resulting constrained optimization problem, so our results don’t imply anything about this regime.
[Figure]
[Figure] [Table] We conduct the same nonlinear experiments described in Section 6 where we ﬁt VAEs to data generated from spheres and linear sigmoid functions. The only change is to clip the decoder variance when it goes below a certain threshold.
In the ﬁgures below, the chosen threshold is e − 4 ≈ 0 .
018 , though we tried also e − 2 and e − 3 , with similar outcomes. We initialize the decoder variance with e − 3 for this set of experiments, so the optimization still can decrease it. With this change, the optimization process on the sigmoid dataset does yield encoder variances with their number of zeros reﬂective of their intrinsic dimensions as in Table 4. For the sphere experiment, this still does not happen, as in Table 5.
In fact, the model consistently recovers one more dimension than the true intrinsic dimension of the manifold and the smaller encoder variances can be as large as 0.1.
We also provide a ﬁgure (Figure 4) in the same style as Figure 1.
We see that training with a clipped decoder variance allows the model to better capture the general shape of the sigmoid function, though the variance of the generated points is high for both of the sphere and sigmoid datasets.
Other training details, such as the general trend of manifold error, encoder variance and VAE loss, can be referred to in Figure 5 and 6.
Overall, the beneﬁt of clipping the decoder variance during training is inconclusive as we see inconsistent results in the sphere and sigmoid datasets.
Designing more algorithms to improve the ability of VAE’s to recover data supported on a low dimensional manifold is an important direction for future work—both empirical and theoretical.
of it
whereas the those generated without decoder variance clipping tend to lie inside the sphere as in Figure
[Figure]
[Figure]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
generated by padding uniformly random samples from a unit r -sphere with zeroes, so that the sphere is embedded in a higher ambient dimension.
We evaluated the manifold error as described in the setup.
The VAE training on this dataset has consistently yielded encoder variances with number of 0.1 entries greater than the number of intrinsic dimension.