Being Bayesian about Categorical Probability

We propose a Bayesian approach to construct the target distribution for classification, called a belief matching framework (BM; Figure 1(b)), in which the categorical probability about a label is regarded as a random variable ${\mathbf{z}}$. Specifically, we express the likelihood of ${\mathbf{z}}$ (given ${\mathbf{x}}$) about the label y as a categorical distribution $p_{{\textnormal{y}}|{\mathbf{x}},{\mathbf{z}}}=\mathcal{P}^{C}({\mathbf{z}}|{\mathbf{x}})$¹¹1In this paper, $p_{{\mathbf{x}}}=\mathcal{P}(\theta)$ is read as a random variable ${\mathbf{x}}$ follows a probability distribution $\mathcal{P}$ with parameter $\theta$.. Then, specification of the prior distribution over ${\mathbf{z}}|{\mathbf{x}}$ automatically determines the target distribution by means of the Bayesian inference: $p_{{\mathbf{z}}|{\mathbf{x}},{\textnormal{y}}}({\bm{z}})\propto p_{{\textnormal{y}}|{\mathbf{z}},{\mathbf{x}}}(y)p_{{\mathbf{z}}|{\mathbf{x}}}({\bm{z}})$.

We consider a conjugate prior for simplicity, i.e., the Dirichlet distribution. A random variable ${\mathbf{z}}$ (given ${\mathbf{x}}$) following the Dirichlet distribution with concentration parameter vector $\beta$, denoted by $\mathcal{P}^{D}(\beta)$, has the following density:

where $\Gamma(\cdot)$ is the gamma function, $\sum_{i}z_{i}=1$ meaning that ${\bm{z}}$ belongs to the $K-1$ simplex $\triangle^{K-1}$, $\beta_{i}>0,\>\forall i$, and $\beta_{0}=\sum_{i}\beta_{i}$. Here, we have that the mean of ${\mathbf{z}}|{\mathbf{x}}$ is $\beta/\beta_{0}$ and $\beta_{0}$ controls the sharpness of the density such that more mass centered around the mean as $\beta_{0}$ becomes larger.

By the characteristics of the conjugate family, we have the following posterior distribution given $\mathcal{D}$:

where the target posterior mean is explicitly smoothed by the prior belief, and the smoothing operation is performed by the principled way of applying Bayes’ rule. Specifically, the posterior mean is given by $\frac{1}{\beta_{0}+\sum_{i}c_{i}^{\mathcal{D}}({\mathbf{x}})}(\beta+c^{\mathcal{D}}({\mathbf{x}}))$, in which the prior distribution acts as adding pseudo counts. We note that the relative strength between the prior belief and the empirical count information becomes adaptive with respect to each data point.

Now, we specify the approximate posterior distribution modeled by the NNs, which aims to approximate $p_{{\mathbf{z}}|{\mathbf{x}},{\textnormal{y}}}$. In this paper, we model the approximate posterior as the Dirichlet distribution. To this end, we use an exponential function $g(x)=\exp(x)$ to transform logits to the concentration parameter of $\mathcal{P}^{D}$, and we let $\alpha^{{\bm{W}}}=\exp\circ f^{{\bm{W}}}$. Then, the NN represents the density over $\triangle^{K-1}$ as follows:

where $\alpha_{0}^{{\bm{W}}}({\mathbf{x}})=\sum_{i}\alpha^{{\bm{W}}}_{i}({\mathbf{x}})$.

From equation 7, we can see that outputs under BM encode much more information compared to those under the softmax. Specifically, it can be easily shown that the approximate posterior mean corresponds to the softmax. In this regard, BM enables neural networks to represent more rich information in their outputs, i.e., the density over $\triangle^{K-1}$ itself not just a single point on it such as the mean. This capability allows capturing more diverse characteristics of predictions at different locations, such as how much concentrate its density around the center mass point, which can be extremely helpful in many applications. For instance, BM gives a more sophisticated measure of the difference between predictions of two neural networks, which can benefit the consistency-based loss for semi-supervised learning as we will show in section 5.4. Besides, BM represents a more sophisticated measure of predictive uncertainty based on the density over simplex, such as mutual information.

From the perspective of learning the target distribution, BM can be considered as a generalization of softmax in terms of changing the moment matching problem to the distribution matching problem in $\mathcal{P}(\triangle^{K-1})$. To understand the distribution matching objective in BM, we reformulate equation 7 as follows:

In the limit of $q_{{\mathbf{z}}|{\mathbf{x}}}^{{\bm{W}}}\rightarrow p_{{\mathbf{z}}|{\mathbf{x}},{\textnormal{y}}}$, mean of the target posterior (equation 6) becomes a virtual label, for which individual ${\bm{z}}$ ought to match; the penalty for ambiguous configuration ${\bm{z}}$ is determined by the number of observations. Therefore, the distribution matching in BM can be thought of as learning to score a categorical probability based on closeness to the target posterior mean, in which exploitation of the closeness information is automatically controlled by the data.

We have defined the target distribution $p_{{\mathbf{z}}|{\mathbf{x}},{\textnormal{y}}}$ and the approximate distribution modeled by the neural network $q_{{\mathbf{z}}|{\mathbf{x}}}^{{\bm{W}}}$. We now present a solution to the distribution matching problem with maximizing the evidence lower bound (ELBO), defined by $l_{EB}({\textnormal{y}},\alpha^{{\bm{W}}}({\bm{x}}))=\mathbb{E}_{q_{{\mathbf{z}}|{\mathbf{x}}}^{{\bm{W}}}}[\log p({\textnormal{y}}|{\mathbf{x}},{\mathbf{z}})]-KL(q^{{\bm{W}}}_{{\mathbf{z}}|{\mathbf{x}}}\parallel p_{{\mathbf{z}}|{\mathbf{x}}})$. Using the ELBO can be motivated by the following equality (Jordan et al., 1999):

where we can see that maximizing $l_{EB}({\textnormal{y}},\alpha^{{\bm{W}}}({\mathbf{x}}))$ corresponds to minimizing $KL(q_{{\mathbf{z}}|{\mathbf{x}}}^{{\bm{W}}}\parallel p_{{\mathbf{z}}|{\mathbf{x}},{\textnormal{y}}})$, i.e., matching the approximate distribution to the target distribution, because the KL-divergence is non-negative and $\log p({\textnormal{y}}|{\mathbf{x}})$ is a constant with respect to ${\bm{W}}$. Here, each term in the ELBO can be analytically computed by:

where $\psi(\cdot)$ is the digamma function (the logarithmic derivative of $\Gamma(\cdot)$), and

where $p_{{\mathbf{z}}|{\mathbf{x}}}$ is assumed to be an input independent conjugate prior for simplicity; that is, $p_{{\mathbf{z}}|{\mathbf{x}}}=\mathcal{P}^{D}(\beta)$. With this analytical solution, we maximizes the ELBO with mini-batch approximation, which gives the following loss function: $\mathcal{L}({\bm{W}})=\mathbb{E}_{{\mathbf{x}},{\textnormal{y}}}[l_{EB}({\textnormal{y}},\alpha^{{\bm{W}}}({\mathbf{x}}))]\approx\frac{1}{m}\sum_{i=1}^{m}l_{EB}(y^{(i)},\alpha^{{\bm{W}}}({\bm{x}}^{(i)}))$. We note that computations of the ELBO and its gradient have a complexity of $\mathcal{O}(K)$ per sample, which is equal to those of softmax. This means that BM can preserve the scalability and the efficiency of the existing baseline. We also note that the analytical solution of the ELBO under BM allows to implement the distribution matching loss as a plug-and-play loss function applied to the logit directly.

The success of the Bayesian approach largely depends on how we specify the prior distribution due to its impact on the resulting posterior distribution. For example, the target posterior mean in equation 6 becomes the counting estimator as $\beta_{0}\rightarrow 0$. On the contrary, as $\beta_{0}$ becomes higher, the effect of empirical counting information is weakened, and eventually disappeared in the limit of $\beta_{0}\rightarrow\infty$. Therefore, considering that most of the inputs are unique in $\mathcal{D}$, choosing small $\beta_{0}$ is appropriate for prevents the resulting posterior distribution from being dominated by the prior²²2In an ideal fully Bayesian treatment, $\beta$ can be modeled hierarchically, and we left this as future research..

However, a prior distribution with small $\beta_{0}$ implicitly makes $\alpha_{0}^{{\bm{W}}}({\mathbf{x}})$ small, which poses significant challenges on the gradient-based optimization. This is because the gradient of the ELBO is notoriously large in the small-value regimes of $\alpha_{0}^{{\bm{W}}}({\mathbf{x}})$, e.g., $\psi^{\prime}(0.01)>10000$. In addition, our various building blocks including normalization (Ioffe & Szegedy, 2015), initialization (He et al., 2015), and architecture (He et al., 2016a) are implicitly or explicitly designed to make $\mathbb{E}[f^{{\bm{W}}}({\mathbf{x}})]\approx\textbf{0}$; that is, $\mathbb{E}[\alpha^{{\bm{W}}}({\mathbf{x}})]\approx\textbf{ 1}$. Therefore, making $\alpha_{0}^{{\bm{W}}}({\mathbf{x}})$ small can be wasteful or requires huge modifications to the existing building blocks. Also, $\mathbb{E}[\alpha^{{\bm{W}}}({\mathbf{x}})]\approx\textbf{1}$ is encouraged in a sense of natural gradient (Amari, 1998), which improves the conditioning of Fisher information matrix (Schraudolph, 1998; LeCun et al., 1998; Raiko et al., 2012; Wiesler et al., 2014).

