Joint Inference for Neural Network Depth and Dropout Regularization
Year: 2021

Abstract:
Dropout regularization methods prune a neural network’s pre-determined backbone structure to avoid overfitting. However, a deep model still tends to be poorly calibrated with high confidence on incorrect predictions. We propose a unified Bayesian model selection method to jointly infer the most plausible network depth warranted by data, and perform dropout regularization simultaneously. In particular, to infer network depth we define a beta process over the number of hidden layers which allows it to go to infinity. Layer-wise activation probabilities induced by the beta process modulate neuron activation via binary vectors of a conjugate Bernoulli process. Experiments across domains show that by adapting network depth and dropout regularization to data, our method achieves superior performance comparing to state-of-the-art methods with well-calibrated uncertainty estimates. In continual learning, our method enables neural networks to dynamically evolve their depths to accommodate incrementally available data beyond their initial structures, and alleviate catastrophic forgetting.

heading: 1 Introduction
text: 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 cannot 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.

heading: 2 Related work
text: 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.

heading: 3 The joint inference framework
text: 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.

heading: 3.1 Network structure with infinite hidden layers
text: 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.

heading: 3.2 Beta process over layer number
text: 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)

heading: 3.3 Marginal likelihood over network structures
text: 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→∞.

heading: 3.4 Efficient inference with SSVI
text: We employ structured stochastic variational inference (SSVI) to approximate the marginal likelihood with a variational lower bound (ELBO) retaining the dependence between network structures and hyperparameters [18, 19, 30]. In particular, we define the variational distribution as
q(Z,ν|{ak}Kk=1, {bk}Kk=1) = q(ν)q(Z|ν) = K∏ k=1 Beta(νk|ak, bk) M∏ m=1 ConBer(zmk|πk) (6)
where πk = ∏k j=1 νj , and {ak, bk}Kk=1 are variational parameters of q(ν), andK denotes a truncation level, which can be relaxed as in [31]. We relax the constraint of the discrete variables with continuous ones by reparameterizing the Bernoulli distribution into a concrete Bernoulli distribution [32, 33]:
ConBer(zmk|πk) = τ πk(zmk)
−τ−1(1− πk)(1− zmk)−τ−1
(πk(zmk)−τ + (1− πk)(1− zmk)−τ )2 (7)
where τ is a temperature controlling the distribution smoothness. We generate random samples from the distribution by first sampling from a logistic distribution as the external source of randomness and then putting the samples through a logistic function as follow:
zmk = 1
1 + exp(−τ−1(log πk − log(1− πk) + )) ∼ Logistic(0, 1) (8)
This reparameterization of the Bernoulli distribution allows us to backpropagate the gradients of the lower bound with respect to the parameters while sampling from it.
We derive a lower bound of the log marginal likelihood in (5) as: log p(D|W, L, α, β) ≥ ∫ q(Z,ν) log p(D|Z,W)dZdν + ∫ q(Z,ν) log p(Z,ν)
q(Z,ν) dZdν (9)
The ELBO from the right-hand side of (9) becomes our objective function:
L(Z,ν,W) = Eq(Z,ν)[log p(D|Z,W)]− KL[q(Z|ν)||p(Z|ν)]− KL[q(ν)||p(ν)] (10) The first term of the ELBO in (10) is the averaged log-likelihood. The second and the third terms are the KL divergence between the variational distribution and the true prior, respectively. In particular, for the second term we relax the prior p(Z|ν) with a concrete Bernoulli as in (7) with different temperatures. For the third term the KL divergence between the variational distribution and the prior distribution over ν is:
KL[q(ν)||p(ν)] = ∑ k ln ( B (α, β) B(ak, bk) ) + (ak − α)ψ(ak)
+ (bk − β)ψ(bk) + (α− ak + β − bk)ψ(ak + bk) (11)
where ψ(·) denotes a di-gamma function.
Theorem 1 Assume Ŵ is the optimized value of the variable, and the likelihood and the induced prior on ΠZ for Z are specified in (2) and (4), in the large sample limit, the limiting of the ELBO in (10) becomes equivalent to the Bayesian information criterion [34], such that BIC(Z) = log p(D|Z,Ŵ)− |Z| 2 log N 2π , where |Z| denotes the number of activated neurons in the network structure.
The proof is in Appendix. This theorem indicates that optimizing the limiting case of our SSVI framework can therefore be seen as optimizing the popular model selection criteria. We adopt the implicit differentiation of beta distribution for backpropagation [35], which allows us to stochastically optimize L with respect to the variational parameters every time we sample a network structure. We iteratively maximize ELBO with respect to the variational parameters and weights W, and approximate the ELBO using random samples from the variational distribution.
We approximate the posterior predictive distribution by sampling from the variational posterior distribution for predictions. Thus, the predictive distribution for new data x∗ is
p(y∗|x∗,Ŵ, {âk}, {b̂k}) = ∫ p(y∗|x∗,Z,Ŵ)q(Z,ν|{âk}, {b̂k})dZdν (12)
where Ŵ is the MAP estimate of the neural network weights based on mini-batch gradient descent.

