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