Mode Collapse in Variational Deep Gaussian Processes

Authors: francisco javier sáez-maldonado
juan  maroñas
daniel  hernández-lobato
Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-09-04 12:22:34

Abstract: Deep Gaussian Processes (DGPs) define a hierarchical model capable of learning complex, non-stationary processes. Exact inference is intractable in DGPs, so a variational distribution is used in each layer. One of the main challenges when training DGPs is the prevention of a phenomenon known as mode collapse where, during training, the variational distribution becomes the prior distribution which is a minimizer of the KL-Divergence term in the ELBO.  There are two main factors that influence the optimization process: the mean function of the inner GPs and the usage of the whitened representation of the variational distribution. In this work, we propose a data-driven initialization of the variational parameters that a) at initialization, predicts an already good approximation of the objective function, b) avoids mode collapse c) is supported by a theoretical analysis of the behavior of the KL divergence and experimental results in real-world datasets.

Main paper: 
 introduction: 
 Gaussian Processes (GPs) [1] are flexible non-parametric models that have shown promising results across many applications, such as molecule optimization [2, 3] or uncertainty estimation in DNNs [4, 5] . They are used as prior distributions over some target function f , i.e., f ∼ GP(µ(•), K ϕ (•, •)), where µ(•) and K ϕ (•, •) are, respectively, the mean and covariance functions of the GP. Deep Gaussian Processes (DGPs) [6] concatenate several GPs across a layered network, defining a hierarchical structure that leads to a more flexible probabilistic model. This concatenation enlarges the class of functions that can be modeled, allowing to capture complex patterns within the data.Bayesian inference in DGPs is intractable. Thus, approximations using variational inducing points such as Double Stochastic Variational Inference (DSVI) [7] must be employed. To alleviate optimization difficulties, two main tools are often used: (i) a whitened representation of the process values at the inducing points [8] to enhance the numerical stability of the model [9] [10] [11] , and (ii) an identity mean function in the inner layers of the model [7] . In [7] , it is claimed that the identity mean function is needed to avoid the pathologies outlined in [12] . However, the reality is that, in practice, this mean function is used even in 2-layer DGPs, which are far from being similar models to those in [12] .Here, we show, via experiments in a toy dataset and a dataset from the UCI repository, that the need for the identity mean function in DGPs is linked to the initial variational parameters and the usage of the whitened representation. Specifically, we found that a zero mean DGP, with the variational initialization used in [7, [9] [10] [11] , can lead to mode collapse during the optimization process, which is undesirable. Namely, the algorithm sets the variational mean and covariance to those of the GP prior (a standard Gaussian in the whitened case). Then, a huge noise variance is set in the observation model to explain the observed data, which is considered pure noise. This is not a desirable behavior.Our main contributions are: 1) we explain why the whitened representation, beyond being designed to enhance mixing in MCMC algorithms [8] , provides numerical stability in the Variational's Inference optimization process; 2) we highlight why the identity mean function in DGPs avoids the mode collapse effect that occurs when the zero mean function is used; 3) we provide a theoretical explanation of the effect of whitening in the optimization process; and, 4) we propose a new initialization of the variational parameters of the zero mean DGP that alleviates the mode collapse problem. 