heading: 4 Experiments
text: We analyze the behavior of our framework by applying it to MLPs and CNNs on a variety of tasks. We study how it enables network depth to adapt to dataset sizes, and capture different types of uncertainty on a synthetic dataset. Our framework outperforms state-of-the-art methods on the UCI datasets with uncertainty estimation. To obtain performance comparison, we perform Bayesian Optimisation and Hyperband (BOHB) to determine the best configurations for each method [36]. We also investigate the effect of truncation level K and M , and show that their settings have no influence on the performance as long as they are reasonably large. The experiments with image datasets show that our framework improves CNN performances with more compact network structures and wellcalibrated predictive distributions. Lastly, we demonstrate our method in two continual learning tasks, and show that by allowing network depth to dynamically augment to accommodate incrementally available information we can effectively alleviate catastrophic forgetting. 3

heading: 4.1 Synthetic experiments
text: We analyze our method’s behaviors given different dataset sizes, and its uncertainty estimation. We incrementally generate 20, 500 and 2000 data points from a periodic function [14]:
y = sin (6x) + 0.4x2 − 0.1x3 − x cos (9 √ exp(x)) + (13)
where ∼ N (0, 0.12) with 20% for validation. The data are shown in Figure 2(a) bottom row. We set the maximum number of neurons per layer M = 20, and use leaky ReLU activation functions in the input and hidden layers with batch normalization to retain stability. We simulate 5 samples per minibatch to approximate the ELBO in (10), and evaluate convergence by computing the cross-correlations of the sample values over 3000 epochs.
As shown in Figure 2(a), to accommodate more information from larger datasets, the neural network becomes deeper with more activated neurons. Although we set the truncation level K = 100 for all the training cases, only the first few hidden layers have activated neurons (i.e., 1 layer, 3 layers, and 4 layers for the respective cases). Figure 2(a) bottom row shows how the predictive distributions in (12) are capturing the true function and estimating uncertainty as the dataset size increases. We compare our method’s performance with vanilla dropout [6] whose backbone structure sizes are the same as the respective ones inferred by our method. The predictive performances on 700 held-out data points in Figure 2(b) indicate that by obtaining the critical balance between network depth and neuron activations our method achieves more stable performance with less variance on the small dataset, and significantly better performance on the larger ones.
3Implementation details of all the experiments are in Appendix.
The datasets generated with a known variance allow us to assess the different types of uncertainty of our method’s predictions. In particular, the epistemic uncertainty is obtained by drawing multiple samples of the inferred depth and neuron activation, and evaluating them on the test set of 700 data points from the same data distribution. Figure 2(c) shows that the epistemic uncertainty decreases as the amount of data points increases. The aleatoric uncertainty approaching the true uncertainty (0.1) shows an increasingly improved estimate as more data is given. Our method’s predictive uncertainty obtained by combining epistemic and aleatoric uncertainties converges to a constant value.

heading: 4.2 Performance comparisons on UCI datasets
text: We evaluate all methods on UCI datasets [37] using standard splits [38], and report their performance in terms of log-likelihood (LL) and root mean square error (RMSE). In particular, we compare the performances with vanilla dropout (Dropout), concrete dropout (CD) [10], Indian buffet process dropout (IBPD) [11], Bayesian architecture learning (BAL) [14], stochastic depth (SDepth) [13], depth uncertainty networks (DUNs) [3], and Deep Ensembles (DeepEns) [39]. We perform hyperparameter optimization with Bayesian Optimisation and Hyperband (BOHB) to determine the best configurations including backbone-structure depths and hyperparameter settings for each method [36], as suggested in [3]. The detailed implementation is in Appendix.
In Figure 3, we rank the methods from 1 to 8 based on their mean performance across each dataset and metric, and report mean ranks with ± one standard deviation. Our joint inference framework outperforms other methods in terms of both test log-likelihood (LL) and RMSE by achieving the best rank across all datasets. In particular, LL measures both accuracy and uncertainty calibration [40]. The narrower standard deviations of our method indicate its robustness. The vanilla dropout achieves the second-best rank followed closely by the Deep Ensembles in terms of LL and concrete dropout in terms of RMSE.

