A variational approximate posterior for the deep Wishart process
Year: 2022

Abstract:
Recent work introduced deep kernel processes as an entirely kernel-based alternative to NNs (Aitchison et al. 2020). Deep kernel processes flexibly learn good top-layer representations by alternately sampling the kernel from a distribution over positive semi-definite matrices and performing nonlinear transformations. A particular deep kernel process, the deep Wishart process (DWP), is of particular interest because its prior can be made equivalent to deep Gaussian process (DGP) priors for kernels that can be expressed entirely in terms of Gram matrices. However, inference in DWPs has not yet been possible due to the lack of sufficiently flexible distributions over positive semi-definite matrices. Here, we give a novel approach to obtaining flexible distributions over positive semi-definite matrices by generalising the Bartlett decomposition of the Wishart probability density. We use this new distribution to develop an approximate posterior for the DWP that includes dependency across layers. We develop a doubly-stochastic inducing-point inference scheme for the DWP and show experimentally that inference in the DWP can improve performance over doing inference in a DGP with the equivalent prior.

heading: 1 Introduction
text: The successes of modern deep learning have highlighted that good performance on tasks such as image classification (Krizhevsky et al., 2012) requires deep models with lower layers that have the flexibility to learn good representations. Up until very recently, this was only possible in featurebased methods such as neural networks (NNs). Kernel methods did not have this flexibility because the kernel could be modified only using a few kernel hyperparameters. However, with the advent of deep kernel processes (DKPs; Aitchison et al., 2021), we now have deep kernel methods that offer neural-network like flexibility in the kernel / top-layer representation. DKPs introduce this flexibility by taking the kernel from the previous layer, then sampling from a Wishart or inverse Wishart centered on that kernel, followed by a nonlinear transformation. The sampling and nonlinear transformation steps are repeated multiple times to form a deep architecture. Remarkably, deep Gaussian processes (DGPs; Damianou & Lawrence, 2013; Salimbeni & Deisenroth, 2017), standard Bayesian NNs, infinite-width Bayesian NNs (neural network Gaussian processes or NNGPs; Lee et al., 2018; Matthews et al., 2018; Novak et al., 2019; Garriga-Alonso et al., 2019) and infinite NNs with finite width bottlenecks (Aitchison, 2020) can be written as DKPs (Aitchison et al., 2021). Indeed, for kernels that can be expressed in terms of operations on Gram matrices, Aitchison et al. (2021) showed that a particular DKP, the deep Wishart process (DWP) has a prior equivalent to a DGP’s prior. In a DGP, the random variables inferred in variational inference are the model’s intermediate features, with kernels computed as a function of these features at each layer. However, in a DWP, there are no features at all: the only random variables are the positive semi-definite kernel
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
matrices themselves, which are sampled directly from Wishart distributions: the DWP works entirely on the kernel matrices implied by the DGP’s features.
Aitchison et al. (2021) argued that DWPs should have considerable advantages over related featurebased models, because feature-based models have pervasive symmetries in the true posterior, which are difficult to capture in standard variational approximate posteriors. For instance, in a neural network, it is possible to permute rows and columns of weight matrices, such that the activations at a given layer are permuted, but the network’s overall input-output function remains the same (MacKay, 1992; Sussmann, 1992; Bishop et al., 1995). These permutations result in network weights with exactly the same probability density under the true posterior, but with very different probability densities under standard variational approximate posteriors, which are generally unimodal. However, these issues do not arise with DWPs, because all permutations of the hidden units correspond to the same kernel (see Appendix D in Aitchison et al. (2021) for more details).
While Aitchison et al. (2021) showed the equivalence between DWPs and DGPs, they were not able to do inference in DWPs because they were not able to find a sufficiently flexible distribution over positive semi-definite matrices to form the basis of an approximate posterior. Instead, they were forced to work with a different DKP: the deep inverse Wishart processes (DIWPs), which was easier because the inverse Wishart itself forms a suitable approximate posterior. While the DIWP also avoids using features, it does not correspond directly to an already-established Bayesian model. Therefore, the DWP is a more important model, as it allows us to directly compare feature-based and kernel-based inference. In this work, we show how to create a sufficiently flexible approximate posterior for DWPs, thereby enabling us to compare directly to their equivalent DGPs. In particular, our contributions are:
• We develop a new family of flexible distributions over positive semi-definite matrices by generalising the Bartlett decomposition (Sec. 3.1).
• We use this distribution to develop an effective approximate posterior for the deep Wishart process which incorporates dependency across layers (Sec. 3.2).
• We develop a doubly stochastic inducing-point inference scheme for the DWP. While the derivation mostly follows that for deep inverse Wishart processes (Aitchison et al., 2021), we need to give a novel scheme for sampling the test/training points conditioned on the inducing points, as this is very different in the DWP compared to the previous DIWP (Sec. 3.3).
• We empirically compare DGP and DWP inference with the exact same prior. This was not possible in Aitchison et al. (2021) as they only derived an inference scheme for deep inverse Wishart processes, whose prior is not equivalent to a DGP prior.
We provide a reference implementation at https://github.com/LaurenceA/bayesfunc.

