PAC Bayesian Performance Guarantees for Deep (Stochastic) Networks in Medical Imaging

Abstract

Application of deep neural networks to medical imaging tasks has in some sense become commonplace. Still, a “thorn in the side” of the deep learning movement is the argument that deep networks are prone to overfitting and are thus unable to generalize well when datasets are small (as is common in medical imaging tasks). One way to bolster confidence is to provide mathematical guarantees, or bounds, on network performance after training which explicitly quantify the possibility of overfitting. In this work, we explore recent advances using the PAC-Bayesian framework to provide bounds on generalization error for large (stochastic) networks. While previous efforts focus on classification in larger natural image datasets (e.g., MNIST and CIFAR-10), we apply these techniques to both classification and segmentation in a smaller medical imagining dataset: the ISIC 2018 challenge set. We observe the resultant bounds are competitive compared to a simpler baseline, while also being more explainable and alleviating the need for holdout sets.

1. Introduction

Understanding the generalization of learning algorithms is a classical problem. Practically speaking, verifying whether a fixed method of inference will generalize may not seem to be a challenging task. Holdout sets are the tool of choice for most practitioners – when sample sizes are large, we can be confident the measured performance is representative. In medical imaging, however, sample sizes are often small and stakes are often high. Thus, mathematical guarantees¹ on the performance of our algorithms are of paramount importance. Yet, it is not abundantly common to provide such guarantees in medical imaging research on deep neural networks; we are interested in supplementing this shortage.

A simple (but effective) guarantee on performance is achieved by applying a Hoeffding Bound to the error of an inference algorithm reported on a holdout set. In classification tasks, Langford provides a useful tutorial on these types of high probability bounds among others [32]. The bounds are easily extended to any bounded performance metrics in general, and we use this methodology as a baseline in our own experimentation (Sect. 3). While effective, Hoeffding’s Inequality falls short in two regards: (1) use of a holdout set requires that the model does not see all available data and (2) the practitioner gains no insight into why the model generalized well. Clearly, both short-comings can be undesirable in a medical imaging context: (1) access to the entire dataset can be especially useful for datasets with rare presence of a disease and (2) understanding why can both improve algorithm design and give confidence when deploying models in a clinical setting. Thus, we desire practically applicable bounds – i.e., competitive with Hoeffding’s Inequality – which avoid the aforementioned caveats.

Unfortunately, for deep neural networks, practically applicable guarantees of this nature can be challenging to produce. Traditional PAC bounds based on the Vapnik-Chervonenkis (VC) dimension [8,46,48,49] accomplish our goals to some extent, but require (much) more samples than parameters in our network to produce a good guarantee.² When our networks are large – e.g., a ResNet-18 [23] with more than 10M parameters – our datasets are thus required to be unreasonably sized to ensure generalization, especially, in medical imaging contexts. In effect, these bounds are vacuous; they are logically meaningless for the sample sizes we observe in practice. Specifically, vacuous describes any bound on error which is larger than 1, and therefore, gives us no new insight on a model’s generalization ability because error rates lie in the range [0, 1]. The term was coined by Dziugaite & Roy [16] who observed that most PAC bounds on error for deep networks were in fact vacuous (when computed). The authors demonstrate a technique to compute non-vacuous bounds for some deep stochastic networks trained on standard datasets (e.g., MNIST) using the PAC-Bayesian framework.³ As discussed later, the motivation for using PAC-Bayesian bounds is the hypothesis that stochastic gradient descent (SGD) trained networks generalize well when their solutions lie in large, flat minima.⁴ Thus, these PAC-Bayes bounds give us exactly what we desire. They provide practically applicable guarantees on the performance of deep networks, while allowing the model access to all data and also giving insight on what network properties lead to good generalization in practice. To borrow an analogy from Arora [2], this approach is “prescriptive”. It is similar to a doctor’s orders to resolve high blood pressure: cut the salt, or in our case, look for flat minima.