mode collapse in variational dgps: 
 Consider the observed data D = {(x i , y i ) | x i ∈ R D , y i ∈ R, i = 1, . . . , N } which is grouped into a matrix X = (x 1 , • • • , x N )f (1) (x) GP y GP f (3) (x) GP f (2) (x)inducing locations Z = (z 1 , • • • , z M ), with z i ∈ R D , have their corresponding inducing values u = f (Z) = (u 1 , • • • , u M ) T . The posterior distribution is approximated using a variational distribution q(f, u) = p(f | u)q(u), with q(u) = N (u | m, S)where m and S are the variational parameters [13] . The prior over u is N (µ Z , K ZZ ). In the whitened representation, we rewrite u = L ZZ v + µ Z where K ZZ = L ZZ L T ZZ and p(v) = N (0, I). Then, we make inference about v instead of u and learn q(v) = N (m v , S v ). Variational parameters are initialized as m = m v = 0, S = S v = 10 -5 I.DGPs concatenate layers of GPs as observed in Fig. 1 . Adding the superscript l to denote the l-th layer, the training objective for the DSVI algorithm, described in [7] , is the ELBO:EQUATIONNon-whitened GPs: In classic SVGPs, the mean of the marginal variational distribution q(f l=1 ) is1 :EQUATIONThe KL between prior and posterior is given by:EQUATIONWhitened GPs: In the whitened representation we have:EQUATIONEQUATIONKL at initialization and mode collapse: In the whitened case, at initialization, the marginal distribution's variational mean for a zero mean DGP is zero and the KL is very close to zero . Thus, the minimization of the objective will likely force m v = 0 and S v = I (since it is the KL minimizer), and will minimize the ELL, i.e., the data-dependent term in (1) , by learning a huge observation noise. However, when using an identity mean function it is simpler to optimize the ELL, since the posterior mean at initialization is x i , the optimization forces m v to move away from its initial value 0 to learn a map from x i to the outputs y. This often avoids mode collapse. For the non-whitened case, at initialization, the variational mean is also 0 in the zero mean DGP. However, we should expect the model not to be so prone to mode collapse, since at initialization the KL depends on K ZZ via the prior, which can differ from the identity matrix (depending on the length-scale value). The contrary may also be true, i.e., one may also expect mode collapse, since now we can freely adapt both the non-whitened GP prior kernel hyper-parameters and Z, and the variational parameters m and S to minimize KL. Therefore, in the non-whitened case, mode collapse depends not only on the initialization of q, but on the initial kernel's hyper-parameter as well. In any case, this model usually suffers from optimization difficulties, so even though mode collapse can be controlled, the model is not useful in practice as it is complicated to optimize.Optimization difficulties: There is a different impact on the KL divergence depending on whether the whitened representation is used or not. In the whitened representation one often observes that: 1) the KL at initialization is close to zero, providing a small learning signal to the objective; and2) the KL only depends on the variational parameters m and S. By contrast, in the non-whitened representation, the KL divergence at initialization is far from zero, both for the zero mean and even more for the identity mean function (due to the term m -µ Z ). Furthermore, the KL depends on both the variational and the model parameters, e.g., Z, length-scales, etc. These two differences complicate optimization and we hypothesize they are the reason why one observes oscillations in the objective during optimization, making the whitened representation the standard one used both in variational GPs and DGPs [9] [10] [11] , see Fig. 6 in the Appendix. 2Understanding the implementation of DGPs is key to a comprehensive view of their optimization process. In GPFLOW, the library used in our experiments, the whitened representation is implemented as outlined above. However, the non-whitened representation is implemented using the reparameterization m ′ = m -µ Z , which changes the predictive mean and KL divergence to:µ qf = K XZ K -1 ZZ m ′ + µ X , KL (q(u) ∥ p(u)) = 1 2 -log|S| -M + log|K ZZ | + Tr K -1 ZZ (S + m ′ m ′ T) .This reparameterization simplifies the implementation since both whitened and non-whitened models require KL computation between q and a zero mean prior. Nevertheless, the optimization difficulties outlined above persist due to the term K ZZ in the KL. For the rest of the work, unless explicitly mentioned, we will be using the aforementioned reparameterized version of the variational mean. 