heading: 2 Background

heading: 2.1 Wishart distribution
text: The Wishart,W (Σ, ν), is a distribution over positive semi-definite P ×P matrices, G, with positive definite scale parameter Σ ∈ RP×P and a positive, integer-valued degrees-of-freedom parameter, ν. The Wishart distribution is defined by taking ν vectors nλ ∈ RP sampled from a zero-mean Gaussian. These vectors can be generated from standard Gaussian vectors, ξλ, by transforming them with the Cholesky, L of the scale parameter, Σ = LLT ,
Lξλ = nλ ∼ N (0,Σ) where ξλ ∼ N (0, I) Both nλ and ξλ can be stacked to form P × ν matrices, N and Ξ,
N = (n1 n2 · · · nν) Ξ = (ξ1 ξ2 · · · ξν) . Wishart samples are defined by taking the sum of the outer products of the nλ’s, which can be written as a matrix-multiplication,
ν∑ λ=1 nλn T λ = NN T = LΞΞTL = LZLT = G ∼ W (Σ, ν) (1)
where Z = ΞΞT is a sample from a standard Wishart (i.e. one with an identity scale parameter,)
ν∑ λ=1 ξλξ T λ = ΞΞ T = Z ∼ W (I, ν) . (2)
Note that therefore the Wishart has mean, E [G] = ν E [ nλn T λ ] = νΣ (3)

heading: 2.2 Bartlett decomposition
text: However, sampling Ξ can be computationally expensive for very large values of ν. Instead, it is possible to sample a Wishart by writing down the distribution over the Cholesky of Z, denoted A (Bartlett, 1934). Taking Z = AAT , the distribution over A is,
P ( A2jj ) = Gamma ( A2jj ;α= ν−j+1 2 , β= 1 2 ) , (4a) P (Aj>k) = N (Ajk; 0, 1) . (4b) i.e. the square of the on-diagonal elements are Gamma distributed and the off-diagonal elements are IID standard Gaussian.