In this context, the end goal of this paper is to validate whether the “flat minima prescription” – with observed evidence in traditional computer vision – may find similar success in explaining generalization for the small-data, non-traditional tasks that are common to medical imaging. Our contributions, in this sense, are primarily experimental. We demonstrate non-vacuous PAC-Bayesian performance guarantees for deep stochastic networks applied to the classification and segmentation tasks within the ISIC 2018 Challenge [13]. Importantly, our results show that PAC-Bayesian bounds are competitive against Hoeffding’s Inequality, offering a practical alternative which avoids the aforementioned caveats. We employ much the same strategies used by Dziugaite et al. [15,16] as well those used by Pérez-Ortiz et al. [43]. With that said, our different setting yields novel experimental results and poses some novel challenges. Specifically, in segmentation, we compute non-vacuous bounds for a fully-sized U-Net using a medical imaging dataset with about 2.3K training samples. To our knowledge, for deep stochastic networks, we are the first to compute non-vacuous bounds in segmentation on such small datasets. Along the way, we offer some practical insights for the medical imaging practitioner including a (mathematically sound) trick to handle batch normalization layers in PAC-Bayes bounds and an experimental technique to “probe” parameter space to learn about the generalization properties of a particular model and dataset. We hope these results promote continued research on the important topic of guarantees in medical imaging.

2. PAC-Bayesian Theory and Generalization

2.1. Formal Setup

In the PAC-Bayes setting, we consider a hypothesis space $\mathcal{H}$ and a distribution Q over this space. Specific to our context, $\mathcal{H}={ℝ}^d$ will represent the space of deep networks with some fixed architecture, and Q will be a distribution over $\mathcal{H}$. Typically, we will set $Q=\mathcal{N}(\mu ,\Sigma )$, a multivariate normal. For some fixed space $\mathcal{Z}=\mathcal{X}\times \mathcal{Y}$, the hypothesis $h\in \mathcal{H}$ defines a mapping $\mathit{x}\mathrm{\mapsto }\mathit{h}\mathrm{(}\mathit{x}\mathrm{)}$ with $x\in \mathcal{X}$, and $h(x)\in \mathcal{Y}$. Given a [0, 1]-bounded loss function $\ell :\mathcal{H}\times \mathcal{Z}\to [0,1]$ and a distribution D over $\mathcal{Z}$, the risk of h is defined ${\mathcal{R}}_{\ell }(h,D)={\mathbf{E}}_{(x,y)~D}\ell (h,(x,y))$. Given instead a sample S ~ D^m over $\mathcal{Z}$, the empirical risk is denoted ${\widehat{\mathcal{R}}}_{\ell }(h,S)$ and is computed as usual by averaging. In all of our discussions, the data distribution D or sample S is usually easily inferred from context. Therefore, we typically write ${\mathcal{R}}_{\ell }(h)={\mathcal{R}}_{\ell }(h,D)$ and ${\widehat{\mathcal{R}}}_{\ell }(h)={\widehat{\mathcal{R}}}_{\ell }(h,S)$. With these definitions, we are interested in quantifying the risk of a stochastic model.⁵ In the context of neural networks, one can imagine sampling the distribution Q over ${ℝ}^d$ and setting the weights before performing inference on some data-point x. Often, we will refer to Q itself as the stochastic predictor. The associated risk for such a model is defined as ${\mathcal{R}}_{\ell }(Q)={\mathbf{E}}_{h~Q}{\mathcal{R}}_{\ell }(h)$ with ${\widehat{\mathcal{R}}}_{\ell }(Q)$ similarly defined. Typically, we cannot exactly compute expectations over Q. For this reason, we also define ${\widehat{R}}_{\ell }(\widehat{Q})$ for a sample $\widehat{Q}~Q^n$ as an empirical variant, computed by averaging. The last components for any PAC-Bayesian bound come from two notions of the Kullback-Leibler (KL) divergence between two distributions Q and P written KL(Q||P) and defined as usual. For numbers q,p ∑ [0,1], we write kl(q||p) as shorthand for the KL divergence between Bernoulli distributions parameterized by q, p. This is typically used to quantify the difference between the risk on a sample and the true risk. In the next section, we put the discussed pieces in play.

2.2. The PAC-Bayesian Theorem

The PAC-Bayesian theory begins primarily with the work of McAllester [37] with similar conceptualizations given by Shawe-Taylor & Williamson [47]. Besides what is discussed in this section, for completeness, readers are also directed to the work of Catoni [11], McAllester [36], Germain et al. [19,20], and the primer by Guedj [21]. We start by stating the main PAC-Bayesian Theorem as given by Maurer [35]. See also Langford & Seeger [33] for the case where l is a 01 loss.

Theorem 1 (Maurer).

Let l be a [0, 1]-bounded loss function, D be a distribution over $\mathcal{Z}$, and P be a probability distribution over $\mathcal{H}$. Then, for δ ∑ (0,1)