avoiding mode collapse via data-driven initialization: 
 We now present a novel initialization that mimics the effect of the identity mean function in DGPs but uses a zero mean function. We achieve this by selecting appropriate values for m, m v and Z.Inner layers: At initialization, in the inner layers (l = 1, • • • , L -1), when using the identity mean function, the mean of the posterior distribution takes the value µ w qf = µ qf = x for any point x ∈ X. We aim at selecting m and m v (in the whitening case) such that, at initialization, we obtain the same posterior mean while using a zero mean DGP. Assuming a 1-dimensional input space (although this can be generalized to arbitrary dimensions), this is achieved by solving the linear systems:EQUATIONThese are linear equations in m and m v , that depend on three independent variables (m, X, Z.).To solve them, X, Z must be given some value, requiring dim(X) = dim(Z) to yield a solvable system. We fix the inducing locations Z to a subset of the train data, i.e., Z ⊂ X. This choice gives the solutions m = Z, and m v = L -1 ZZ Z, and it guarantees that the initialization will predict identity values at the inducing points, i.e., when X = Z. In both cases, we initialize S = 10 -5 I so that the initial forwarded samples in DSVI do not differ from the mean.Output layer: DGPs usually use a zero mean GP in its output layer with m, m v initialized to 0. We propose to extend the idea of the inner layers to the output layer. However, in this case, we want the posterior mean to resemble the target y values3 . To this end, we replace X by y in both l.h.s. of Eq. ( 7). This gives m = y, and m v = L -1 ZZ y, when using whitening. By selecting the inducing locations Z as a subset of X, we obtain precise posterior predictive means at those data points at initialization. This leads to an initial good solution when using enough inducing points, see Fig. 2 .Let us remark the difference between our proposed initialization and that of standard DGP configurations. Consider a non-whitened DGP with the aforementioned reparameterization of the variational mean m ′ . Assume an identity mean function in the inner layers, a zero mean function in the output layer, and that we initialize m ′ = 0. This is the typical setup considered in the literature for DGPs [7] . The output of the inner layers of this model model at Z, is Z, i.e., the identity, as specified by the prior mean µ Z . See Eq. ( 6). However, since the identity mean function is not used in the last layer [7] , the predictive mean at the inducing points Z will be zero, unlike in the proposed method, described in the previous paragraph. Something similar will happen in the case of the whitened representation, when m v = 0. See Eq. ( 4). That is, the initial DGP predictive mean, without training the model, will be zero for every input.Inducing points: The initial inducing points may impact the initial solution. Selecting them at random from the training set may lead to regions of the input space being unrepresented, leading to a poor initial solution in those areas. See Fig. 2 . To ensure a good initialization, we use a two-step algorithm:1. We compute M centroids C = {c 1 , • • • , c M } of X train using k-means.2. We select the inducing locations z j from X train using the cosine distance:EQUATIONFig. 2 compares the predictive distribution (without training) of the standard initialization of a DGP with that of the proposed initialization with a random selection of the inducing points among the training points, and when selecting the inducing locations as described above. We observe that the proposed initialization already explains quite well the observed data with no training of the objective.Optimization difficulties: As mentioned, with the standard initialization, the KL in the whitened case provides a more stable learning signal (Sec. 2). In the proposed initialization, in the whitened case we have p(v) = N (0, I) and q(v) = N (L -1 ZZ Z, 10 -5 • I). In the non-whitened (reparameterized version) we have p(u) = N (0, K ZZ ) and q(u) = N (Z, 10 -5 • I). Therefore, the whitened parameterization is expected to lead to a more stable optimization process since, although both KL are different from 0 at initialization, the gradient of the KL for the non-whitened case depends on both the model parameters (e.g., length-scales and inducing points) and the variational parameters.  
experiments: 
 To validate the proposed initialization, we compare it with the standard initialization. We perform an extensive qualitative evaluation on the toy dataset presented in Fig. 2 and on the Yatch dataset. The results show that the proposed initialization has two benefits: a) it leads to a faster convergence than the standard initialized DGP; and b) the proposed initialization avoids mode collapse when using a 5-layer DGP. We perform an extensive analysis of 8 different UCI datasets. See Appendix B.2. The most interesting case is the Yacht dataset analyzed below. The data-dependent term of the objective (ELL), as a function of the training epochs, is shown in Fig. 3 (left). The KL is also reported (center). We observe a clear example of the mode collapse in the 5-layer DGP when using the zero mean function with the standard initialization. Specifically, the test log-likelihood results displayed in Fig. 3 (right) are significantly worse for the zero mean function with the standard initialization DGPs. We observe the ELL and the KL (left and center of the figure). In the standard initialization, when using the zero mean function, the KL quickly falls to zero when whitening is not used. Besides this, we also observe that the KL has a high variance when the whitened representation is used. This indicates that there are some data splits where mode collapse happens, even when using whitening. By contrast, the proposed initialization solves the mode collapse problem, preventing the KL from falling to zero by choosing a good initial solution. It also leads to an overall higher final KL. The ELL chart (left) shows that, in the whitened case, our proposal exhibits a fast convergence and the final results pair up with the DGP with identity mean function. The non-whitened case has a slower convergence because the initial solution has a low ELL, caused by the initial value of the kernel's length-scale 4 . We observe this behavior in other datasets. Appendix B.2 has a figure with all the results of LL and RMSE.As a limitation of the proposed initialization, we remark that the correct inversion of the kernel evaluated at the inducing points highly influences the initial solution. The parameters of the kernel (in our case, the length-scale ℓ of the RBF kernel) must be properly selected so that the covariance matrix is not ill-conditioned, making the proposed initialization properly predict the corresponding inducing value at the selected inducing locations. An example of this can be observed in Fig. 4 . However, this limitation can be surpassed by performing an initial evaluation of the model using different values of the kernel length-scale and selecting the one with higher ELL or lower RMSE. 