In order to resolve the gradient-based optimization challenge in learning the posterior distribution while preventing dominance of the prior distribution, we set $\beta=\textbf{1}$ for the prior distribution and then multiply $\lambda$ to the KL divergence term in the ELBO: $l^{\lambda}_{EB}({\textnormal{y}},\alpha^{{\bm{W}}}({\mathbf{x}}))=\mathbb{E}_{q_{{\mathbf{z}}|{\mathbf{x}}}}\left[\log p({\textnormal{y}}|{\mathbf{x}},{\mathbf{z}})\right]-\lambda KL(q^{{\bm{W}}}_{{\mathbf{z}}|{\mathbf{x}}}\parallel\mathcal{P}^{D}(\textbf{1}))$. This trick significantly stabilizes the optimization process, while making a local optimal point remains unchanged. To see this, we can compare the gradients of the ELBO and the lambda multiplied ELBO:

where $\tilde{\alpha}_{k}^{{\bm{W}}}({\mathbf{x}})=\lambda\alpha_{k}^{{\bm{W}}}({\mathbf{x}})$. Here, we can see that a local optimal in equation 12 is achieved when $\alpha^{{\bm{W}}}({\mathbf{x}})=\beta+\tilde{{\mathbf{y}}}$ and a local optima for equation 13 is $\alpha^{{\bm{W}}}({\mathbf{x}})=1+\frac{1}{\lambda}\tilde{{\mathbf{y}}}$. Therefore, a ratio between $\alpha_{i}^{{\bm{W}}}({\mathbf{x}})$ and $\alpha_{j}^{{\bm{W}}}({\mathbf{x}})$ equal to those of a local optimal point in equation 12 for every pair of $i$ and $j$. In this regard, searching for $\lambda$ with $l^{\lambda}_{EB}({\textnormal{y}},\alpha^{{\bm{W}}}({\mathbf{x}}))$ and then multiplying $\lambda$ after training corresponds to the process of searching for the prior distribution’s parameter $\beta$ with $\mathcal{L}({\bm{W}})$.

We evaluate the generalization performance of BM on CIFAR (Krizhevsky, 2009) with the pre-activation ResNet (He et al., 2016b). CIFAR-10 and CIFAR-100 contain 50K training and 10K test images, and each 32x32x3-sized image belongs to one of 10 categories in CIFAR-10 and one of 100 categories in CIFAR-100. Table 1 lists the classification error rates of the softmax cross-entropy loss, BM, and MC-dropout. In all configurations, BM consistently achieves the best generalization performance on both datasets. On the other hand, last-layer MC dropout sometimes results in higher generalization errors than softmax and all-layer MC-dropout significantly increases error rates even though they consume 100x more computations for inference.

We next perform a large-scale experiment using ResNext-50 32x4d and ResNext-101 32x8d (Xie et al., 2017) on ImageNet (Russakovsky et al., 2015). ImageNet contains approximately 1.3M training samples and 50K validation samples, and each sample is resized to 224x224x3 and belongs to one of the 1K categories; that is, the ImageNet has more categories, a larger image size, and more training samples compared to CIFAR, which may enable a more precise evaluation of methods. Consistent with the results on CIFAR, BM improves test errors of softmax (Table 2). This result is appealing because improving the generalization error of deep NNs on large-scale datasets by adopting a Bayesian principle without computational overhead has rarely been reported in the literature.

One of the most attractive benefits of Bayesian methods is their ability to represent the uncertainty about their predictions. In a naive sense, uncertainty representation ability is the ability to “know what it doesn’t know.” For instance, models having a good uncertainty representation ability would increase some form of predictive uncertainty on misclassified examples compared to those on correctly classified examples. This ability is extremely useful in both real-world applications and downstream tasks in machine learning. For example, underconfident NNs can produce many false alarms, which makes humans ignore the predictions of NNs; conversely, overconfident NNs can exclude humans from the decision-making loop, which results in catastrophic accidents. Also, better uncertainty representation enables balancing exploitation and exploration in reinforcement learning (Gal & Ghahramani, 2016) and detecting OOD samples (Malinin & Gales, 2018).

