Mean-Field Approximation to Gaussian-Softmax Integral with Application to Uncertainty Estimation

The main steps of the approximation scheme Eq.(11) in this subsection also appeared in (Daunizeau, 2017), though the author there did not use it or other forms Eq.(9,10) for uncertainty estimation. Addiditionally, we discuss the intuitions of those approximations (§3.2), extend them with two temperatures for additional control (§4), and showed that the simpler form of Eq.(9) is equally successful (§6).

The case $K=2$ is well-known (Bishop, 2006, p. 219)

where the softmax becomes the sigmoid function $\sigma(\cdot)$. $\lambda_{0}$ is a constant and is usually chosen to be $\pi/8$ or $3/\pi^{2}$. To generalize to $K>2$, we rewrite the integral Eq. (29) so that we compare $a_{i}$ to $a_{k}$ with $k\neq i$, drawing close to Eq. (7),

where “mean field” (MF) “pushes” the integral through each sigmoid term independently²²2This approximation is prompted by the MF approximation $\mathbb{E}\left[f(\cdot)\right]\approx f(\mathbb{E}\left[\cdot\right])$ for nonlinear functions $f(\cdot)$. Moreover, we have factorized the $f()$ to be integrated as “pairwise factors”, i.e. $(a_{k}-a_{i})$, and apply $p(a_{i},a_{k})$ to them – this is the structural mean-field approach by factorizing joint distributions into “smaller” factors (Koller and Friedman, 2009).. Next we plug in the approximation Eq. (7) to compute the inside expectations.

Any of the 3 forms has approximation errors that depend on how “nonlinear” the integrand is. If one class is dominant over others, the softmax is similar to the sigmoid, where the approximation is well-known and applied in practice. We leave more precise quantification of those errors to future studies and focus in this work on their empirical utility. In particular, the simple forms of mfGS make it possible to understand them intuitively.

First, note that when $\bm{S}\rightarrow\mathbf{0}$, i.e., $p(\bm{a})$ reduces to a $\delta$ distribution, mfGS approximations become exact

Namely, if there is no uncertainty associated with $\bm{a}$, the model should output the regular softmax — this is a desirable sanity check of the approximation schemes. However, when $\bm{a}$ has variance, each scheme handles it differently.

mf0 is especially interesting. The $\tilde{e}_{k}(0)$ is similar to the regular softmax with a temperature

However, as opposed to the regular temperature scaling factor (Guo et al., 2017), $T_{k}$ is adaptive as it depends on the variance of a particular sample to be predicted. If the activation on the $k$-th class has high variance, the temperature for that category is high, reducing the corresponding “probability” $\tilde{e}_{k}$. Specifically,

In other words, the scaling factor is category and input-dependent, providing additional flexibility to a global temperature scaling factor.

mf1 and mf2 differ from mf0 in how much information from $i\neq k$ is considered: the variance $s_{i}^{2}$ or additionally the covariance $s_{ik}$. The effect is to measure how much the $k$-th class should be confused with the $i$-th class after taking into consideration the variances in those classes. Note that mf0 and mf1 are computationally preferred over mf2 which uses $K^{2}$ covariances, where $K$ is the number of classes.

We propose to use a Gaussian for $\tilde{p}(\theta)$ in Eq. (1):

where $\widehat{\theta}_{n}$ is the neural network parameter of the solution (Eq. (2)). ${T_{\mathsf{ens}}}$ is a global ensemble temperature controlling the spread of the Gaussian by scaling the covariance matrix. The covariance $\bm{\Sigma}_{n}$ is

where the Hessian $H_{n}$ and the variance of the gradients $J_{n}$ are defined as follows, computed at $\widehat{\theta}_{n}$,

The Gaussian is used not only because it is convenient but also the forms of $\bm{\Sigma}_{n}$ have been occurring frequently in the literature as ways to capture the variations in estimations of model parameters. In §4.3, we will discuss the theoretical rationales behind the choices, after we describe how to use them for uncertainty estimation. (As a preview to our results, $\bm{\Sigma}_{n}=H_{n}^{-1}$ performs the best empirically.)

To compute Eq. (1), we make one final approximation, borrowing the linear perturbation idea from the influence function (Cook and Weisberg, 1982), or infinitessimal jackknife, see Eq. (27) in §4.3. Consider the logits $\bm{a}(\theta)$ (Eq. (4)) when $\theta$ is close to $\widehat{\theta}_{n}$,

where $\nabla\bm{a}$ is the gradient w.r.t. $\theta$. Under $\theta\sim\tilde{p}(\theta)$, we have $\bm{a}(\theta)$ as a Gaussian as well,

