Joint Inference for Neural Network Depth and Dropout Regularization

1 Introduction . Both dropout regularization and network depth are critical to the success of deep neural networks ( DNNs ) [ 1 , 2 , 3 , 4 ] . Finely tuned or selected structures empower DNNs with proper model capacity that can not only efficiently capture statistical regularities in data but also avoid being deceived by random noise into “ discovering ” non-existent relationships . Dropout as an effective regularization method sparsifies a neural network ’ s structure by randomly deleting neurons along with their connections in training to prevent it from overfitting [ 5 , 6 ] . Recent studies extend the method from a Bayesian perspective to allow automatic tuning of the dropout probability in large models [ 4 , 7 , 8 , 9 , 10 , 11 ] . Although dropout and its variants can effectively govern model capacity , without uncertainty calibration deep models tend to be overly confident with their predictions [ 3 , 12 ] . Probabilistic inference approaches for network structure selection propose to bypass or down-weight certain hidden layers to reduce training cost and build smaller models that perform just as well as larger ones [ 3 , 13 , 14 , 15 ] . These methods lead to more efficient or better-calibrated models only by reducing the depth of a pre-determined network structure . They can not scale the network structure up to accommodate incrementally available data beyond the upper limit of its capacity . We thus propose a novel Bayesian model selection framework to jointly infer network depth and perform dropout regularization simultaneously . In particular , we model the depth of a network structure as a stochastic process by defining the beta process [ 16 , 17 ] over the number of hidden layers to enable it to go to infinity in theory , as in Figure 1 ( a ) ( left ) . The beta process induces layer-wise activation probabilities , which allows its conjugate Bernoulli process to generate a binary ∗Corresponding author 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . vector per hidden layer to prune neurons for regularization , as in Figure 1 ( b ) . The probabilistic inference framework automatically balances network depth and dropout regularization by computing a marginal likelihood over the hidden layers and their neuron activations , and provides well-calibrated uncertainty estimates for predictions . The exact computation of the marginal likelihood is intractable due to the nonlinear nature of neural networks with a potentially infinite number of hidden layers . We thus employ the structured stochastic variational inference [ 18 , 19 ] with a continuous relaxation on the Bernoulli variables to efficiently approximate the integral with a continuum of lower bounds . The relaxation allows to maintain differentiability and mitigate overfitting via model averaging of the sampled hidden layers and neuron activations . Benefiting from theoretical analysis , we readily integrate the joint inference with the parameter learning process through an alternating approach to efficiently update the parameters , and perform hidden-layer sampling and dropout regularization . We analyze the behavior of our joint inference framework over multilayer perceptrons ( MLPs ) and convolutional neural networks ( CNNs ) , and evaluate their performance across domains . The experiments show that our method leads to a compact neural network by balancing its depth and dropout regularization with uncertainty calibration , and achieves superior performance comparing to state-of-the-art dropout and structure selection methods . We also demonstrate our method on a continual learning task . By enabling both network depth and neuron activations to dynamically evolve to accommodate incrementally available data , we can alleviate catastrophic forgetting . 2 Related work . Dropout is an effective neural network regularization technique [ 6 , 21 ] . By randomly pruning neurons and their connections from a network structure with a dropout rate during training , it can prevent DNNs from overfitting [ 22 , 23 , 24 ] . [ 7 ] interprets dropout from a Bayesian perspective to quantify uncertainty . Variational dropout proposes a stochastic gradient variational inference approach to learn the dropout rate from data with a constraint on large dropout rate values [ 8 ] . [ 9 ] extends the variational dropout to set unbounded individual dropout rate per weight . Concrete dropout proposes a continuous relaxation of the dropout ’ s discrete masks to automatically tune the dropout probability , and obtain its uncertainty estimates [ 10 ] . Bayesian nonparametrics are also applied to extend the variational dropout approaches by defining an Indian buffet process ( IBP ) over neurons per hidden layer [ 11 ] or over channels in CNNs [ 20 ] to infer network widths , as in Figure 1 ( a ) ( middle ) . Another body of work leverage the non-parametric Bayesian framework to infer the dimensionality of latent representation encoded by DNNs [ 25 , 26 ] . Structure selection methods improve training efficiency by reducing the depth of a pre-arranged network structure , as in Figure 1 ( a ) ( right ) . Some lead to well-performed smaller networks , and others achieve better uncertainty calibration . In order to build smaller models that perform just as well , [ 2 ] proposes to reduce a large neural network by merging adjacent hidden layers with only linear relationships being activated in between . A stochastic depth method randomly drops a subset of hidden layers from a large network by bypassing them with an identity function in each mini-batch to speed up the training session , but it still deploys the large network structure at test time [ 13 ] . [ 15 ] shows that the multiplicative noise from dropout induces structured shrinkage priors over a neural network ’ s weights , and they further extend the shrinkage framework based on ResNet structures [ 27 ] to model the probabilities of hidden layers being used . In particular , the work also indicates that the framework is equivalent to the stochastic depth regularization [ 13 ] . [ 14 ] proposes to tune a hidden layer ’ s influence on a prediction by learning a bypass variable between adjacent-layer connections and skip connections with a variational Bayes method . The method also prunes neurons separately . [ 3 ] proposes a Bayesian model averaging method to down-weight deeper layers by directly connecting every hidden layer to the output layer , and achieve competitive performances with uncertainty calibration . [ 28 ] achieves model effectiveness and computational efficiency by inferring a distribution of connections and units in the context of local winner-takes-all DNNs with IBP . 3 The joint inference framework . We propose to model network depth as a stochastic process , and jointly perform dropout regularization upon the inferred hidden layers . Theoretical analysis shows that by approximating its marginal likelihood the inference framework achieves an optimal balance between network depth and neuron activations as Bayesian model selection . 3.1 Network structure with infinite hidden layers . Let hl denote the l-th hidden layer composed of neurons ( i.e. , non-linear activation functions ) σ ( · ) . The neural network has the form : hl = σ ( Wlhl−1 ) ⊗ zl + hl−1 l ∈ { 1 , 2 , ... , ∞ } ( 1 ) where Wl ∈ RM×M is the layer l ’ s weight matrix with a Gaussian prior p ( W ) = N ( W|0 , s2wI ) , ⊗ denotes element-wise multiplication of two vectors , and M is the maximum number of neurons in a layer . For simplicity , we set M to be the same for all hidden layers . In particular , we prune the l-th layer ’ s outputs by multiplying them elementwisely by a binary vector zl where its element variable zml ∈ { 0 , 1 } . Each random variable zml takes the value 1 with πl ∈ [ 0 , 1 ] indicating activation probability of the l-th layer , as in Figure 1 ( b ) . The combination of the previous layer ’ s outputs with the current layer ’ s via skip connections not only avoids vanishing gradient but also propagates the output of the last hidden layer with activated neurons directly to the output layer f ( · ) .2 Given a dataset D = { ( xn , yn ) } Nn=1 with N input-output pairs ( xn , yn ) as training examples , for a regression task we express the likelihood of the neural network as : p ( D|Z , W ) = N∏ n=1 N ( yn|f ( xn ; Z , W ) , s2I ) ( 2 ) where Z is a binary matrix whose l-th column is zl , and W = { Wl } denotes the set of the weight matrices . s2 denotes likelihood noise with I as an identity matrix with the same dimensionality as the training examples ’ outputs . For classification tasks , we replace f ( · ) with a softmax function and the normal distribution with a multinoulli distribution . 3.2 Beta process over layer number . We treat the number of hidden layer functions hl as a stochastic process by defining a beta process over their space H = { hl } ∞l=1 , as shown in Figure 1 ( b ) . Conceptually , we generate a beta process B from its base measure B0 as B ∼ BP ( c , B0 ) by drawing a set of samples ( hl , πl ) ∈ H× [ 0 , 1 ] from a non-homogeneous Poisson process [ 16 ] . Let B = ∑ l πlδhl , where δhl is a unit point mass at hl . In our setting , a pair ( hl , πl ) corresponds to a hidden layer function hl ∈ H and its activation probability πl ∈ [ 0 , 1 ] . We define a conjugate 2A graphical demonstration of the network structure is in Appendix . Bernoulli process Zm· ∼ BeP ( B ) as Zm· = ∑ l zmlδhl at the same locations δhl as B where zml are independent Bernoulli variables with a probability that zml = 1 equals to πl . As in ( 1 ) , zml = 1 activates the m ’ th neuron in layer l. Computationally , we employ the stick-breaking construction of beta-Bernoulli processes [ 29 , 17 ] as zml ∼ Ber ( πl ) , πl = l∏ j=1 νj , νl ∼ Beta ( α , β ) ( 3 ) where νl are sequentially drawn from a beta distribution . According to the formulation , πl is decreasing with l. α and β are the hyperparameters governing preference over a balance between network depth and dropout regularization . In particular , if α > β > 1 , the hidden layers tend to have lower activation probabilities with less number of activated neurons , as in Figure 1 ( b ) left . The setting thus prefers a narrower but deeper network structure . On the other hand , β > α > 1 favors shallower but wider network structure , where fewer hidden layers tend to have higher activation probabilities with more neurons being activated , as in Figure 1 ( b ) right . We thus define a prior over Z via the beta process as p ( Z , ν|α , β ) = p ( ν|α , β ) p ( Z|ν ) = ∞∏ l=1 Beta ( νl|α , β ) M∏ m=1 Ber ( zml|πl ) ( 4 ) 3.3 Marginal likelihood over network structures . By combining the beta process prior over hidden layers in ( 4 ) with the neural network structure in ( 2 ) , we can infer the number of hidden layers L by integrating over Z = { zl } Ll=1 and ν = { νl } Ll=1 : p ( D|W , L , α , β ) = ∫ p ( D|Z , W ) p ( Z , ν|α , β ) dZdν ( 5 ) The exact computation of this marginal likelihood is intractable due to the non-linearity of the neural network and L→∞ .

Summary: The authors propose to perform joint inference over the depth and effective width of a deep neural network. The depth is modelled as a stochastic process (in particular a Beta process), which allows for the depth to be infinite in theory. Given a fixed maximum width, the effective width is determined using a conjugate Bernoulli distribution which models the dropout probability for each neuron in each layer.