Improved Depth Estimation of Bayesian Neural Networks

Authors: bart  van erp
bert  de vries
Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-08-30 09:16:30

Abstract: This paper proposes improvements over earlier work by Nazareth and Blei (2022) for estimating the depth of Bayesian neural networks. Here, we propose a discrete truncated normal distribution over the network depth to independently learn its mean and variance. Posterior distributions are inferred by minimizing the variational free energy, which balances the model complexity and accuracy. Our method improves test accuracy on the spiral data set and reduces the variance in posterior depth estimates.

Main paper: 
 introduction: 
 Determining the optimal neural network architecture for a given problem is a challenging task, typically involving manual design iterations or automated grid searches. Both approaches are time-consuming and resource-intensive. A critical aspect of this process is balancing the model's complexity to prevent overfitting while ensuring high accuracy.The seminal work of Nazareth and Blei [1] introduced a variational inference scheme to network depth estimation. By treating the layer depth of the model as a latent variable, they can infer its posterior distribution. Importantly, their variational free energy provided an excellent objective for balancing the model complexity against the model accuracy.Although the approach presented in [1] offers a refreshing perspective, some areas could be improved. For instance, using a truncated Poisson distribution for layer depth results in the mean and variance being approximately equal, which can lead to significant uncertainty in determining the appropriate number of layers, especially for networks of increasing depth and complexity. Moreover, although the methodology in [1] is based on variational principles, certain simplifying assumptions undermine the probabilistic nature of their model. Specifically, the first-order linearization approximation over expectations neglects uncertainties over the parameters. This paper focuses exclusively on Bayesian neural networks and builds on the work by [1] , addressing the aforementioned areas of improvement. Specifically, we make the following contributions:• We propose a discrete truncated normal distribution over the number of hidden layers of a Bayesian neural network, enabling variance reduction in the posterior estimates of the appropriate number of layers;• Parameter estimation and structure learning are jointly performed by minimization of the variational free energy, explicitly taking the uncertainties over variables into account.In Section 2 the probabilistic model is specified, after which the inference procedure is elaborated in Section 3. Section 4 discusses the results obtained, and Section 5 concludes the paper.f 0 f 1 f 2 f 3 g 0 g 1 g 2 x n Ω 0 (x n ) Ω 1 (x n ) Ω 2 (x n )Figure 1 : Visualization of the non-linearity Ω L in (2) . Deeper models reuse parts of shallower models. 
model specification: 
 Let D = {(x n , y n )} N n=1 be a dataset of N labeled observations. We define the likelihood function of a Bayesian neural network asEQUATIONfor regression and classification, respectively. N (• | µ, Σ) represents a normal distribution with mean µ and covariance Σ and Cat(• | p) is a categorical distribution with event probabilities p, with σ(•) denoting the softmax function. The underlying non-linearity Ω L is parameterized by parameters θ, is visualized in Figure 1 and is defined as the compositionEQUATIONwith input transformation f 0 , latent transformations {f l } L l=1 and output transformations {g l } L l=0 . We treat the model depth L ∈ N 0 as an unknown variable. Therefore, a suitable discrete prior must be selected, with limited support and enabling efficient inference. The truncated Poisson distribution proposed in [1] has a variance and support that grows in network depth, preventing it from converging to a single value for the depth.Alternative discrete distributions suffer from similar problems, such as the negative binomial distribution whose variance is always larger or equal to its mean. Others do not have continuous parameters, such as the hypergeometric distribution with integer parameters. For the categorical distributions used in [2] , the support needs to be bounded. The generalized Poisson distribution [3] enables situations where its mean exceeds its variance, however, in those situations the distribution quickly becomes ill-defined [4] . Furthermore, the Conway-Maxwell-Poisson distribution does not require closed-form expressions for its normalization constant [5, 6] .Here, we propose to use a discrete truncated normal distribution, whose mean and variance are decoupled, which enables us to model both over-and under-dispersed distributions. LetN ≥0 (x | µ, σ 2 ) △ ∝ N (x | µ, σ 2 )1[x | x ≥ 0]denote a normal distribution truncated to the positive real line. Based on this truncated normal distribution, we define the prior over L as its discrete counterpartEQUATIONWe intentionally do not choose a discrete Gamma distribution [7] here, despite its positive domain, because computing derivatives to the shape parameter after truncation is difficult due to the presence of the lower incomplete gamma function in its cumulative density function.To complete the model specification, the prior over the parameters is chosen to fully factorize asEQUATIONwhere an explicit distinction in made between the parameters in the input and hidden layers {θ f l } L l=0 , which are shared amongst different model depths, and in the depth-specific output layers {θ g l } L l=0 .-1 With the model specified, the next step involves specifying the variational posterior distribution. We factorize the variational posterior distribution as0 1 -1 0 1 ω = 0 -1 0 1 ω = 2 -1 0 1 ω = 5 -1 0 1 ω = 10 -1 0 1 ω = 20EQUATIONTo retain tractability, we further truncate the variational posterior distribution over L to its lower and upper quantiles defined by p l , p u to ensure a limited support by definingEQUATIONUsing this expression the variational posterior distribution over the network depth is formulated asEQUATIONwhere the • accent identifies the variational parameters in ( 6) and ( 7). 