conclusions: 
 DGPs may suffer from mode collapse (i.e., convergence to the prior) as the number of layers grows.Here, we have presented an initialization of SVGPs and DGPs that uses the training data. This initialization fixes the inducing points and the variational parameters to predict an initial good solution of the targets. The results obtained, when using this initialization, show that it has benefits both in the convergence speed DGPs and in avoiding the mode collapse problem. A limitation is that it requires solving a linear system that depends on the initial kernel parameters, which may be problem-dependent. 
 
 Appendix: We include an extended background section for a clearer understanding of the problem statement and notation.We tackle the standard regression problem, where the goal is to infer an unknown function f :R D → R, using the observed data D = {(x i , y i ) | x i ∈ R D , y i ∈ R, i = 1, . . . , N }. We denote X = (x 1 , • • • , x N ) and y = (y 1 , • • • , y N ) . Gaussian Processes [1] place a multivariate Gaussian prior distribution over f , that is, f = (f (x 1 ), • • • , f (x N )) ∼ N (µ(X), K ϕ (X, X)), where µ : R D → R is the mean function, and K ϕ (X, X) is the covariance matrix determined by a covariance function K ϕ : R D × R D → R. We write µ x = µ(x) and K xx = K ϕ (x, x) to shorten notation. The exact posterior distribution p(f | D) is intractable since it scales cubically with N , so we use sparse variational Gaussian Processes (SVGPs) [13] to scale to large datasets. = f (Z) = (u 1 , • • • , u M ), where p(u) = N (u | µ(Z), K ϕ (Z, Z)).The posterior distribution is approximated using a variational distribution q(f, u) = p(f | u)q(u), where q(u) = N (u | m, S) where m and S are the variational parameters. Using this variational distribution, the posterior distribution is Gaussian q(f) = N (f | µ qf , K qf ), where:EQUATIONEQUATIONThe presented formulation is sometimes referred to as the unwhitened representation of a GP. Another common form is the whitened prior GP representation [8, 14] . This consists of expressing u = L ZZ v with L ZZ L T ZZ = K ZZ and v = N (v | 0, I). The original variational parameters are recovered as m = µ Z + L ZZ m v and S = L ZZ S v L T ZZ , so the parameters to estimate are now m v , S v . Using this representation, the posterior q w (f) = N (f | µ w qf , K w qf ) is computed as:EQUATIONEQUATIONWhitening the GP prior leads to a much faster convergence in many cases [14] .Following the fashion of deep neural networks, DGPs concatenate layers of GPs to generalize them and extend their flexibility. Formally, the output of the GP in layer l -1 is used as the input of layer l, defining a prior recursively on each layer F 1 , • • • , F L . Each layer l is defined by its inducing locations and values (Z l , u l ), and the corresponding variational parameters (m l , S l ). The joint distribution of the DGP factorizes as:EQUATION)The exact posterior remains intractable in DGPs, so it is approximated using a variational distribution that mimics the one in SVGPs and assumes layer independence:q({F l , u l } L l=1 ) = L l=1p(F l | u l ; F l-1 , Z l-1 )q(u l ).With these considerations, we obtain an ELBO that can be optimized using two sources of stochasticity [7] :L = N n=1 E q(f L n ) log p(y n | f L n ) - L l=1KL q(u l ) ∥ p(u l ; Z l-1 (15) results which are at least paired to the results obtained using the identity mean function in the inner layers of the DGP.In Fig. 8 we represent the results using the RMSE. The conclusions are similar to the ones obtained from the LL results. Again, the poor solutions in terms of RMSE can be clearly seen in the case of a zero mean DGP with standard initialization, with those cases being solved with our initialization.When l >2, X is replaced by points f l-1 sampled from q(f l-1 ) at the previous layer.We have not explicitly found that these are the reasons why these libraries use whitening in practice, beyond being recommended in its source code, as seen for example in[11] in this comment.In classification we would use a one-hot encoding of the class labels.In the UCI datasets we have not selected a good initial kernel length-scale for each dataset. Due to this, the faster convergence is not appreciated in Fig.3. 



Reviews: review: title: title: Review
main_review: main_review: The paper investigates mode collapse in Deep Gaussian Processes (DGPs) and proposes a new initialization method to prevent this problem. It identifies the limitations of the current approaches, particularly with the use of identity mean functions, and provides empirical evidence on toy datasets and UCI benchmarks.

**Strengths:**
- Clear Goal: The paper has a well-defined objective, focusing on addressing mode collapse in DGPs and demonstrating the effectiveness of the proposed initialization method.

- Interesting Observations: The insights into why the identity mean function prevents mode collapse are valuable, as well as the exploration of whitening effects on optimization. These observations are potentially useful for practitioners implementing DGPs in real-world applications.

**Weaknesses:**
- Lack of Theoretical Guidance: While the paper provides empirical results, it lacks a strong theoretical framework to guide the proposed approach. The explanations are mainly heuristic and would benefit from a more rigorous theoretical analysis.

- Experiments Limited to Toy Datasets: Although the paper is empirical, most experiments are conducted on toy datasets and a small number of UCI datasets. For broader applicability, more experiments on real-world datasets with larger scales would strengthen the paper's contributions.
title: title: Overall is a good paper to be accept for workshop.
main_review: main_review: The document discussing the issue of mode collapse in Variational Deep Gaussian Processes (DGPs) and proposing a solution to this problem. An evaluation based on the given context is as follows:

Quality:

The addressed problem is well identified within a specific area of interest, Bayesian machine learning. In particular, the problem related to mode collapse for DGPs. It clearly presents the problem statement and the proposed solution to that effect, which is a new initialization strategy for the variational parameters. The paper presents experimental results to substantiate its claims; typically, such work implies good quality research.

Clearness:

Appearantly, the paper is well-structured: it introduces the problem, it elaborates on the proposed solution, and discusses the findings in a conclusive manner. The use of figures and tables (not included in the text provided) suggests that the authors have made effort to present their results in accessible ways.

Originality:

The paper claims to develop a new initialization strategy that prevents mode collapse, which would be an original contribution to the field. Another contribution is the theoretical explanation of the whitening effect in the optimization process, something rather new.

Impact:

The work can be impactful since it will address one known problem with DGPs, which are powerful models for function approximation and uncertainty estimation. In case the proposed solution works well, we will have more stable and reliable DGP models, adding value to the machine learning community.

prors:
- The paper tackles a relevant problem in the field of DGPs and provides a potential solution.
- It includes experimental validation, which adds credibility to the proposed method.
- The paper appears to be well-written and structured, making it accessible to readers.

 cons:
- The paper may require a good understanding of Gaussian Processes and variational inference, which could limit its accessibility to a broader audience.
While this is an effective solution, the dependence on the choice of the initial kernel parameters could be a flaw if not handled appropriately.
 Due to the lack of a complete experimental setup and results within the paper, complete robustness testing of the proposed method for different datasets and scenarios cannot be done.

In general, this paper appears to be a good contribution in the Bayesian machine learning area and might enable DGP models to become more stable and give better performance. However, without the complete paper showing details and results of experiments, it would be more appropriate to assess the full significance and impact of the work.