Applying any form of mfGS in §3 of computing Gaussian-softmax integral Eq. (29), we obtain a mean-field approximation to Eq. (1). Note that the approximation in Eq. (19) is exact if $\bm{a}(\theta)$ is a linear function of the parameter, for instance, in the last-layer setting where all layers but the last are treated as being deterministic.

Readers who are familiar with classical results in theoretical statistics and Bayesian statistics can easily identify the origins of the forms in Eq. (17). An accessible introduction is provided in (Geyer, 2013). Despite being widely known, they have not been applied extensively to deep learning uncertainty estimation. We summarize them briefly as follows.

In the asymptotic normality theory of maximum likelihood estimation, for identifiable models, if the model is correctly specified and the true parameter is $\theta^{*}$, the maximum likelihood estimation converges in distribution to a Gaussian

where $I_{1}(\theta^{*})$ is the Fisher information matrix, and equals to the variance of the gradients $J_{1}(\theta^{*})$³³3The notation $V(\cdot)$ is used in (Geyer, 2013). The subscript 1 refers to sample size 1. $I_{n}=nI_{1}$ and $J_{n}=nJ_{1}$.. In practice, the observed Fisher $J_{1}(\theta^{*})$ is often used, and only requires the first-order derivatives:

If the model is misspecified, the estimator converges to the so-called “sandwich” estimator

The true model parameter $\theta^{*}$ is unknown and so we use a plug-in estimation such as $\widehat{\theta}_{n}$, resulting Eq. (16).

Neural networks are not identifiable so the asymptotic normality theory does not apply. However, the following infinitessimal jackknife technique can also be used to motivate the use of those Gaussians. Jackknife is a well-known resampling method to estimate the confidence interval of an estimator (Tukey, 1958; Efron and Stein, 1981). Each element $z_{i}$ is left out from the dataset $\mathcal{D}$ to form a unique “leave-one-out” Jackknife sample $\mathcal{D}_{i}=\mathcal{D}-\{z_{i}\}$, giving rise to

We obtain an ensemble of $n$ such samples $\{\tilde{\theta}_{i}\}_{i=1}^{n}$ and use them to estimate the variances of $\widehat{\theta}_{n}$ and the predictions made from $\widehat{\theta}_{n}$. However, it is not feasible to retrain modern neural networks $n$ times. Infinitesimal jackknife approximates the retraining (Jaeckel, 1972; Koh and Liang, 2017; Giordano et al., 2018). The basic idea is to approximate $\tilde{\theta}_{i}$ with a small perturbation to $\widehat{\theta}_{n}$, i.e., infinitessimal jackknife

where $H_{n}$ is the Hessian Eq. (18). We fit the $n$ copies of $\tilde{\theta}_{i}$ with a Gaussian distribution

which is very similar to the asymptotic normality results. The perturbation idea can also be used to approximate any differentiable function $f(\tilde{\theta}_{i})$, e.g. the logits in Eq. (19).

If we extend the above procedure and analysis to bootstrapping (i.e., sampling with replacement), we will arrive at the following Gaussian to characterize infinitessimal bootstraps:

More details can be found in  (Giordano et al., 2018). Finally, the Laplace approximation in Bayesian neural networks uses $\theta\sim\mathcal{N}(\theta^{*},H^{-1})$ to approximate the posterior distribution where $\theta^{*}$ is the maximum-a-posterior estimate.

In short, all those Gaussian approximations to the true distribution of $\widehat{\theta}_{n}$ (i.e. the epistemic uncertainty) share the similar form as Eq. (17) with an added temperature ${T_{\mathsf{ens}}}$ controlling the scaling factor (such as $\nicefrac{{1}}{{n}}$).

We provide key information in the following. For more details, please see the Suppl. Material.

NLL measures the KL-divergence between the empirical label distribution and the predictive distribution of the classifiers. ECE measures the discrepancy between the histogram of the predicted probabilities by the classifiers and the observed ones in the data – properly calibrated classifiers should yield matching histograms. Both metrics are commonly used in the literature and the lower the better.

On the task of out-of-distribution (OOD) detection, we assess how well $\tilde{p}(x)$, the classifier’s output being interpreted as probability, can be used to distinguish invalid samples from normal in-domain ones. Following the common practice (Hendrycks and Gimpel, 2016; Liang et al., 2017; Lee et al., 2018), we report two threshold-independent metrics: area under the receiver operating characteristic curve (AUROC), and area under the precision-recall curve (AUPR). Since the precision-recall curve is sensitive to the choice of positive class, we report both “AUPR in:out” where in-distribution and out-of-distribution samples are swapped. We also report detection accuracy, the optimal accuracy achieved among all thresholds in classifying in-/out-domain samples. All three metrics are the higher the better.

