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 paper proposes a comprehensive but effective method to simultaneously perform inference for the NN depth and width. As the idea might not be new, the execution of the method is solid. The methodology is also linked to Bayesian information criterion with a theorem.