heading: 2.3 Deep Gaussian proceses (DGPs)
text: In a DGP, we progressively sample features, F`, from a Gaussian process, conditioned on features from the previous layer,
P (F`|F`−1) = ∏ν` λ=1N ( f `λ; 0,Kfeatures (F`−1) ) with F0 = X, (5a)
P (Y|FL+1) = ∏νL+1 λ=1 N ( yλ; f L+1 λ , σ 2I )
(5b)
where X ∈ RP×ν0 is the input and F` ∈ RP×ν` are the features. We use P for the number of input points and ν` for the width of layer `; thus ν0 is the number of inputs and νL+1 is the number of outputs. In addition, the features and targets can be written as a stack of vectors, f `λ ∈ RP and yλ ∈ RP , i.e.
F` = (f ` 1 f ` 2 · · · f `ν`) Y = (y1 y2 · · · yνL+1).
The function Kfeatures (F`−1) takes the features at the previous layer and returns the corresponding P × P kernel matrix. We consider isotropic kernels such as the squared exponential, which can be written as a function of R`−1ij , the distance between input features i and j,
Kfeatures;ij = k(R `−1 ij ),
R`−1ij = 1 N` ∑N` λ=1 ( F `−1iλ − F `−1jλ )2 .

heading: 2.4 Deriving equivalent deep Wishart processes
text: Following Aitchison et al. (2021), we show how the DGP model of Eq. (5) can be expressed as a deep Wishart process. We first consider the P × P Gram matrices defined as
G` = 1 ν` F`F T ` = 1 ν` ν∑̀ λ=1 f `λ(f ` λ) T ,
where f `λ are IID and multivariate-Gaussian distributed conditioned on the features at the previous layer (Eq. 5a). Thus, G` follows the definition of the Wishart (Eq. 1), and we can sample G` directly,
P (G`|F`−1) =W ( G`;
1 ν` Kfeatures(F`−1), ν`
) .
To work entirely with Gram matrices rather than features, we need to be able to compute the kernel, Kfeatures(F`−1) as a function of the Gram matrix at the previous layer, G`−1. Remarkably, this is possible for a large family of practically relevant kernels (Aitchison et al., 2021), for instance
isotropic kernels and the arc-cosine (or ReLU) kernel. In particular, as we focus on isotropic kernels, note that it is possible to recover distances from the Gram matrix:
R`ij = 1 N` ∑N` λ=1 (( F `iλ )2 − 2F `iλF `jλ + (F `jλ)2) = G`ii − 2G`ij +G`jj .
Thus, since isotropic kernels depend only on the distance, it is possible to obtain K(·), which takes the Gram matrix from the previous layer and returns the same kernel matrix as that returned by applying Kfeatures to the features from the previous layer:
Kfeatures(F`−1) = K(G`−1) = K( 1 ν` F`−1F T `−1).
By using the equivalent kernel written as a function of the Gram matrix at the previous layer, we can entirely eliminate intermediate layer features, resulting in a deep Wishart process,
P (G`|G`−1) =W ( G`;
1 ν` K(G`−1), ν` ) with G0 = 1ν0 XX T , (6a)
P (FL+1|GL) = ∏νL+1 λ=1 N ( fL+1λ ; 0,K (GL) ) , (6b)
P (Y|FL+1) = ∏νL+1 λ=1 N ( yλ; f L+1 λ , σ 2I ) . (6c)

heading: 2.5 The DWP formulation captures true-posterior symmetries while DGP does not
text: We now have two equivalent generative models: one phrased in terms of features, F` and another phrased in terms of Gram matrices, G`. Is there any reason to prefer one over the other? It turns out that there is. In particular, consider a transformation of the features, F′` = F`U where U is a unitary matrix, such that UUT = I. Remarkably, the true posterior is symmetric under these transformations, in the sense that all unitary transformations of the underlying features have the exact same true-posterior probability density (see Aitchison et al., 2021, Appendix D.2),
P (F′1, . . . ,F ′ L,FL+1|X,Y) = P (F1, . . . ,FL,FL+1|X,Y) .
It would be desirable for variational approximate posteriors to capture these true posterior symmetries. However, the usual family of Gaussian approximate posteriors over features fails to capture these symmetries because they use non-zero means. Worryingly, the failure to capture these symmetries can bias variational inference to focus on low-mass areas of the true posterior (Moore, 2016; Pourzanjani et al., 2017).
In contrast, the deep Wishart process sidesteps this issue by phrasing posteriors entirely in terms of Gram matrices, G` = 1ν`F`F T ` . Critically, the Gram matrix is invariant to unitary transformations of the features,
G` = 1 ν` F`F T ` = 1 ν` F`U`U T ` F T ` = 1 ν` F′`F ′T ` .
As such, DWP approximate posteriors written in terms of G` implicitly respect this unitary symmetry over the features.

heading: 2.6 Equivalent DWP posteriors will have better ELBOs and generalisation
text: Let us consider the ELBOs implied by using a DGP posterior Q ({F`}`)) and the DWP posterior implied by this posterior, which for this section we denote Q ({G`}`). Using D = (X,Y):
LDGP = EQ[P (Y|FL+1)]−KL(Q ({F`}`) ||P ({F`}`|D)), LDWP = EQ[P (Y|FL+1)]−KL(Q ({G`}`) ||P ({G`}`|D)).
By the data processing inequality (see e.g. Thm 6.2 in Polyanskiy & Wu (2014)), since G` is a deterministic transformation of F`, it is straightforward to show that LDWP ≥ LDGP. This result is one of the main motivations of the recent trend towards function-space inference (Sun et al., 2018; Ma et al., 2019); for a deeper theoretical understanding of this result we refer the reader to Burt et al. (2021), This fact can also be used to derive better PAC Bayes generalisation bounds (Section 6.1.3 of Alquier (2021)).
Note that this analysis relied on the assumption that Q ({G`}`) would be the distribution implied by Q ({F`}`)). In practice, this is will not be the case, as we wish to directly specify our distribution in terms of Gram matrices. However, given a sufficiently flexible posterior over Gram matrices, we should still see this advantages in practice. We now turn to developing such a distribution.

heading: 3 Methods
text: As detailed in Aitchison et al. (2021), the key difficulty in obtaining a variational inference scheme for DWPs is the difficulty of providing a sufficiently flexible approximate posterior. In particular, as we are working with a probabilistic process, the number of input points, P , can be arbitrarily large, and thus there is always the possibility that ν < P and hence that our sampled Gram matrices are low-rank. We therefore need to form flexible variational approximate posteriors over rank ν Gram matrices. An obvious first choice is the Wishart distribution itself with degrees of freedom ν, so as to match the rank of matrices sampled from the prior. However, for fixed degrees of freedom the Wishart variance,
V [Gij ] = ν ( Σ2ij + ΣiiΣjj ) cannot be specified independently of the mean (Eq. 3), which is essential for a variational approximate posterior that can flexibly capture potentially narrow true posteriors. An alternative approach would be to work with a non-central Wishart, which is defined by taking Ξ, which is IID standard Gaussian in the case of the Wishart, to have non-zero mean. However, the non-central Wishart has a probability density function that is too difficult to evaluate in the inner loop of a deep learning algorithm (Koev & Edelman, 2006). Instead, we develop a new Generalised Singular Wishart distribution, based on the Bartlett decomposition, which modifies the Wishart to give independent control over the mean and variance of sampled matrices.

heading: 3.1 The Generalised Singular Wishart
text: To define the Generalised Singular Wishart distribution, we first need to generalise the Bartlett construction to potentially singular matrices (i.e. those for which ν < P ). Remembering that Z = AAT , in the singular case A is given by
A =  A11 . . . 0 ... . . . ... Aν1 . . . Aνν
... ... ... AP1 . . . APν  , (7a) P ( A2jj ) = Gamma ( A2jj ; ν−j+1 2 , 1 2 ) , P (Ai>j) = N (Aij ; 0, 1) . (7b)
Recalling that G = LAATLT and by applying the results of Appendices B and C, we have that
P (G) =  P∏ j=1
1
L min(j,ν) jj min(P,ν)∏ j=1 Gamma ( A2jj ; ν−j+1 2 , 1 2 ) AP−jjj L P−j+1 jj P∏ i=j+1 N (Aij ; 0, 1) .
In Appendix D we prove that this corresponds to the known full rank and singular Wishart distribution. Equipped with the singular Bartlett, we can now develop a generalisation of the Wishart distribution:
Definition 1 The Generalised Singular Wishart, W (G; Σ, ν,α,β,µ,σ), is a distribution over positive semi-definite P × P matrices, G, with positive definite scale matrix Σ = LLT ∈ RP×P , a positive, integer-valued degrees-of-freedom parameter ν, and Bartlett-generalising parameters α,β,µ,σ. These latter parameters modify the Bartlett decomposition as follows:
Q ( A2jj ) = Gamma ( A2jj ;αj , βj ) for j ≤ ν,
Q (Ai>j) = N ( Aij ;µij , σ 2 ij ) for j ≤ ν.
This implies a distribution over G = LAATLT with density
Q (G) =  P∏ j=1
1
L min(j,ν) jj min(P,ν)∏ j=1
1
AP−jjj L P−j+1 jj
Gamma ( A2jj ;αj , βj ) P∏ i=j+1 N ( Aij ;µij , σ 2 ij ) .
Note that we use Q (·) to reflect the fact that we will use the Generalised Singular Wishart as the basis for our approximate posterior. The density is derived using the same transformations and Jacobians as for the singular Wishart above.

heading: 3.2 Full approximate posterior distribution
text: Unlike the deep inverse Wishart process in Aitchison et al. (2021), it is not possible to obtain an optimal last-layer posterior for the deep Wishart process. Therefore, we choose a form that mimics the form of the DIWP posterior, allowing for similar across-layer dependencies:
Q (G`|G`−1) =W ( G`; (1− q`) 1ν`K(G`−1) + q`V`V T ` , ν`, α`, β`, µ`, σ` ) , (9)
where the approximate posterior parameters are {V`,α`,β`,µ`,σ`, q`}L`=1, where 0 < q` < 1 is a scalar, and V` ∈ RP×P . Here, the learnable parameter q` allows us to trade the off influence on the scale matrix from the prior and the learned covariance V`VT` , which allows for similar across-layer dependencies as in (Aitchison et al., 2021). ν` is fixed, as it determines the width of the layer; the remaining parameters are simply the parameters from the Bartlett generalisation.

heading: 3.3 Doubly stochastic inducing-point variational inference in deep inverse Wishart processes
text: For efficient inference in high-dimensional problems, we take inspiration from the DGP literature (Salimbeni & Deisenroth, 2017) by considering doubly-stochastic inducing-point deep Wishart processes. We begin by decomposing all variables into inducing and training (or test) points Xi ∈ RPi×N0 and Xt ∈ RPt×N0 where Pi is the number of inducing points, and Pt is the number of testing/training points,
X = ( Xi Xt ) , FL+1 = ( FL+1i FL+1t ) , G` = ( G`ii G ` it G`ti G ` tt ) ,
where e.g. G`ii is Pi × Pi and G`it is Pi × Pt. The full ELBO including latent variables for all the inducing and training points is
L = E [ log P (Y|FL+1) + log P ( {G`}L`=1,FL+1|X ) Q ( {G`}L`=1,FL+1|X )] , (10) where the expectation is taken over Q ( {G`}L`=1,FL+1|X ) . The prior is given by combining all terms in Eq. (6) for both inducing and test/train inputs,
P ( {G`}L`=1,FL+1|X ) = [∏L `=1 P (G`|G`−1) ] P (FL+1|GL) ,
where the X-dependence enters on the right because G0 = 1ν0 XX T . Taking inspiration from Salimbeni & Deisenroth (2017), the full approximate posterior is the product of an approximate posterior over inducing points and the conditional prior for train/test points,
Q ( {G`}L`=1,FL+1|X ) =
Q ( {G`ii}L`=1,FL+1i |Xi ) P ( {G`it}L`=1, {G`tt}L`=1,FL+1t |{G`ii}L`=1,FL+1i ,X ) . (11)
And the prior can be written in the same form, P ( {G`}L`=1,FL+1|X ) =
P ( {G`ii}L`=1,FL+1i |Xi ) P ( {G`it}L`=1, {G`tt}L`=1,FL+1t |{G`ii}L`=1,FL+1i ,X ) . (12)
We discuss the second terms (the conditional prior) in Eq. (15). The first terms (the prior and approximate posteriors over inducing points), are given by combining terms in Eq. (6) and Eq. (9),
P ( {G`ii}L`=1,FL+1i |Xi ) = [∏L `=1 P ( G`ii|G`−1ii )] P ( FL+1i |GLii ) , (13)
Q ( {G`ii}L`=1,FL+1i |Xi ) = [∏L `=2 Q ( G`ii|G`−1ii )] Q ( FL+1i |GLii ) . (14)
Substituting Eqs. (11–14) into the ELBO (Eq. 10), the conditional prior cancels and we obtain,
L = E log P (Y|FL+1t )+ log [∏L `=1 Q ( G`ii|G`−1ii )] Q ( FL+1i |GLii )[∏L `=1 P ( G`ii|G`−1ii )] P ( FL+1i |GLii )  .
The first term is a summation across test/train datapoints, and the second term depends only on the inducing points, so as in Salimbeni & Deisenroth (2017) we can compute unbiased estimates of the expectation by taking only a minibatch of datapoints. We also never need to compute the density of the conditional prior in Eq. (12), we only need to be able to sample from it,
P ( {G`ti,G`tt}L`=1,FL+1t |{G`ii}L`=1,FL+1i ,X ) =
P ( FL+1t |FL+1i ,GL )∏L `=1 P ( G`ti,G ` tt|G`ii,G`−1 ) . (15)
The first distribution, P ( FL+1t |FL+1i ,GL ) , is a multivariate Gaussian, and can be evaluated using methods from the GP literature (Rasmussen & Williams, 2006; Salimbeni & Deisenroth, 2017). Specifically, we use the global inducing point scheme from Ober & Aitchison (2021). The second distribution is more difficult to sample from. To address this issue, we go back to features,
F`F T ` = νG` ∼ W (K (G`−1) , ν) , (16)
with
F` = ( F`i F`t ) K (G`−1) = ( Kii K T ti Kti Ktt ) ,
where F` ∈ R(Pi+Pt)×ν` , Fi ∈ RPi×ν` and Ft ∈ RPt×ν` . Our goal is to sample G`it and G`tt given G`ii. Our approach is to note that, Ft conditioned on Fi is given by a matrix normal, (Eaton et al., 2007, page 310),
P ( F`t ∣∣F`i ) =MN (KTti K−1ii Fi,Ktt·i, I) , (17)
where
Ktt·i = Ktt −KTit K−1ii Kit. Note that we sample each test/train point one-at-a-time/independently, in which case, Pt = 1 and Σ22·1 is scalar.
Then G`, which includes G`it and G ` tt is given by,
G` =
( G`ii G ` it
G`ti G ` tt
) = 1ν ( F`i ( F`i )T F`i ( F`t )T
F`t ( F`i )T F`t ( F`t )T ) = 1νF`F T `
For Fi, we can use any value as long as G`ii = F ` i ( F`i )T
, as the resulting distribution over G` arising from Eq. (16) does not depend on the specific choice of Fi (App. E). Remembering that to sample Gii, we explicitly sample its potentially low-rank Cholesky, L`A`, we can directly use
F`i = L`A`
However, this only works if ν ≤ Pi, in which case, L`A` ∈ RPi×ν . In the unusual case where we have fewer inducing points than degrees of freedom, Pi < ν, then L`A` ∈ RPi×Pi , so we need to pad to achieve the required size of Pi × ν`,
F`i = (L`A` 0) .
Finally, note that we can optimise all the variational parameters using standard reparameterised variational inference (Kingma & Welling, 2014; Rezende et al., 2014). For an algorithm, see Alg. 1.

heading: 3.4 Computational complexity
text: Recalling that ν` is the width of the `th layer, Pi is the number of inducing points, and Pt is the number of train or test points, the computational complexity of one DWP layer is given by O(P 3i + PtP 2 i ). This is a decrease of a factor of ν`+1 over the complexity for standard DGP inference, such as doubly stochastic variational inference (Salimbeni & Deisenroth, 2017), which has complexity O(ν`+1(P 3 i + PtP 2 i )). The difference arises from the fact that in a DGP, ν`+1 Gaussian processes are sampled in each layer, whereas for a DWP we sample a single Gram matrix.
Algorithm 1 Computing predictions/ELBO for one batch
P parameters: {ν`}L`=1. Q parameters: {V`, q`,α`,β`,µ`,σ`}L`=1,Xi. Inputs: Xt; Targets: Y combine inducing and test/train inputs X = (Xi Xt) sample first Gram matrix and update ELBO G0 =
1 ν0 XXT
for ` in {1, . . . , L} do sample inducing Gram matrix and its Cholesky, L`A` and update ELBO L`A`A T ` L T ` = G ` ii ∼ Q ( G`ii|G`−1ii ) L ← L+ log P ( G`ii|G`−1ii ) − log Q ( G`ii|G`−1ii
) sample full Gram matrix from conditional prior Σ = 1ν`K(G`−1) Σtt·i = Σtt −ΣTit Σ−1ii Σit F`i = L`A`
F`t ∼MN ( ΣTti Σ −1 ii F̂i,Σtt·i, I ) G` = ( G`ii F̂ ` i (F̂ ` t ) T
F̂`t (F̂ ` i ) T F̂`t (F̂ ` t ) T ) end for sample GP inducing outputs and update ELBO FL+1i ∼ Q ( FL+1i |GLii
) L ← L+ log P ( FL+1i |GLii ) − log Q ( FL+1i |GLii
) sample GP predictions conditioned on inducing points FL+1t ∼ Q ( FL+1t |GL,FL+1i
) add likelihood to ELBO L ← L+ log P ( Y|FL+1t )

heading: 4 Results
text: The DWP prior is equivalent to a DGP prior (Sec. 2.4) (Aitchison et al., 2021); the only difference is that in a DGP, we use features as the latent variables, whereas in the DWP we use Gram matrices. Using Gram matrices as in the DWP should be beneficial as the true posteriors are expected to be simpler than in the DGP (Sec. 2.5). We perform a comparison of the posterior intermediate-layer features implied by a DWP and those of a DGP on a toy problem in App. G; these show the benefits of being more able to model the true posterior.
For more quantitative experiments, we trained a DWP and a DGP with the exact same generative model with a squared exponential kernel on the UCI datasets from Gal & Ghahramani (2015). We trained both models for 20000 gradient steps using the Adam optimizer Kingma & Ba (2015); we detail the exact experimental setup in Appendix F. We report ELBOs and test log likelihoods for depth 5 in Table 1; we report other depths and quote the relevant results from Aitchison et al. (2021) for reference in Appendix H. We found that the DWP often outperformed the DGP model, especially evident if we look at the ELBOs and smaller datasets (boston, concrete, energy, and yacht). On larger datasets, the benefits often disappear, as accurate uncertainty modelling is less relevant. Note that we compared against the recently introduced DGP method based on global inducing points (Ober & Aitchison, 2021). Global inducing point methods were particularly important in our setting because we use a standard feedforward architecture without skip connections to ensure equivalence between the DGP and DWP. Standard DSVI has considerable difficulties with optimizing the approximate posterior in such models; to get optimization to work effectively Salimbeni & Deisenroth (2017) were forced to modify the prior to introduce skip connections.

heading: 4.1 Runtimes & training curves
text: In Sec. 3.4, we showed that DWPs have a lower computational complexity than DGPs, because of the need for DGPs to sample ν` features in each layer, whereas DWPs only need to sample one Gram matrix. Here, we briefly discuss the runtimes of our implementations. We show a plot of the training curves for one split of boston with a 5-layer DGP and DWP in Fig. 1, plotted against both runtime and epoch. From these plots, we make two observations. First, the DWP trains much more quickly than the DGP in terms of runtime. However, it seems to require slightly more epochs than the DGP to converge (note that the spike at the start of the DGP curve is an artifact of the tempering scheme we use). In Appendix H, we provide a table of time per epoch, which shows that we obtain faster runtime for protein and for shallower models, although the gains are slightly more modest due to the models being shallower and the fact that we run protein on a GPU, as opposed to a CPU for boston.

heading: 5 Related Work
text: The DWP prior was introduced by Aitchison et al. (2021). However, they were not able to do variational inference with the DWP because they did not have a sufficiently flexible approximate posterior over positive semi-definite matrices. Instead, they were forced to work with a deep inverse Wishart process, which is easier because the inverse Wishart itself is a suitable approximate posterior. Here, we give a flexible generalised Wishart distribution over positive semi-definite matrices which is suitable for use as a variational approximate posterior in the DWP. As the deep Wishart process prior is equivalent to a DGP prior, we were able to directly compare DGP and DWP inference in models with the exact same prior. Such a comparison with equivalent priors was not possible in Aitchison et al. (2021), because their deep inverse Wishart process priors are not equivalent to DGP priors.
There is an alternative line of work using generalised Wishart processes (Wilson & Ghahramani, 2011, as opposed to our deep Wishart processes). Note that the “generalised Wishart process” terminology does not seem to have spread as widely as it should, but it is very useful in our context. A generalised Wishart process specifies a distribution over an infinite number of finitedimensional Wishart-distributed matrices. These matrices might represent e.g. the noise covariance in a dynamical system, in which case there might be an infinite number of such matrices, one for each time or location in the state-space (Wilson & Ghahramani, 2011; Heaukulani & van der Wilk, 2019; Jorgensen et al., 2020). In contrast, the Wishart process (Dawid, 1981; Bru, 1991) describes finite dimensional marginals of a single, potentially infinite dimensional matrix. In our context, we stack
(non-generalised) Wishart processes to form a deep Wishart process. Importantly, these generalised Wishart priors do not have the flexibility to capture a DGP prior because the underlying features at all locations are jointly multivariate Gaussian (Sec. 4 in Wilson & Ghahramani, 2011) and therefore lack the required nonlinearities between layers. Further, not only do the underlying stochastic processes (deep vs generalised Wishart process) differ, inference is also radically different. In particular, work on the generalised Wishart performs inference on the underlying multivariate Gaussian feature vectors (Eq. 1 e.g. Eq. 15-18 in Wilson & Ghahramani 2011, Eq. 12 in Heaukulani & van der Wilk 2019 or Eq. 24 in Jorgensen et al. 2020). Unfortunately, variational approximate posteriors defined over multivariate Gaussian feature vectors fail to capture symmetries in the true posterior (Sec. 2.5). In contrast, we define approximate posteriors directly over the symmetric positive semi-definite Gram matrices themselves, which required us to develop new, more flexible distributions over these matrices.

heading: 6 Limitations
text: There are a few limitations of our work. First, it is only possible to derive equivalent DWPs for certain kernels - namely, those where we can skip the feature representation and work entirely in Gram matrices. While this holds for a large range of kernels, such as isotropic kernels and the arc-cosine kernel, it does not hold for some common kernels such as automatic relevance determination (ARD) kernels. However, we note that we are able to use any kernel for the first layer, and that in practice we did not find that ARD kernels in intermediate layers significantly improved performance, as all the features have a shared prior. The second limitation is that it is not currently possible to incorporate modifications to the basic DGP model such as skip connections (Duvenaud et al., 2014; Salimbeni & Deisenroth, 2017). Such modifications would require the use of the non-central Wishart distribution, which is difficult to evaluate in the inner loop of a deep learning algorithm (Koev & Edelman, 2006). Third, the generalisation of our work to more complex architectures such as convolutional models is non-trivial, and will be left to future work. Finally, it seems that the performance is not as competitive for larger datasets, where uncertainty representation is of more limited use - we will explore improving this in future work.

heading: 7 Conclusions
text: We introduced a flexible distribution over positive semi-definite matrices which formed the basis of a variational approximate posterior for the deep Wishart process. We adapted the doubly stochastic variational inference scheme from Aitchison et al. (2021) to the deep Wishart process. Thus, we were able to directly compare the performance for inference in a DWP vs. DGP with exactly the same prior. This isolates the effects on performance of the prior and the inference procedure. We found DWPs often have improved performance over their corresponding equivalent DGP models, particularly in the ELBOs, indicating improved uncertainty representation.
There are no anticipated social impacts as the work is largely theoretical.

heading: Acknowledgements
text: The authors thank David R. Burt, Mark van der Wilk, Samuel Power, and Adam X. Yang for helpful discussions, and would like to thank the reviewers for their comments.
Funding Transparency
SWO acknowledges funding from the Gates Cambridge Trust for his doctoral studies. The authors have no competing interests or additional funding to declare.

Review summaries:
The authors introduce effective inference strategies for deep Wishart processes (DWP). This class of deep models is of interest because it moves away from the featurized representation implicit in most deep models to one that is fully kernelized. Inference relies on a doubly stochastic inducing-point scheme familiar from variational Deep Gaussian Processes and makes use of a custom generalization of the Bartlett decomposition adapted to the DWP setting in order to define a flexible variational family. In experiments on UCI data the method shows small but consistent improvements over Deep Gaussian Processes.

The authors introduce an approximate inference procedure for the deep Wishart process. In order to do so, they define a generalised singular Wishart process and use this as a variational posterior. The DWP appears to be a model consisting of layers of Wishart process, which are distributions over PSD matrices. This allows one to define a Wishart process whose parameters are the output of another Wishart process. When the a deep GP is constructed with a kernel that is isotropic, the, DGP is a DWP.  However, the usual approximate posterior used in the DGP fails to capture a certain orthogonal invariance property in the features that is present in the true posterior. The DWP on the other hand, deals with Gram matrices instead of features and this requirement from the approximate posterior is automatically satisfied. 

Deep kernel process (DKP, Aitchison et al., 2020) is a new probabilistic model that provides more flexible modelling power than a shallow kernel method. The deep Wishart process is equivalent to deep Gaussian process, but there is no inference algorithm in the original paper. This paper proposes a new variational inference algorithm for it. It proposes a generalised singular Wishart distribution to model the posterior distribution of the random Gram matrix in DWPs and apply the doubly-stochstic inducing-point scheme. This is a nice contribution to the progress of research along the line of DKPs.