probabilistic inference: 
 Estimation of the variational posterior distributions, which encompasses both parameter estimation and structure learning, is achieved by minimization of the variational free energyEQUATIONwhere the expectation over parameters can be further decomposed asEQUATIONAlthough the expectation over the network depth seems computationally involved, the limited support as a result of the truncation in (7) reduces this operation to a finite summation asE q(L) [f (• | L)] = l∈supp{q(L)} q(l)f (• | l).Furthermore, since hidden layers are reused in networks of varying depth as illustrated in Figure 1 , most computations can be reused in computing the expected log-evidence. 
experiments: 
 All experiments1 have been implemented in Julia [8] to explore its excellent metaprogramming capabilities as required by the dynamic nature of the unbounded models. We closely follow the experimental design of [1] and generate a train, validation and test set of 1024 samples each of the spiral dataset [1, 9] for binary classification as described in Appendix A.1. This dataset is parameterized by a rotation speed ω, which captures the difficulty of the dataset as shown in Figure 2 .The input layer f 0 : R 2 → R 32 and latent layers f l : R 32 → R 32 ∀ l ≥ 1 each consist of a linear transformation followed by a LeakyReLU [10] . The output layers only involve a linear transformation g l : R 32 → R 2 ∀ l ≥ 0, where the non-linearity appears in (1b). We compare our approach to [1] which uses a Poisson(0.5) prior, where the variational posterior distribution is initialized by the Poisson(1.0) distribution, truncated to the 0.95-quantile. We select a similarly shaped normal distribution (µ L = 0, μL = 0, σ L = 1.15 and σL = 1.8), whose truncation is defined by p l = 0.025 and p u = 0.975. Appendix A.2 shows the resemblance between these priors.We jointly learn the parameters of the probabilistic model and its variational posterior through stochastic variational inference [11] by minimizing the variational free energy in (8) using the Adam optimizer [12] until convergence. Appendix A.3 specifies the hyperparameter settings. Inference in the model is performed using Bayes-by-backprop [13] with local reparameterization [14] . The model that achieves the lowest variational free energy on the validation set is saved and evaluated on the test set by forming predictions according toEQUATIONFigure 3 shows the achieved predictive accuracy on the test set and the inferred posterior distributions over the model depth. From this we conclude that the discrete truncated normal distribution outperforms the Poisson distribution on the spiral classification task. The normal-based model achieves a higher accuracy, which becomes increasingly significant when the complexity of the data increases. Furthermore, as expected, the posterior distribution over the model depth in the normal-based model has a reduced variance in comparison to the Poisson-based model, as its mean and variance are naturally decoupled during training. In practice this leads to computational savings when making predictions using (10) as the narrow support of q(L) requires less output layers g l to be active. 
