Deep Bayesian Structure Networks

1 INTRODUCTION . Bayesian deep learning aims at equipping the flexible and expressive deep neural networks with appropriate uncertainty quantification ( MacKay , 1992 ; Neal , 1995 ; Hinton & Van Camp , 1993 ; Graves , 2011 ; Blundell et al. , 2015 ; Gal & Ghahramani , 2016 ) . Traditionally , Bayesian neural networks ( BNNs ) introduce uncertainty in the network weights , addressing the over-fitting issue which standard neural networks ( NNs ) are prone to . Besides , the predictive uncertainty derived from the weight uncertainty is also of central importance in practical applications , e.g. , medical analysis , automatic driving , and financial tasks . Modeling the uncertainty on network weights is plausible and well-evaluated ( Blundell et al. , 2015 ; Ghosh et al. , 2018 ) . However , BNNs usually preserve benefits of Bayesian principles such as wellcalibrated predictions at the expense of compromising performance and hence are impractical in real-world applications ( Osawa et al. , 2019 ) , due to various reasons . On one hand , specifying a sensible prior for networks weights is difficult ( Sun et al. , 2019 ; Pearce et al. , 2019 ) . On the other hand , the flexible variational posterior of BNNs comes with inference challenges ( Louizos & Welling , 2017 ; Zhang et al. , 2018 ; Shi et al. , 2018 ) . Recently , the efficient particle-based variational methods ( Liu & Wang , 2016 ) have been developed with promise , but they still suffer from the particle collapsing and degrading issues for BNNs due to the high dimension of the weights and the overparameterization nature of such models ( Zhuo et al. , 2018 ; Wang et al. , 2019 ) . In this work , we investigate a new direction of Bayesian deep learning that performs Bayesian reasoning on the structure of neural networks while keeping the weights as point estimates . We propose an approach , named Deep Bayesian Structure Networks ( DBSN ) . Specifically , in the spirit of differentiable neural architecture search ( NAS ) ( Liu et al. , 2019 ; Xie et al. , 2019 ) , DBSN builds a deep network by repeatedly stacking a computational cell in which any two nodes ( i.e . tensors ) are connected by redundant transformations ( see Figure 1 ) . The network structure is defined as the gating weights on these transformations , whose distribution is much easier to capture than those of the high-dimensional network weights . To jointly optimize the network weights and the parameterized distribution of the network structure , we adopt a stochastic variational inference paradigm ( Blundell et al. , 2015 ) and use the reparameterization trick ( Kingma & Welling , 2013 ) . One technical challenge is driving DBSN to achieve satisfying convergence , since the network weights can hardly fit all the structures sampled from the structure distribution . To overcome this challenge , we propose two techniques . First , we advocate reducing the variance of the sampled structures with a simple mod- ification of the sampling procedure . Second , we suggest using a more compact structure learning space than that of NAS , to make the training more feasible and more efficient . There are at least two motivations that make DBSN an appealing choice : 1 ) DBSN bypasses the frustrating difficulties of characterizing weight uncertainty and enables the performance-enhancing structure learning ( Zoph & Le , 2016 ; Liu et al. , 2019 ) , so DBSN shall have better predictive performance than classic BNNs . 2 ) Previous analysis ( Wang et al. , 2019 ) shows that due to the overparametrization nature of BNNs , the state-of-the-art inference algorithms for weight uncertainty can suffer from mode collapsing , as multiple configurations of weights with a fixed structure correspond to one single function . In contrast , DBSN compactly models the uncertainty of structure and performs inference in a much lower-dimensional space , avoiding this issue and hence being able to exhibit more calibrated predictive uncertainty . Moreover , in the perspective of NAS , DBSN is also promising as it provides another principled way to learn network structures by resorting to the Bayesian formalism instead of the widely used meta-learning formalism in differentiable NAS . To empirically validate these hypotheses , we evaluate DBSN with extensive experiments . We first testify the data fitting and structure learning ability of DBSN on challenging classification and segmentation tasks . Then , we compare the quality of predictive uncertainty estimates via calibration , which is a common concern in the community . We further evaluate the predictive uncertainty on adversarial examples and out-of-distribution samples , drawn from shifted distributions from the training data , to verify whether the model knows what it knows . At last , we perform an experiment to validate a promising application of DBSN in the one-shot NAS ( Bender et al. , 2018 ; Guo et al. , 2019 ) . Surprisingly , across all the tasks , DBSN consistently achieves comparable or even better results than the strong baselines . 2 BACKGROUND . We first review the necessary background for DBSN and then elaborate DBSN in the next section . 2.1 STOCHASTIC VARIATIONAL INFERENCE FOR BNNS . Let D = { ( xi , yi ) } Ni=1 be a set of N data points . BNNs are typically defined by placing a prior p ( v ) on some variables of interest ( e.g. , network weights or network structure ) and the likelihood is p ( D|v ) . Directly inferring the posterior distribution p ( v|D ) is intractable because it is hard to integrate w.r.t . v exactly . Instead , variational BNNs ( Hinton & Van Camp , 1993 ; Graves , 2011 ; Blundell et al. , 2015 ) suggest approximating p ( v|D ) with a θ-parameterized distribution q ( v|θ ) by minimizing the Kullback-Leibler ( KL ) divergence between them : min θ DKL ( q ( v|θ ) ‖p ( v|D ) ) = −Eq ( v|θ ) [ log p ( D|v ) ] +DKL ( q ( v|θ ) ‖p ( v ) ) + log p ( D ) , ( 1 ) where log p ( D ) is a constant w.r.t . θ and usually omitted in the minimization . To solve problem ( 1 ) , the most commonly used method is the low-variance reparameterization trick ( Kingma & Welling , 2013 ; Blundell et al. , 2015 ) , which replaces the sampling procedure v ∼ q ( v|θ ) with the corresponding deterministic transformation v = t ( θ , ) with a sample of parameter-free noise , to enable the direct gradient back-propagation through θ . 2.2 CELL-BASED DIFFERENTIABLE NEURAL ARCHITECTURE SEARCH ( NAS ) . Cell-based NAS has shown promise ( Zoph et al. , 2018 ; Pham et al. , 2018 ) and been developed to be differentiable for better scalability ( Liu et al. , 2019 ; Xie et al. , 2019 ; Weng et al. , 2019 ) . Generally , the network in cell-based differentiable NAS1 is composed of a sequence of cells ( e.g. , modules ) which have the same internal structure and are separated by upsampling or downsampling modules . Every cell contains B sequential nodes ( i.e. , tensors ) : N1 , . . . , NB . Each node N j is connected to all of its predecessors N i so long as i < j by K possible redundant operations o ( i , j ) 1 , . . . , o ( i , j ) K , e.g. , convolution , skip connection , pooling . The network structure is defined as α = { α ( i , j ) |1 ≤ i < j ≤ B } where α ( i , j ) ∈ ∆K−1 corresponds to the gating weights on the K available operations from N i to N j . Therefore , the information gathered from N i to N j is a weighted sum of the outputs from K different operations on N i ( we denote the set including the parameters of all the operations in the network as w ) : N ( i , j ) = K∑ k=1 α ( i , j ) k · o ( i , j ) k ( N i ; w ) . ( 2 ) Then , the nodeN j is calculated by summing all the information from its predecessors : N j = ∑ i < j N ( i , j ) . ( 3 ) Meta-learning-like gradient descent is adopted for optimization to reduce the prohibitive computational cost needed by RL or evolution ( Liu et al. , 2019 ; Xie et al. , 2019 ) . However , the goal of the optimization is the network structure instead of the model performance . Thus , after training , this kind of NAS needs to prune the searched structure and re-train a new network model with the compact structure for performance comparison , which is labor-intensive and is avoided in our work . 3 DEEP BAYESIAN STRUCTURE NETWORKS . In this work , we propose a novel Bayesian structure learning approach for the deep neural networks . Concretely , we follow the network design of NAS but we view α as Bayesian variables and w as point estimates ( see the graphical model in Figure 1 ) . To infer the posterior distribution p ( α|D , w ) = p ( α ) p ( D|α , w ) p ( D ) , where p ( α ) is the prior ( we omit its dependency on the hyperparameter θ0 here ) , we adopt the techniques in Section 2.1 . We assume both the prior and the introduced variational are fully factorizable categorical distributions , namely , p ( α ) = ∏ i < j p ( α ( i , j ) ) and q ( α|θ ) = ∏ i < j q ( α ( i , j ) |θ ( i , j ) ) , where θ = { θ ( i , j ) ∈ RK |1 ≤ i < j ≤ B } denotes the trainable categorical logits . We rewrite Eq . ( 1 ) and obtain the negative evidence lower bound ( ELBO ) : L ( θ , w ) = −Eq ( α|θ ) [ log p ( D|α , w ) ] +DKL ( q ( α|θ ) ‖p ( α ) ) . ( 4 ) Notably , minimizing L w.r.t . θ and w corresponds to Bayesian inference on α and maximum a posteriori ( MAP ) estimation of w2 , respectively . Thus , the optimization of the network structure and network weights can be unified as minθ , w L ( θ , w ) . To resolve this , we relax both p ( α ( i , j ) ) and q ( α ( i , j ) |θ ( i , j ) ) to be the concrete distributions ( Maddison et al. , 2016 ) . Then , samples α from q ( α|θ ) are generated via the softmax transformation : α = g ( θ , ) = { softmax ( ( θ ( i , j ) + ( i , j ) ) /τ ) } , ( 5 ) where = { ( i , j ) ∈ RK | ( i , j ) k ∼ Gumbel i.i.d . } are the Gumbel variables and τ ∈ R+ is the temperature . Then we derive the following gradient estimators : ∇θL ( θ , w ) = E [ −∇θ log p ( D|g ( θ , ) , w ) +∇θ log q ( g ( θ , ) |θ ) −∇θ log p ( g ( θ , ) ) ] , ( 6 ) ∇wL ( θ , w ) = E [ −∇w log p ( D|g ( θ , ) , w ) ] . ( 7 ) The first term in Eq . ( 6 ) corresponds to the gradient of the negative log likelihood and we leave how to estimate the last two terms ( i.e . log densities ) in the next section . In practice , we approximate the expectation in Eq . ( 6 ) and Eq . ( 7 ) with T Monte Carlo ( MC ) samples , and update the structure and the weights w simultaneously . After training , we gain the following predictive distribution : p ( y|xnew , w∗ ) = Eq ( α|θ∗ ) [ p ( y|xnew , α , w∗ ) ] , ( 8 ) where θ∗ and w∗ denote the converged parameters . Eq . ( 8 ) implies that the model predicts by ensembling the predictions of the networks whose structures are randomly sampled . 1We will refer to the cell-based differentiable NAS as NAS for short if there is no misleading . 2This is because we use regularizor on weights , e.g. , weight decay , to alleviate over-fitting .

Summary: The paper combines ideas from neural architecture search (NAS) and Bayesian neural networks. Instead of maintaining uncertainty in network weights, the authors propose to retain uncertainty in the network structure. In particular, building on cell-based differentiable NAS, the authors infer a distribution over the gating weights of different cells incident onto a tensor while relying on point estimates for the weights inside each cell.