Technical Summary: An Analysis of Posterior Collapse, Parameterization and Initialization in Variational Deep Gaussian Processes

1. Problem Statement

Deep Gaussian Processes (DGPs) are hierarchical probabilistic models that concatenate Gaussian Processes (GPs) across multiple layers to capture complex, non-linear data patterns. However, exact Bayesian inference in DGPs is intractable, necessitating the use of Variational Inference (VI), specifically the Doubly Stochastic Variational Inference (DSVI) algorithm.

A critical challenge in training DGPs via DSVI is posterior collapse, a phenomenon where the variational posterior collapses to the prior distribution. In the context of DGPs, this implies the model learns to explain the observed data entirely as noise, failing to capture the underlying signal.

The paper identifies that the prevalence of posterior collapse is strongly linked to two common design choices in the DSVI framework:

1. The use of a ZERO prior mean function in inner layers, initialized with variational means set to zero ($m=0$).
2. The initialization of variational parameters, particularly the variational covariance ($S$) and the noise parameter ($\sigma^2$).

While the community has widely adopted a linear (PCA) prior mean function in inner layers to mitigate these issues, the prevailing belief was that this choice was necessary to avoid "non-injective pathologies" in very deep DGPs (where the mapping from input to output becomes non-injective). This paper challenges that assumption, arguing that the success of the PCA prior mean is not due to architectural necessity but rather to how it conditions the optimization landscape at initialization.

2. Methodology and Analysis

2.1 Theoretical Analysis of Posterior Collapse

The authors analyze the training dynamics of DGPs at initialization using a toy regression problem and coordinate update simulations. They demonstrate that:

• Virtual Dataset Effect: When a DGP with a ZERO prior mean is initialized with $m=0$, the output of the inner layers (which serve as input to the next layer) tends to be close to zero for all training inputs. Consequently, the final layer effectively sees a "virtual dataset" of pairs $\{(0, y_n)\}$. Since the input to the final layer is constant (zero) regardless of the original input $x_n$, the model cannot learn a mapping between $x$ and $y$.
• Optimization Pathology: To maximize the Expected Log-Likelihood (ELL) term in the Evidence Lower Bound (ELBO) under these conditions, the optimizer is driven to increase the noise variance $\sigma^2$ to explain the targets $y_n$ as pure noise. Simultaneously, the Kullback-Leibler Divergence (KLD) term is minimized by setting the variational parameters to match the prior ($m=0, S=K_{ZZ}$), resulting in posterior collapse.
• Role of PCA Prior Mean: A DGP with a PCA prior mean (where the prior mean is the identity function or a PCA projection) initializes the inner layer outputs to be close to the input $x_n$. This prevents the "virtual dataset" from collapsing to a constant zero vector, allowing the ELL term to provide meaningful gradients to update the variational parameters away from the prior.

2.2 Parameterization Analysis

The paper evaluates three common parameterizations of DGPs:

1. Standard Non-Whitened: Used in libraries like GPYTORCH and GPJAX.
2. Whitened: A reparameterization where inducing point values are transformed to have zero mean and unit variance ($p(v) = \mathcal{N}(0, I)$).
3. Non-Whitened with Reparameterized Mean: Used in GPFLOW, where the variational mean is defined as the difference between the variational mean and the prior mean ($m_r = m - \mu_Z$).

Key Findings on Parameterization:

• Posterior Collapse: Collapse can occur in all parameterizations if the initialization leads to a zero-input signal at the final layer. The parameterization itself is not the root cause.
• Optimization Stability: The whitened parameterization consistently provides more stable convergence. The authors attribute this to the removal of prior-induced correlations, which yields better-conditioned gradients and a Hessian that is less dependent on the inducing point locations ($Z$).
• Reparameterized Mean: The GPFLOW reparameterization does not inherently improve stability over the standard non-whitened form but interacts with the PCA prior mean to mitigate collapse more effectively than the standard form.

2.3 Proposed Initialization Strategy