heading: 4.3 Effect of truncation level
text: We further investigate the effect of truncation level K on our method’s performance, and compare it with the backbone-structure depth L of dropout variants and structure selection methods.
All the methods have the same maximum width M = 64 (i.e., the maximum number of feature maps in convolutional layers). In Figure 4, we apply all the methods on CNNs with the backbone-structure depth L and the truncation level K over the range L = K ∈ {5, 10, 15, 20, 25, 30, 35, 40, 45, 50} for classification tasks of four image datasets: MNIST [41], FashionMNIST [42], SVHN [43], and CIFAR10 [44]. When L = K = 5, all the methods have an underfitting problem due to the limited model capacity, although our method still achieves competitive or better performances. As the truncation level becomes reasonably large (≥ 10), it no longer affects our method’s performance. The results are consistent with Theorem 1, as our joint inference essentially works as Bayesian model selection on a network structure truncation. On the contrary, the depths of backbone structure L
affect the other methods to a great extent. The increase of the test errors shows that they all have an overfitting problem when the backbone-structure depths become large. Individual method reaches its own best performance, as L increases. Our method outperforms them for all these cases with orders of magnitude less computing resources across the four datasets. The respective evaluations in Figure 6 (i.e., M = 64) show that our method only activates neurons in 6, 6, 10 and 6 hidden layers for predictions on MNIST, FashionMNIST, SVHN, and CIFAR10, respectively. Only 10% of total neurons in the truncation level are activated for all the cases. This suggests by jointly inferring network depth and neuron activations our method is quite robust to overfitting, and the balance between the two also leads to superior performances with compact network structures. In contrast, the backbone-structure sizes of the dropout variants and structure selection methods have to be set carefully to determine the best configuration for a given dataset.
We also evaluate the uncertainty calibration of our framework and other methods with their best settings in Figure 4. In Figure 5(a), we train all methods on MNIST and evaluate their predictive distributions on increasingly rotated digits [45]. The uncertainty calibration of our method is more robust to the data with large angle rotations than other methods, and it achieves the performance with the lowest expected calibration error (ECE) [46]. In Figure 5(b), to make a fair and rigorous comparison as in [3], we also evaluate the corruption robustness of all methods on the CIFAR10 dataset with 16 types of algorithmically generated corruptions [47]. Each type of corruption has 5 levels of severity. The Deep Ensembles outperform other methods, and our model achieves comparable performance at all corruption levels.

heading: 4.4 Effect of M
text: We next assess the influence of the maximum number of neurons/feature maps (M ) on our method with classification tasks on the four image datasets. Figure 6 shows the evolution of the number of inferred convolutional layers as M increases. When M = 16, we tend to have deeper network structures to compensate for the relatively narrow layers. As M increases, the structures become
shallower. When M is reasonably large, it has no influence on the inferred structure depth. The percentage of activated neurons in the truncation level remains relatively stable across different values of M . This suggests that our method can adapt both the network depth and the neuron activations to maintain the best performances as M changes.

heading: 4.5 Continual learning
text: Continual learning is an important application since real-world tasks are dynamic and non-stationary. Machine learning models need to learn consecutive tasks without forgetting how to perform previously trained ones. Assuming the pre-determined backbone structures are sufficient to accommodate all information from the continual tasks, network regularization approaches alleviate catastrophic forgetting by regularizing the updates of neural weights [48]. Dropout and its variants are applied to learn an implicit gating mechanism that activates different gates for different tasks [49, 50]. However, the rigid backbone structure constrains the applicability of these methods in real-world settings.
We slightly modify our method to enable network depth and neuron activations to dynamically evolve to accommodate incrementally available data. Given a set of sequentially arriving datasets {Dt} where each may contain a single datum, we update our ELBO in (10) as:

heading: Lt = Eqt(Z,ν)[log p(Dt|Z,ν,W)]− KL[qt(Z|ν)||qt−1(Z|ν)]− KL[qt(ν)||qt−1(ν)] (14)
text: by replacing the priors p(Z|ν) and p(ν) with their variational approximation for the previous dataset qt−1(Z|ν) and qt−1(ν). We initialize them as q0(Z|ν) = p(Z|ν) and q0(ν) = p(ν). The variational distribution in (6) after seeing the t-th dataset is recursively updated by taking the distribution after seeing the (t− 1)-th dataset, multiplying by the likelihood and re-normalizing. We compare the performance of our method with baseline continual learning methods including elastic weight consolidation (EWC) [48], orthogonal gradient descent (OGD) [51], and stochastic gradient descent with dropout (SGD + dropout). We conduct an extensive grid search on the hyperparameter setting of these methods, and evaluate them on two popular continual learning benchmarks: permuted MNIST [48] and rotated MNIST datasets [52]. Permuted MNIST at each time step Dt consists of labeled MNIST images whose pixels undergo a fixed random permutation. Rotated MNIST is generated by the rotation of the original MNIST images. Figure 7(a) shows our method outperforms the other methods on the first four tasks, and is only slightly worse than some (i.e., SGD for Permuted MNIST). Over the five tasks, our method achieves the highest average accuracies (95.16% on Rotated MNIST and 92.34% on permuted MNIST). Although some methods perform well on the fifth task after being trained on it, they fail to preserve the knowledge learned in previous tasks. In contrast, our method can preserve the knowledge from old tasks and achieves comparable or better accuracy on new tasks. Figure 7(b) shows the evolution of network structures for the five sequential tasks, as our method activates more neurons in the inferred hidden layers. In particular, the first hidden layer tends to be full from the first task with its activation probability close to one, the activation probabilities of the second and the third layers go up as the tasks are coming.

heading: 5 Conclusions
text: The proposed general joint inference framework can be applied to various neural networks. The experiments on MLPs and CNNs show that our method balances network depth and neuron activations to achieve superior performance. By enabling neural network structures to dynamically evolve to accommodate incrementally available data, we effectively alleviate catastrophic forgetting.

heading: 6 Funding
text: This material is based upon work supported by the National Science Foundation under Award No. 2045804 and Award No. 1850492.

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

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.