On the Convergence Analysis of Over-Parameterized Variational Autoencoders: A Neural Tangent Kernel Perspective

In this work, we adopt a standard notation to represent vectors, matrices, and scalars. Specifically, we use bold-faced letters for vectors and matrices and non-bold letters for scalars. To denote the Euclidean norm of a vector or the spectral norm of a matrix, we use the notation $\|\cdot\|_{2}$. The Frobenius norm of a matrix is represented by $\|\cdot\|_{F}$. We use the notation $[n]={1,2,\ldots,n}$ to represent the set of integers from 1 to $n$. Besides, we represent a matrix as a set of row vectors, i.e., ${\bf W}=[{\bf w}^{\top}_{1},{\bf w}^{\top}_{2},\dots,{\bf w}^{\top}_{m}]^{\top}$, where ${\bf w}_{r}$ with $r\in[m]$ is a column vector of the matrix. Finally, we denote the least eigenvalue of a matrix by $\lambda_{0}(\boldsymbol{\Theta})$, which is equivalent to $\lambda_{\min}(\boldsymbol{\Theta})$.

A Variational auto-encoder (VAE) [1], as a directed probabilistic graphical model (DPGM), is designed to learn a latent variable model. Its primary objective is to maximize the log-likelihood of the training data $\{{\bf x}_{i}\}^{n}_{i=1}$ via variational inference, where $n$ is the number of training samples. The VAE introduces a distribution $q_{\phi}(\bf z|x)$ to approximate the intractable true posterior $p(\bf z|x)$, where $\phi$ are neural network parameters that can be learned in the encoder. Then, the decoder takes ${\bf z}$ as input to generate ${\bf x}^{\prime}$ as a reconstruction for ${\bf x}$.

The common training objective of the VAE is to maximize the Evidence Lower Bound (ELBO), given by:

where ${\bf z}\sim q_{\phi}({\bf z}|{\bf x}_{i})$, and $\phi$ and $\theta$ represent the parameters in encoder and decoder, respectively. The first term in the ELBO measures the reconstruction loss between the generated ${\bf x}^{\prime}$ and the original ${\bf x}$. The second term represents the Kullback-Leibler (KL) divergence between the approximate posterior $q(\bf z|x)$ and the prior $p({\bf z})$, where $p({\bf z})$ is often chosen to be an isotropic multivariate Gaussian distribution.

Consider a stochastic neural network (SNN) $\mathbf{f}\in\mathbb{R}^{d}$, where $d$ is the input dimension. In the context of this work, our SNN is defined as follows:

where ${\bf x}^{(e)}\in\mathbb{R}^{d}$ is the encoded representation derived from the input $\mathbf{x}$, ${\bf W}^{(\mu)},{\bf W}^{(\sigma)}\in\mathbb{R}^{m\times d}$ are weight matrices employed to construct the latent Gaussian representation. Here $m$ represents the width of the network, indicating the number of neurons, $\sigma(\cdot)$ is the non-linear activation function, $\psi(\cdot)$ is the decoder representation function, and ${\bf W}^{(d)}\in\mathbb{R}^{m\times d}$ is the linear weight matrix utilized in the final layer. A visual representation of the SNN under study is depicted in Figure 1.

In the construction of the latent representation, we employ the re-parametrization trick, a technique that allows for the backpropagation of gradients through random nodes. In particular, the latent variable can be expressed as:

where $\mathbf{W}^{(\mu)}$ and $\mathbf{W}^{(\sigma)}$ represent the mean and variance weights, respectively. Besides, $\boldsymbol{\zeta}$ is a random variable drawn from a standard normal distribution.

Given the structure of the SNN, our objective function considered in this work is defined as:

where $\hat{\mathbf{f}}(\mathbf{x}_{i})\triangleq\mathbb{E}_{\boldsymbol{\zeta}}[\mathbf{f}(\mathbf{x}_{i},\boldsymbol{\zeta})]$, and $\mathbf{z}_{i}(t)$ is the latent representation for input $\mathbf{x}_{i}$ at time $t$. Besides, $\beta$ is an adjustable hyperparameter that balances latent channel capacity and independence constraints with reconstruction accuracy [45]. The first term $\sum_{i=1}^{n}\ell(\hat{\mathbf{f}}(\mathbf{x}_{i}),\mathbf{x}_{i})$ is called the reconstruction loss. In this study, we utilize the mean squared error as our reconstruction loss, following seminal theoretical works [35, 17, 27, 26]. The second term ${\rm KL}(\cdot)$ is a Kullback–Leibler (KL) divergence, where prior distribution is the Gaussian distribution of latent variable at initialization, and the posterior is the distribution of latent variable after training, $\mathbf{z}_{i}\sim\mathcal{N}({\bf W}^{(\mu)}{\bf x}^{(e)}_{i},\mathrm{diag}({\bf W}^{(\sigma)}{\bf x}^{(e)}_{i}))$. It’s worth noting that our KL is tailored to align with our theoretical analysis for constructing kernel ridge regression.