To address the root cause (optimization-driven constraints) without forcing a specific prior mean function (which should be dictated by modeling assumptions), the authors propose a principled initialization scheme for ZERO prior mean DGPs.

The core idea is to initialize the variational parameters of a ZERO prior mean DGP such that its predictive distribution at initialization mimics that of a PCA prior mean DGP.

• Mechanism: The authors derive a system of linear equations to solve for the initial variational means ($m, m_v, m_r$) such that the predictive mean at the inducing points matches the identity function (or PCA projection).
• Output Layer Initialization: The strategy is extended to the output layer, initializing the variational mean to predict the observed targets $y$ directly at the inducing points, rather than starting at zero.
• Kernel Length-scale: A heuristic is proposed to select the initial kernel length-scale ($\ell$) by minimizing the Root Mean Square Error (RMSE) of the initialized model on a validation split, ensuring the initial predictive distribution is representative of the data.

3. Key Contributions

1. Reinterpretation of PCA Prior Mean: The paper demonstrates that the effectiveness of the linear (PCA) prior mean function in DGPs stems from its ability to condition the optimization landscape at initialization, preventing the model from seeing a degenerate zero-input signal, rather than from avoiding non-injective pathologies in deep networks.
2. Diagnosis of Posterior Collapse: The authors provide a rigorous analysis showing that posterior collapse in DGPs arises from the interaction between standard variational parameter initialization ($m=0$) and the inner-layer mean function, leading to a local optimum where data is explained as noise.
3. Parameterization Analysis: The study clarifies the distinct optimization behaviors of whitened, standard non-whitened, and reparameterized non-whitened DGPs, offering the first detailed explanation for why whitening yields more stable convergence (better-conditioned gradients).
4. Novel Initialization Scheme: A computationally efficient initialization strategy is proposed for ZERO prior mean DGPs. This allows for the successful training of DGPs under principled ZERO mean priors, removing the need for optimization-driven architectural constraints (i.e., forcing a linear prior mean).

4. Experimental Results

The proposed methods were validated on a synthetic toy dataset and eight real-world UCI regression benchmarks.

• Toy Dataset: Visualizations confirmed that ZERO prior mean DGPs with standard initialization suffer from posterior collapse, especially with fewer inducing points and deeper architectures. The proposed initialization successfully prevented collapse and produced predictive distributions comparable to PCA prior mean models.
• UCI Benchmarks:
  • DGPs with ZERO prior mean and standard initialization frequently collapsed or performed poorly.
  • The proposed initialization (ZERO-Points-MY) consistently stabilized optimization, avoided collapse, and achieved predictive performance (RMSE) comparable to, and in some cases better than, DGPs using the PCA prior mean function.
  • The whitened parameterization demonstrated superior stability across all configurations.
  • Increasing the number of inducing points was shown to alleviate collapse, but the proposed initialization was effective even with moderate numbers of inducing points.

5. Significance and Claims

The paper claims that the choice of prior mean function in DGPs should be driven by modeling assumptions rather than optimization convenience. By identifying that the "PCA prior mean" is essentially a heuristic to fix a poor initialization, the authors propose a solution that decouples the model's statistical structure from its optimization trajectory.

The significance lies in:

• Theoretical Clarity: Correcting the misconception that linear priors are necessary to avoid non-injectivity in deep GPs, instead attributing their success to initialization dynamics.
• Practical Utility: Providing a drop-in initialization strategy that enables the use of ZERO prior mean DGPs (which are theoretically simpler and more flexible for certain modeling tasks) without sacrificing training stability or predictive performance.
• Optimization Insight: Highlighting the critical role of the whitened parameterization and specific initialization of variational covariances ($S \ll K_{ZZ}$) in ensuring stable gradient descent, offering guidelines for robust DGP training.

The authors remain modest, noting that their analysis focuses on the DSVI algorithm and that while the proposed initialization solves the collapse problem for ZERO mean priors, the statistical model space remains unchanged; the improvement is strictly in the ability to reach a good solution from a better starting point.