Variational autoencoders in the presence of low-dimensional data: landscape and implicit bias

1 INTRODUCTION . 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 , fine-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 . ( 2016 ) ) . 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 0 , ( 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 . ( 2017 ) ) . 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 : Rr → Rd , 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 findings 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 . 2 SETUP . 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-defined “ 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 , Ir ) , and output x = f ( z ) , for a suitably chosen f . In the linear case , f ( z ) = Az , for some matrix A ∈ Rd×r . In the nonlinear case , f ( z ) will be a nonlinear function f : Rr → Rd . 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 , Ir ) and output x ∼ N ( f ( z ) , I ) , for trainable f , ; the encoder given input x will output z ∼ N ( g ( x ) , D ) , where D ∈ Rr×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 ) . 3 OUR RESULTS . 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 , Ir ) through a linear generator x = Az with AAT = Σ ; 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 + 2I = Σ . 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 — confirming predictions made in Dai & Wipf ( 2019 ) . We then explain theoretically why this happens , where it turns out we find 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 fixed matrices such that A = ÃB̃A and suppose that # { i : Ãi = 0 } > r − d , i.e . the number 1They also argued this would hold in the nonlinear case , but our simulations show this is generally false in that setting , even for the solutions chosen by gradient descent with a standard initialization — see Section 6. of zero columns of Ã is strictly larger than r − d. Then there exists ̃t → 0 and positive diagonal matrices D̃t such that limt→∞ L ( Ã , B̃ , D̃t , ̃t ) = −∞ . Also , these are all of the asymptotic global minima : any convergent sequence of points ( Ãt , B̃t , D̃t , ̃t ) along which the loss L goes to −∞ satisfies Ãt → Ã , B̃t → B̃ with A = ÃB̃A 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 LLT = Σ with L : d × r∗ and defining A : d × r by A = [ L 0d×r−r∗ ] . The larger the gap is between the manifold/intrinsic dimension r∗ and the ambient dimension d , the more flexibility 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 satisfied if r < d. To summarize , the asymptotic global minima satisfy A = ÃB̃A , 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 = [ Ir∗ 0 0 0 ] , then Ã = [ Ir∗+1 0 0 0 ] and B̃ = [ Ir∗+1 0 0 0 ] is a perfectly valid asymptotic global optima of the landscape , even though decoder Ã generates a different higher-dimensional Gaussian distribution N ( 0 , [ Ir∗+1 0 0 0 ] ) 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 fit 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 flow . The above theorem indicates that studying the minima of the loss landscape alone can not 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 flow , 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 ( spanA ) singular values of Ã decay at an exponential rate , as long as : ( 1 ) the gradient flow continues to exist2 , and ( 2 ) the gradient flow does not go off to infinity , i.e . neither Ã or ̃ go to infinity ( in simulations , the decoder Ã converges to a bounded point and ̃ → 0 so the latter assumption is true ) . To simplify the proof , we work with a slightly modified loss which “ eliminates ” the encoder variance by setting it to its optimal value : L1 ( Ã , B̃ , ̃ ) : = minD̃ L ( Ã , B̃ , ̃ , 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̃ , ̃ then L will behave like the simplified loss L1 . ) 2We remind the reader that the gradient flow on loss L ( x ) is a differential equation dx/dt = −∇L ( x ) . Unlike discrete-time gradient descent , gradient flow in some cases ( e.g . dx/dt = x2 ) has solutions which exist only for a finite time ( e.g . x = 1/ ( 1 − t ) ) , which “ blows up ” at t = 1 ) , so we need to explicitly assume the solution exists up to time T . Theorem 2 ( Implicit Bias of Gradient Flow , Informal version of Theorem 5 ) . Let A : d × r be arbitrary and define W to be the span of the rows of A , let Θ̃ ( 0 ) = ( Ã ( 0 ) , B̃ ( 0 ) , ̃ ( 0 ) ) be an arbitrary initialization and define the gradient flow Θ̃ ( t ) = ( Ã ( t ) , B̃ ( t ) , ̃ ( t ) ) by the ordinary differential equation ( ODE ) dΘ̃ ( t ) dt = −∇L1 ( Θ̃ ( t ) ) ( 1 ) with initial condition Θ0 . If the solution to this equation exists on the time interval [ 0 , T ] and satisfies maxt∈ [ 0 , T ] maxj [ ‖ ( Ãt ) j‖2 + ̃2t ] ≤ K , then for all t ∈ [ 0 , T ] we have d∑ k=dim ( W ) +1 σ2k ( Ã ( t ) ) ≤ C ( A , Ã ) e−t/K ( 2 ) where C ( A , Ã ) : = ‖PW⊥ÃT ( 0 ) ‖2F and PW⊥ 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 spanA . 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 .

Summary: The authors study further on the conjectures of Dai & Wipf (2019). For the linear case, they provide the proof that the conjecture is true. For the non-linear case, the paper disagrees with the conjecture, and they argue that the VAE training frequently learns a higher-dimensional manifold which is a superset of the ground truth manifold.