$$\underset{S~D^m}{\mathbf{Pr}}\left(\forall Q:\text{kl}\left({\widehat{\mathcal{R}}}_{\mathcal{l}}(Q)\Vert {\mathcal{R}}_{\mathcal{l}}(Q)\right)\le \frac{\text{KL}(Q\Vert P)+\ln \frac{1}{\delta }+\ln \sqrt{4m}}{m}\right)\ge 1-\delta .$$ (1)

By way of Pinsker’s inequality, the above may be loosened for the purpose of interpretation [19]

$${\mathcal{R}}_{\ell }(Q)\le {\widehat{\mathcal{R}}}_{\ell }(Q)+\sqrt{\frac{\mathrm{KL}(Q\Vert P)+\ln \frac{1}{\delta }+\ln \sqrt{4m}}{2m}}.$$ (2)

In Sect. 3, we compute a much tighter formulation of the bound given in Eq. (2) which handles the term kl $\left({\widehat{\mathcal{R}}}_{\ell }(Q)\Vert {\mathcal{R}}_{\ell }(Q)\right)$ directly. We provide the derivation of this bound in the Appendix. Various insights and results used to build the final bound are of course due to Langford & Caruana [31]; Dziugaite et al. [15,16]; and Pérez-Ortiz et al. [43] who have all computed similar (or identical) bounds on stochastic neural networks before us. In Sect. 3, we compute this bound for classification and segmentation tasks. For classification, we take to be the 01-loss defined l₀₁(h, (x, y)) = 1[h(x) = y] where 1 is the indicator function. Precisely, ${\mathcal{R}}_{{\ell }_{01}}(h)$ is equal to 1 minus the accuracy. For segmentation, we pick to be l_DSC(h, (x, y)) = 1 - DSC(h, (x, y)) where DSC is the [0, 1]-valued Dice similarity coefficient. These upperbounds trivially yield corresponding lowerbounds for both the accuracy and the Dice similarity coefficient, respectively.

Selecting the Prior.

Often P is referred to as the prior and Q as the posterior. Still, it is important to note that no restriction on the form of P and Q is required (e.g., as it is in Bayesian parameter estimation). What is required is that P be fixed before observing the sample S that is used to compute the bound. Albeit, P can depend on D and samples independent of S. In fact, it is not uncommon for the prior P to be data-dependent.⁶ That is, P may be trained on a sample which is disjoint from that which is used to compute the bound; i.e., disjoint from S. On the other hand, the bound holds for all posteriors Q regardless of how Q is selected. So, the datasets used to train Q and P may actually intersect. All in all, we can train Q with all available data without violating assumptions. We must only ensure the datasets used to train P and compute the bound do not intersect. In effect, we avoid the first caveat of Hoeffding’s Inequality.

Interpretation.

The bound also offers insight into why the model generalizes. Intuitively, we quantify the complexity of the stochastic predictor Q in so much as it deviates from some prior knowledge we have on the solution space (e.g., from the data-dependent prior). This is captured in the term KL(Q||P). Dziugaite & Roy [16] also relate PAC-Bayesian bounds to the flat-minima hypothesis. To understand their observation, we consider the case where Q is a normal distribution $\mathcal{N}\left(\mu ,{\sigma }^{2}I\right)$ with σ a constant and I the identity matrix. The model Q on which we bound the error is stochastic: each time we do inference, we sample from $\mathcal{N}\left(\mu ,{\sigma }^{2}I\right)$. Because the distribution has some variance (dictated by σ), we sample network weights in a region around the mean. Thus, when performance of the stochastic model Q is good, there must be a non-negligible area around the mean where most networks perform well, i.e., a flat minimum around the mean. We know the variance is non-negligible in the posterior network Q because a small upperbound implies small KL divergence with the prior P which itself has non-negligible variance (we pick this value). So, to reiterate, small KL divergence and small empirical risk imply a flat-minimum of appropriate size around the mean of Q. In this sense the bound is explainable: Q generalizes well because it does not deviate much from prior knowledge and it lies in a flat minimum.

Additional Context.

Dziugaite & Roy [16] provide a nice synopsis of the history behind the flat-minima hypothesis including the work of Hochreiter & Schmidhuber [25], Hinton & Van Camp [24], Baldassi et al. [3,4], Chaudhari et al. [12], and Keskar et al. [27]. Since then, large scale empirical studies – e.g., Jiang et al. [26]; Dziugaite, Drouin, et al. [14] – have continued to indicate that measures of sharpness of the minimum may be good indicators of neural network generalization in practice. For completeness, we also point out some other theoretically plausible indicators of deep network generalization. These include small weight norms – e.g., Bartlett [5,6]; Neyshabur et al. [40,41] – and the notion of algorithmic stability proposed by Bousquet & Elisseeff [10] which focuses instead on the SGD algorithm – e.g., Hardt et al. [22]; Mou et al. [39]; Kuzborskij & Lampert [29].