discussion and conclusion: 
 This paper introduces a discrete truncated normal distribution for modeling the depth of a Bayesian neural network and demonstrates how to infer its posterior distribution through the minimization of variational free energy. Compared to methods using a Poisson prior [1] , our approach results in reduced variance in posterior estimates and improved test accuracy on the spiral classification task.The results presented in this paper show promising improvements in estimating the depth of Bayesian neural networks. However, additional experiments are required involving more complex models and tasks. Network width estimation and parameter pruning [13, [15] [16] [17] offer valuable opportunities for further expanding the methodology presented here.  
 
 Appendix: This appendix outlines the implementation details corresponding to experiments in Section 4.The spiral dataset used in the experiments of Section 4 and visualized in Figure 2 are generated according to the following sampling procedure2 :EQUATIONxEQUATIONFor selecting the prior and initial variational posterior distributions in the experiments of Section 4, we manually align the discrete truncated distribution with the Poisson distributions. Figure 4 shows a comparison of the log probability mass function of both functions as comparisons. Most important are the segments with a high log-probability, where the priors align relatively well from visual inspection.It should be noted that some discrepancies are inevitable, but at the same time negligible as these distributions only serve as a starting point and can be optimized over.Below we describe the training procedure. Here we tried to stay as close to the experimental design of [1] as possible.For each run we set a random seed equal to the run index. We then independently sample a train, validation and test set consisting of 1024 samples each for ω = 0, 1, . . . , 30 according to Appendix A. (l | 0, 1.8 2 ) dl.• Optimizer: Adam [12] with default hyperparameters (β 1 = 0.9 and β 2 = 0.999).• Learning rate: 0.005 (0.0005 for μL and σ2 L , and for the rate parameter of the posterior Poisson distribution λL ).• Number of epochs: 20.000.• Batch size: 256 (randomly shuffled per epoch).• Leaky ReLU: max(αx, x), where α = 0.1.• Reparameterization: strictly positive parameters are transformed using the softplus function for unconstrained optimization.Workshop on Bayesian Decision-making and Uncertainty, 38th Conference on Neural Information Processing Systems (NeurIPS 2024).All experiments are publicly available at https://github.com/biaslab/DepthEstimationBNN.The variance in the last step seems to differ from the original description in [1, Appendix B.1], however, the value reported there (0.02) refers to the standard deviation of the normal distribution, as verified with their publicly available experiments. 



Reviews: review: title: title: The paper proposes a novel approach for estimating the depth of BNNs, which improves upon existing methods.
main_review: main_review: The paper presents a novel approach for estimating the depth of Bayesian Neural Networks (BNN) by using a discrete truncated normal distribution instead of a Poisson distribution. The proposed method reduces the variance in posterior estimates of the appropriate number of layers and improves test accuracy compared to existing methods. The paper also discusses the probabilistic model and variational inference procedure, which are essential for understanding the method's effectiveness.

However, the paper has few weaknesses. First, the authors do not provide a thorough analysis of the proposed method's computational efficiency. Second, the authors only evaluated the method on a single dataset (spiral dataset), which may not be sufficient to demonstrate its generalizability. Third, the authors do not provide a detailed discussion of the method's limitations.
title: title: Overall good for workshop.
main_review: main_review: Quality:

The paper presents a novel approach to depth estimation in Bayesian neural networks by introducing a discrete truncated normal distribution. This method allows for independent learning of the mean and variance of the network depth, which is a significant improvement over previous methods .
The authors utilize variational inference to infer posterior distributions, effectively balancing model complexity and accuracy through the minimization of variational free energy. This indicates a solid methodological foundation.


Clarity:

The writing is generally clear, with a structured approach to presenting the problem, methodology, and results. The introduction effectively sets the stage by referencing earlier work and highlighting the improvements made.
However, some areas could benefit from further elaboration, particularly in explaining the implications of the results and the significance of the variational free energy in practical applications.

Originality:

The introduction of a discrete truncated normal distribution for modeling depth is a unique contribution that distinguishes this work from existing literature. This originality is crucial in advancing the field of Bayesian neural networks.
The paper builds on the foundational work of Nazareth and Blei, but it extends their ideas in a meaningful way, showcasing the authors' innovative thinking.

Significance:

The improvements in test accuracy on the spiral dataset and the reduction in variance of posterior depth estimates are significant outcomes that demonstrate the practical applicability of the proposed method.
The findings could have broader implications for various applications in machine learning where depth estimation is critical, thus enhancing the relevance of this research.

pros:

Innovative Approach: The use of a discrete truncated normal distribution is a novel contribution to the field.

Improved Accuracy: The method shows improved test accuracy and reduced variance in posterior estimates, which are critical for practical applications.

Solid Methodology: The use of variational inference and the Adam optimizer for training indicates a robust methodological framework.


cons:

Clarity Issues: Some sections could be clearer, particularly regarding the implications of the results and the methodology.

Limited Scope: The experiments are primarily focused on a specific dataset (spiral dataset), which may limit the generalizability of the findings.

In summary, this work presents a significant advancement in the field of Bayesian neural networks, with a clear methodology and promising results, although it could benefit from improved clarity and broader experimental validation.