Classifiers For MNIST, we train a two-layer MLP with 256 ReLU units per layer, using Adam optimizer for 100 epochs. For CIFAR-10, we train a ResNet-20 with Adam for 200 epochs. On CIFAR-100, we train a DenseNet-BC-121 with SGD optimizer for 300 epochs. For ImageNet, we train a ResNet-50 with SGD optimizer for 90 epochs.

We validate mfGS by comparing it to Monte Carlo sampling for computing the integral Eq. (29). Since there is no ground-truth to this integral, we validate its utility in estimating predictive uncertainty as a “downstream task”. To this end, we use the MNIST and NotMNIST datasets and perform uncertainty estimation with only the last layer.

In Table 1, results from all mean-field approximation schemes are better than those from sampling, in both in-domain and out-domain metrics. Note that while sampling could be improved by increasing samples (thus, more computational costs and memory demands), the mean-field approaches still manage to outperform, especially on the OOD metrics, probably because sampling requires a significantly more samples to converge (the observed improvement is at a fairly slow rate, preventing us from trying more samples.).

Among the 3 schemes, mf0 performs better in ECE, an in-domain task metric while being identical to others. This might be an interesting observation worthy of further investigation. One possible explanation is that for OOD samples, the variances of the logits are such that $s_{i}\approx s_{k}\gg s_{ik}$, i.e. $\bm{S}=g(x;\phi_{n})H_{n}^{-1}g(x;\phi_{n})^{\top}$⁴⁴4Eq. (20) after proper vectorization and padding. is diagonally dominant – the features of an OOD $x$ is in the null space of $H$. This is intuitive as $H$ should contain vectors maximizing inner products with features from the in-domain samples.

Among 3 choices of selecting $\bm{\Sigma}_{n}$ in Eq. (17), $H^{-1}$ has a clear advantage over the other two $J^{-1}$ and $H^{-1}JH^{-1}$, though $J^{-1}$ is reasonably good. We also experimented with the covariance matrix derived from the samples on the training trajectories, i.e. SWAG. SWAG performs similar as mf0 with $H^{-1}JH^{-1}$. This is re-assuring as it confirms the theoretical analysis that SWAG converges to the sandwich estimator when its trajectory is in the vicinity of the MLE (Maddox et al., 2019).

In short, the study suggests the mfGS approximation has high quality, measured in the metrics of the “downstream” task of uncertainty estimation. Thus, in the rest of the paper, we focus on mf0 using $H^{-1}$ as $\bm{\Sigma}_{n}$.

In this section, we study the effect of using different number of layers of parameters in mfGS for uncertainty estimation. We report both sampling and mf0 results in Table 2.

mf0 consistently outperforms Monte Carlo sampling. We also observe that using more layers than the last layer does not further help. This is in line with recent studies (Zeng et al., 2018; Kristiadi et al., 2020) showing that last-layer Bayesian approximation can yield good performances in uncertainty tasks. Thus we focus on mfGS with the last layer for the rest of uncertainty estimation experiments.

Ashukha et al. suggest to assess new uncertainty estimate methods by checking the number of models in Deep E. needed to have a matching performance (Ashukha et al., 2020). In OOD metrics, mf0 is about 500 models in Deep E., representing a significant reduction in computational cost and memory requirement — it took about 30 hours on a TitanX GPU to train Deep E. (500) while mf0 needs only one model and a few seconds of inference time. In in-domain metrics, mf0 needs to be improved in accuracy, which is left as future work.

Table 5 provides an extended version of Table 3 in the main text to include results from more approaches. mf0 outperforms all other approaches in OOD metrics including ensemble methods, single-pass methods, and BNN methods. For in-domain tasks, mf0 performs better than either the single-model or BNN methods in ECE. Ensemble performs the strongest in terms of accuracy and NLL.

Tables 6 and 7 supplement the main text with additional experimental results on the CIFAR-10 dataset with both in-domain and out-of-distribution detection tasks, and on the CIFAR-100 with out-of-distribution detection using the SVHN dataset. BNN(LL-VI) refers to BNN with stochastic variational inference applied on the last layer only, following Snoek et al. (2019).

The results support the same observations in the main text: mf0 noticeably outperforms other approaches on out-of-distribution detection, but is not as strong as ensemble method in terms of accuracy and NLL. It does in general outperform all other approaches in those two metrics.