3. Experiments

In this section, we first evaluate PAC-Bayesian bounds within a self-bounded learning setting. Specifically, a self-bounded learner must both learn and provide a guarantee using the same dataset.⁷ As noted, providing guarantees with our trained networks can bolster confidence in small data regimes. We compare the PAC-Bayesian bounds discussed in Sect. 2 to a simple baseline for producing performance guarantees: application of Hoeffding’s Inequality to a holdout set.⁸ We show PAC-Bayesian bounds are competitive with Hoeffding’s Inequality, while also alleviating some caveats discussed in the previous sections. This result (in medical imaging) compliments those previously shown on natural image datasets by Pérez-Ortiz et al. [43]. Specifically, we compute bounds on the Lesion Segmentation and the Lesion Classification Tasks in the ISIC 2018 Skin Lesion Challenge Dataset [13]. The data in these sets used for training (and bound computation) totals 2.3K and 9K labeled examples, respectively, which is much smaller than previous computation [15,16,43] using MNIST [34] or CIFAR-10 [28].⁹ Our second contribution in this section comes from tricks and tools which we hope prove useful for the medical imaging practitioner. We demonstrate an experiment to probe the loss landscape using PAC-Bayesian bounds and also devise a strategy to handle batch normalization layers when computing PAC-Bayesian bounds. Our code is available at: https://github.com/anthonysicilia/PAC-Bayes-In-Medical-Imaging.

3.1. Setup

Models.

For segmentation, we use U-Net (UN) [45] and a light-weight version of U-Net (LW) with 3% of the parameters and no skip-connections. For classification, we use ResNet-18 (RN) [23]. Probabilistic models use the same architecture but define a multivariate normal distribution over network weights $\mathcal{N}(\mu ,\Sigma )$ with a diagonal covariance matrix Σ. The distribution is sampled to do inference.

Losses.

For segmentation, we train using the common Dice Loss [38] which is a smooth surrogate for 1 - DSC. For classification, we use the negative log-likelihood. Probabilistic models used modified losses we describe next.

Training Probabilistic Models.

Recall, in the PAC-Bayes setting we define both the prior P and the posterior Q as both are needed to compute bounds and Q is needed for inference. The prior P is a probabilistic network defined by $\mathcal{N}\left(\mu ,{\sigma }^{2}I\right)$ where I is the identity matrix and σ_p is a constant. In this text, we use a data-dependent prior unless otherwise noted (see Sect. 2). To pick the prior, μ_p is learned by performing traditional optimization on a dataset disjoint from that which is used to compute the bound (see Data Splits). The parameter σ_p = 0.01 unless otherwise noted. The posterior Q is initialized identically to P before it is trained using PAC-Bayes with Backprop (PBB) as proposed by Rivasplata et al. [43,44]. This training technique may be viewed as (mechanically) similar to Bayes-by-Backprop (BBB) [9]. In particular, it uses a re-parameterization trick to optimize a probabilistic network through SGD. Where PBB and BBB differ is the motivation, and subsequently, the use of PAC-Bayes upperbounds as the objective to optimize.¹⁰ Note, PAC-Bayes bounds are valid for all [0, 1]-bounded losses, and thus, are valid for the Dice Loss or normalized negative log-likelihood. The upperbound used for our PBB objective is the Variational bound of Dziugaite et al. [15].

Probabilistic Batch Norm Layers.

While generally each weight in the probabilistic networks we consider is sampled independently according to a normal distribution, batch norm layers¹¹ must be handled with special care. We treat the parameters of batch norm layers as point mass distributions. Specifically, the parameter value has probability 1 and 0 is assigned to all other values. The posterior distribution for these parameters is made identical to the prior by “freezing” the posterior batch norm layers during training and inference. In effect, we avoid sampling the means and variances in our batch norm layers, and importantly, the batch norm layers do not contribute to the KL-divergence computation. In the Appendix, we provide a derivation to show this strategy is (mathematically) correct; it relies primarily on the independence of the weight distributions.

Optimization Parameters.

Optimization is done using SGD with momentum set to 0.95. For classification, the batch size is 64, and the initial learning rate is 0.5. For segmentation, the batch size is 8, and the initial learning rate is 0.1 for LW and 0.01 for U-Net. All models are initialized at the same random location and are trained for 120 epochs with the learning rate decayed by a factor of 10 every 30 epochs. In the PAC-Bayesian setting, the data-dependent prior mean μ_p is randomly initialized (as other models) and trained for 30 epochs. The posterior Q is initialized at the prior P and trained for the remaining 90 epochs. The learning rate decay schedule is not reset for posterior training.