To optimize the objective function given by (4), we adopt a gradient descent rule:

where $\eta$ is the learning rate. Note that while the weights in the encoder and decoder remain fixed, we specifically optimize the mean weights $\mathbf{W}^{(\mu)}$, variance weights $\mathbf{W}^{(\sigma)}$, and the weights in the final layer $\mathbf{W}^{(d)}$. This optimization strategy is primarily adopted for the sake of theoretical simplicity. It’s worth noting that this choice does not compromise or alter our final conclusions.

For the purpose of our optimization analysis, we introduce the concept of the neural tangent kernel for a stochastic neural network:

A few remarks on Definition 1 are in order. Unlike standard (deterministic) neural networks, the VAE comprises two sets of parameters in the latent layer, namely, ${\bf W}^{(\mu)}$ and ${\bf W}^{(\sigma)}$. Due to the reparameterization trick, gradient descent is executed on each of these parameters independently. Consequently, we observe two distinct tangent kernels corresponding to each parameter set. Secondly, The scenario with multiple outputs in variational autoencoder networks presents added complexity compared to networks with a single output [16, 26]. Given that the output dimension of the stochastic neural network is $d$, the neural tangent kernel is a matrix of size $\mathbb{R}^{nd\times nd}$. As we delve deeper in the subsequent sections, it will become evident that the non-diagonal NTK across the output index is zero, and the diagonal NTK remains consistent across the output index. This uniformity allows us to employ Kronecker products, facilitating the derivation of NTKs.

Next, we impose some technical conditions on the activation function, which is stated as follows:

This assumption ensures that we can interchange the operations of integration and differentiation over the activation function. Subsequently, we present technical conditions on both the activation function and the decoder representation function:

These conditions are important in demonstrating the stability of the training process within the framework of the NTK.

For the sake of simplification, we focus on the optimization of the stochastic neural network as described in (2), emphasizing solely on the reconstruction loss. This means we are setting aside the KL divergence term for the time being. Additionally, given that we’re adopting a squared loss without KL divergence, the objective function (4) reduces to:

Then the gradient flow dynamics of output function $\hat{f}_{k}$ are governed by:

Equation (8) implies that the dynamics of output function are governed by the neural tangent kernels. Furthermore, as we will show later, the neural tangent kernels will stay constant during the training process in the infinite-width limit. In this way, Equation (8) reduces to an ordinary differential equation (ODE):

where we define the neural tangent kernel of an infinitely-wide SNN by:

To demonstrate the convergence result induced by Equation (9), we perform an in-depth concentration analysis. This analysis focuses on the convergence of stochastic neural networks in a non-asymptotic manner, i.e., with a large but finite width. We present our main result in the following theorem:

The proof sketch of Theorem 1 will be given in Section 5. Theorem 1 establishes that if $m$ is large enough, the expected training error converges to zero at a linear rate. In particular, the least eigenvalue of NTK governs the convergence rate.

By Theorem 1, we establish the global convergence of stochastic neural networks with a large width in VAE. Building on this foundation, we further consider full objective function (4) which incorporates an additional KL divergence term.

After a detailed calculation of the KL divergence for two Gaussian distributions, we simplify our analysis by making certain assumptions. Specifically, we assume that $\mathbf{W}^{(\sigma)}$ remains constant and select a prior $\mathbf{x}^{(e)}_{i}$ such that the objective function (4) is transformed to:

Building on this, we further analyze the regularization effect of the KL term when training VAEs and present our findings in the subsequent theorem:

The proof of Theorem 2 will be given in the Appendix. Note that the error term is bounded by the difference between the output function of the finite network and the infinitely-wide network. This difference is further decomposed into the initial difference and the difference during training. The latter can be bounded by $\frac{\sqrt{n}}{\lambda_{0}+\beta}\mathcal{E}_{\Theta}$, where $\sqrt{n}$ comes from the input and $\frac{1}{\lambda_{0}+\beta}$ results from the integration over the training time. Besides, the necessity of fixing the variance weight in Theorem 2 arises because we are seeking a closed-form solution under the NTK regime. Theorem 2 reveals the regularization effect of the KL divergence on the convergence of over-parameterized VAEs and makes a connection between solution of training a VAE and kernel ridge regression.