A variational approximate posterior for the deep Wishart process

1 Introduction . 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 . 2 Background . 2.1 Wishart distribution . 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 ) 2.2 Bartlett decomposition . 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 . 2.3 Deep Gaussian proceses ( DGPs ) . 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 . 2.4 Deriving equivalent deep Wishart processes . 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 )

Summary: 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.