Bound Details.

Note, in all cases, PAC-Bayes bounds are computed using a data-dependent prior. Bounds are computed with 95% confidence (δ = 0.05) with data sample size given in the next section. For PAC-Bayes bounds, the number of models sampled is either 1000 (in Fig. 1a) or 100 (in Fig. 1b, c, d).

Fig. 1.

(a) DSC/ACC (red) and lowerbounds (blue). (b, c, d) Modulation of prior variance for U-Net (b) and LW (c) and ResNet-18 (d).

Data Splits for Self-bounded Learning.

Each method is given access to a base training set (90% of the data) and is expected to both learn a model and provide a performance guarantee for this model when applied to unseen data. To evaluate both the model and performance guarantee in a more traditional fashion, each method is also tested on a final holdout set (10% of the data) which no model sees. Splits are random but identically fixed for all models. For probabilistic networks trained using PBB, we split the base training data into a 50%-prefix set¹² used to train the prior and a disjoint 50%-bound set. Both the prefix-set and bound-set are used to train the posterior, but recall, the PAC-Bayes bound can only be computed on the bound-set (Sect. 2). For the baseline non-probabilistic networks, we instead train the model using a traditional training set (composed of 90% of the base training set) and then compute a guarantee on the model performance (i.e., a Hoeffding bound) using an independent holdout set (the remaining 10% of the base training set). In this sense, all models are on an equal footing with respect to the task of a self-bounded learner. All models have access only to the base training set for both training and computation of valid performance guarantees. In relation to Fig. 1, performance metrics such as DSC are computed using the final holdout set. Lowerbound computation and training is done using the base training set.

3.2. Results

Comparison to Hoeffding's Inequality.

As promised, we observe in Fig. 1a performance guarantees by both bounds are comparable and performance on the final holdout set is also comparable. Hoeffding’s Inequality does have a slight advantage with regards to these metrics, but as mentioned, PAC-Bayes bounds possess some desirable qualitative properties which make it an appealing alternative. For PAC-Bayes bounds the posterior Q sees all training data, while for the Hoeffding Bound, one must maintain an unseen holdout set to compute a guarantee. Further, we may explain the generalization of the PBB trained model through our interpretation of the PAC-Bayes bound (see Sect. 2). These qualities make PAC-Bayes appealing in medical imaging contexts where explainability is a priority and we often need to maximize the utility of the training data.

Flat Minima and Their Size.

As discussed, the application of PAC-Bayesian bounds may be motivated by the flat minima hypothesis (see Sect. 2). We explore this idea in Fig. 1b, c, d by modulating the prior variance σ_p across runs. Informally, these plots can give us insight into our loss landscape. The reasonably tight lowerbounds – which are slightly looser than in Fig. 1a only due to fewer model samples – imply small KL-Divergence and indicate the prior and posterior variances are of a similar magnitude. Likewise, the difference between the prior and posterior means should not be too large, relative to the variance. A fixed random-seed ensures priors are identical, so each data-point within a plot should correspond to roughly the same location in parameter space; i.e., we will assume we are analyzing the location of a single minimum.¹³ For U-Net, we see stable performance and a sudden drop as the prior variance grows. Before the drop at σ_p = 0.04, consistently high DSC and a high lowerbound indicate the network solution lies in a flat minimum (as discussed in Sect. 2). So, we may conclude a flat minimum proportional in size to σ_p = 0.03. For LW and ResNet-18, we instead see consistent performance degradation as the prior variance grows. For these networks, the minima may not be as flat. Informally, such sensitivity analysis can tell us “how flat” the minima are for a particular network and dataset as well as “how large”. Practically, information like this can be useful to practitioners interested in understanding the generalization ability of their models. Namely, it is hypothesized “larger” flat minima lead to better generalization because less precision (fewer bits) is required to specify the weights [25].

4. Conclusion

As a whole, our results show how PAC-Bayes bounds can be practically applied in medical imaging contexts – where theoretical guarantees (for deep networks) would appear useful but not commonly discussed. With this said, we hope for this paper to act primarily as a conversation starter. At 2.3K examples, the segmentation dataset we consider is still larger than commonly available in some application domains (e.g., neuroimaging applications). It remains to be considered how effective these bounds can be in ultra-low resource settings.