We evaluate the uncertainty representation ability on both in-distribution and OOD datasets. Specifically, we measure the calibration performance on in-distribution test samples, which examines a model’s ability to match its probabilistic output associated with an event to the actual long-term frequency of the event (Dawid, 1982). The notion of calibration in NNs is associated with how well its confidence matches the actual accuracy; e.g., we expect the average accuracy of a group of predictions having the confidence around $0.7$ to be close to $70\%$. We also examine the predictive uncertainty for OOD samples. Since the examples belong to none of the classes seen during training, we expect neural networks to produce outputs of “I don’t know.”

BM adopts the Bayesian principle outside the NNs, so it can be applied to models already trained on different tasks, unlike BNNs. In this regard, we examine the effectiveness of BM on the transfer learning scenario. Specifically, we downloaded the ImageNet-pretrained ResNet-50, and fine-tune weights of the last linear layer for 100 epochs by the Adam optimizer (Kingma & Ba, 2015) with learning rate of 3e-4 on three different datasets (CIFAR-10, Food-101 (Bossard et al., 2014), and Cars (Krause et al., 2013).

Table 3 compares test error rates of softmax and BM, in which BM consistently achieves better performances compared to softmax. Next, we examine the predictive uncertainty for OOD samples (Figure 6). Surprisingly, we observe that BM significantly improves the uncertainty representation ability of pretrained-models by only fine-tuning the last layer weights. These results present a possibility of adopting BM as a post-hoc solution to enhance the uncertainty representation ability of pretrained models without sacrificing their generalization performance. We believe that the interaction between BM and the pretrained models is significantly attractive, considering recent efforts of the deep learning community to construct general baseline models trained on extremely large-scale datasets and then transfer the baselines to multiple down-stream tasks (e.g., BERT (Devlin et al., 2018) and MoCo (He et al., 2019)).

BM enables NNs to represent rich information in their predictions (cf. section 3.2). We exploit this characteristic to benefit consistency-based loss functions for semi-supervised learning. The idea of consistency-based losses employs information of unlabelled samples to determine where to promote robustness of predictions under stochastic perturbations (Belkin et al., 2006; Oliver et al., 2018). In this section, we investigate two baselines that consider stochastic perturbations on inputs (VAT; Miyato et al., 2018) and networks ($\Pi$-model; Laine & Aila. 2017), respectively. Specifically, VAT generates adversarial direction ${\bm{r}}$, then measures KL-divergence between predictions at ${\bm{x}}$ and ${\bm{x}}+{\bm{r}}$:

where the adversarial direction is chosen by ${\bm{r}}=\operatorname*{arg\,max}_{\parallel{\bm{r}}\parallel\leq\epsilon}KL(\phi(f^{{\bm{W}}}({\bm{x}}))\parallel\phi(f^{{\bm{W}}}({\bm{x}}+{\bm{r}})))$, and $\Pi$-model measures $L^{2}$ distance between predictions with and without enabling stochastic parts in NNs:

where $\bar{f}$ is a prediction without the stochastic parts.

We can see that both methods achieve the perturbation invariant predictions by minimizing the divergence between two categorical probabilities under some perturbations. In this regard, BM can provide a more delicate measure of the prediction consistency–divergence between Dirichlet distributions–that can capture richer probabilistic structures, e.g., (co)variances of categorical probabilities. This generalization the moment matching problem to the distribution matching problem can be achieved by replacing only the consistency measures in equation 15 with $KL(q_{{\mathbf{z}}|{\bm{x}}}^{{\bm{W}}}\parallel q_{{\mathbf{z}}|{\bm{x}}+{\bm{r}}}^{{\bm{W}}})$ and in equation 16 with $KL(\bar{q}_{{\mathbf{z}}|{\bm{x}}}^{{\bm{W}}}\parallel q_{{\mathbf{z}}|{\bm{x}}}^{{\bm{W}}})$.

We train wide ResNet 28-2 (Zagoruyko & Komodakis, 2016) via the consistency-based loss functions on CIFAR-10 with 4K/41K/5K number of labeled training/unlabeled training/validation samples. Our experimental results show that the distribution matching metric of BM is more effective than the moment matching metric under the softmax on reducing the error rates (Table 4). We think that the improvement in semi-supervised learning with more sophisticated consistency measures shows the potential usefulness of BM on other useful applications. This is because employing the prediction difference of neural networks is a prevalent method in many domains such as knowledge distillation (Hinton et al., 2015) and model interpretation (Zintgraf et al., 2017).