We study the roles of ensemble and activation temperatures in mfGS. ${T_{\mathsf{ens}}}$ controls the how wide we want our Gaussian is and thus determines how many models we would want to integrate. On the other hand, ${T_{\mathsf{act}}}$ controls how much we want each model in the ensemble to be confident about their predictions.

To this end, on MNIST and NotMNITST datasets we grid search the two temperatures and generate the heatmaps of NLL and AU ROC on the held-out sets, shown in  Figure 1. Note that $({T_{\mathsf{ens}}}=\infty,{T_{\mathsf{act}}}=1)$ correspond to MLE. What is particularly interesting is that for NLL, higher activation temperature (${T_{\mathsf{act}}}\geq 1$) and lower ensemble temperature (${T_{\mathsf{ens}}}\leq 10^{-2}$) work the best. For area under ROC, however, lower temperatures on both work best. Using ${T_{\mathsf{act}}}>1$ for better calibration was also observed in (Guo et al., 2017). On the other end, for OOD detection,  (Liang et al., 2017) suggests a very high activation temperature (${T_{\mathsf{act}}}=1000$ in their work, likely due to using a single model instead of an ensemble).

Examining both plots horizontally, we can see when the ${T_{\mathsf{ens}}}$ is higher, the effect of ${T_{\mathsf{act}}}$ is more pronounced than when ${T_{\mathsf{ens}}}$ is lower. Vertically, we can see when the ${T_{\mathsf{act}}}$ is lower, the effect of ${T_{\mathsf{ens}}}$ is more pronounced than when ${T_{\mathsf{act}}}$ is higher, especially on the OOD tasks.

Snoek et al. (2019) points out that many uncertainty estimation methods are sensitive to distributional shift. Thus, we evaluate the robustness of mf0 on rotated MNIST images, from $0$ to $180^{\circ}$. The ECE curves in Figure 1(c) show mf0 is better or as robust as other approaches.

The statistics of the benchmark datasets used in our experiment are summarized in Table 9.

NLL is defined as the KL-divergence between the data distribution and the model’s predictive distribution,

where $y_{c}$ is the one-hot embedding of the label.

ECE measures the discrepancy between the predicted probabilities and the empirical accuracy of a classifier in terms of $\ell_{1}$ distance. It is computed as the expected difference between per bucket confidence and per bucket accuracy. All predictions $\{p(x_{i})\}_{i=1}^{N}$ are binned into $S$ buckets such that $B_{s}=\{i\in[N]|p(x_{i})\in I_{s}\}$ are predictions falling within the interval $I_{s}$. ECE is defined as,

where $\text{conf}(B_{s})=\frac{\sum_{i\in B_{s}}p(x_{i})}{|B_{s}|}$ and $\text{acc}(B_{s})=\frac{\sum_{i\in B_{s}}\bm{1}\ [\hat{y}_{i}=y_{i}]}{|B_{s}|}$.

Hyper-parameters in training  Table 10 provides key hyper-parameters used in training deep neural networks on different datasets.

Hyper-parameters in uncertainty estimation  For mf0 method, we use $\lambda_{0}=\frac{3}{\pi^{2}}$ in our implementation. For the Deep E. approach, we use $M=5$ models on all datasets as in (Snoek et al., 2019). For RUE, Dropout, BNN(VI) and BNN(KFAC), where sampling is applied at inference time, we use $M=500$ Monte-Carlo samples on MNIST, and $M=50$ on CIFAR-10, CIFAR-100 and ImageNet. We use $B=10$ buckets when computing ECE on MNIST, CIFAR-10 and CIFAR-100, and $B=15$ on ImageNet.

For in-domain uncertainty estimation, we use the NLL on the held-out sets to tune hyper-parameters. For the OOD detection task, we use area under ROC curve on the held-out to select hyper-parameters. We report the results of the best hyper-parameters on the test sets. The key hyper-parameters we tune are the temperatures, regularization or prior in BNN methods, and dropout rates in Dropout.

Other implementation details When Hessian needs to be inverted, we add a dampening term $\tilde{H}=(H+\epsilon I)$ following (Ritter et al., 2018; Schulam and Saria, 2019) to ensure positive semi-definiteness and the smallest eigenvalue of $\tilde{H}$ be 1. For BNN(VI), we use Flipout (Wen et al., 2018) to reduce gradient variances and follow (Snoek et al., 2019) for variational inference on deep ResNets. On ImageNet, we compute the Kronecker-product factorized Hessian matrix, rather than full due to high dimensionality. For BNN(KFAC) and RUE, we use mini-batch approximations on subsets of the training set to scale up on ImageNet, as suggested in (Ritter et al., 2018; Schulam and Saria, 2019).