Out of Distribution Detection, Generalization, and Robustness Triangle with Maximum Probability Theorem
^††thanks: This project is supported by the NIH funding: R01-CA246704 and R01-CA240639, and Florida Department of Health (FDOH): 20K04.

Regularization. Training machine learning models requires careful tuning of hyper parameters, architectures, and the loss functions. An effective tuning of the non-trainable aspects of models is mainly justified by empirical evaluations which costs time and is a source of uncertainty in evaluation of models and algorithms. Regularization of models is widely accepted as a way of stabilizing the training process and improving models generalization performance.

In the Bayesian view of machine learning, the popular regularization approach is to define an explicit prior over parameters of the model. Parallel to this view, a recent approach, called Maximum Probability Theorem (MPT), is introduced [33]. MPT considers a model as an event in the probability space and provides an upper-bound for the probability of the event representing the model. Increasing the probability of a model in MPT leads to regularization effects. Unlike other regularization strategies, MPT eliminates the need for explicitly defined priors over the model’s parameters. Instead, MPT requires only a prior on the observables (e.g. input and output) to determine the prior on the parameters or hidden random variables.

Our proposal. We hypothesize that MPT can act as a black-box regularization in vision applications such as OOD detection, and can lead to improved robustness, generalization, and performance in such tasks. We first construct some objective functions encapsulating MPT regularization, and demonstrate their relation with cross entropy loss. Next, we show how to incorporate MPT regularization into training of CNNs and their energy based variants. Finally, we discuss the details of OOD experiments and the results on generalization and robustness.

Why Energy Based Models (EBMs)? The recent line of work by [46] showed that treating CNNs as EBMs increase the generalization performance of the CNN model and provide additional abilities for the network [11] such as:

• •

  EBMs can be used to detect anomalous input (Out of Distribution Detection),

• •

  EBMs can be trained without labels,

• •

  EBMs can be used as generative models.

Hence, in this study, we treat CNNs as an EBM [26, 42] and also experiment with effects of MPT on CNN-EBMs.

The summary of our contributions is below:

1. 1.

   We successfully adapt MPT regularization in the energy based models,

2. 2.

   We reconsider the OOD detection problem with the proposed MPT based regularization,

3. 3.

   We demonstrate the robustness and generalization effects of MPT regularization on three datasets and 1080 trained networks, with obtained promising results.

Our work relies mainly on three lines of research, regularization methods, EBMs, and OOD methods.
Regularization and Priors. Regularization and priors of probabilistic models have a rich history, including Jaynes’ Maximum Entropy Principle [21, 19, 20], Jeffreys’ uninformative priors [22], and reference priors [4, 6] generalizing Jeffreys priors. Practically, the well accepted approach to regularize probabilistic models is to impose explicit priors on parameters or other variables in the model. In deep learning, the explicit prior is reflected in the popular regularization schemes. For example, The $\ell_{1},\ell_{2}$ norm regularization reflects a Gaussian and Laplacian prior on variables [24, 31, 9]. The logic behind reparameterization invariant priors in [4, 22], is that our prior knowledge should not change based on how a parameter is represented or how the model is constructed. The prior on parameters in reference priors, are determined to maximize the mutual information of parameters and the observable random variables [2, 3, 4, 5]. Authors in [34] propose a gradient based optimization of models with reference priors, which is a promising ground for integrating uninformative priors. However, automatic determination of the reference priors for parameters of complex models is not fully understood [34].

MPT [33], a recent addition to the probabilistic models literature, offers a black box regularization of probabilistic models while being invariant to reparameterizations. To use MPT, one does not need to explicitly determine the priors over the parameters. In other words, the parameters do not need to be modeled as random variables.

Energy Based Models. EBMs assume some non-normalized log probability distribution over the space of their input [41, 26]. This non-normalized log probability is referred to as energy. In the conventional loss functions such as Log-Likelihood, the loss depends on the imposed log probability distribution. Training EBMs is possible by estimating the gradients and sampling from their imposed probability distribution [41, 14, 10]. The gradients of the log probability with respect to the model’s parameters, is in the form of an expectation with respect to the EBMs distribution. If the domain of EBMs are arbitrary large, unbiased and low variance sampling from their underlying distribution is not trivial. To sample from the EBM probability distribution multiple sampling techniques can be used, including the class of MCMC sampling techniques [36, 35]. The major drawback with the MCMC techniques is that they induce bias in calculation of gradients during their burnout period [32, 38].

Also, depending on the domain of EBMs, the estimated gradients can have large variance [12]. The variance in estimation of gradients is a double edge sword. On one hand, the variance helps with escaping plateaus in the objective function landscape. On the other hand, large variance in gradients, disrupts convergence of parameters to the optimal point. Despite the bias and variance drawbacks, MCMC methods such as Stochastic Langevin Dynamics [43] had been shown to be successful in training EBMs [46].

Out of Distribution Detection. Out of distribution (OOD) detection task is motivated by understanding reliability of prediction of a model on some input data [1]. Within the existing methods for OOD [16, 31, 29, 39] and many others, our work focuses on probabilistic OOD methods and EBMs. We hypothesize that if the model fits a probability distribution to its input, the model’s probability distribution function (pdf) on input can be used to identify out-distribution data. Existing OOD detection methods commonly rely on a scoring function that derives statistics from the output layer of the neural network. Popular probabilistic measures for OOD are Softmax score [13], Energy score [11]. ODIN [28], GODIN[16] can be viewed as modification of Energy score and Softmax Score. MOOD [29] is an architectural modification using similar probabilistic principles for training and evaluation.

Notation. We use the notation according to the MPT paper: a probability space is defined as $(\Omega,\Sigma,P)$ where $\Omega$ is the sample space, $\Sigma$ is the $\sigma$-algebra and $P$ is the probability measure. A parametric model, parameterized by $\theta$ in some arbitrary space, is defined as $M_{\theta}\in\Sigma$, an event in the $\sigma$-algebra. We use upper case letters for random variables such as $X,Y$ and the range of random variables is denoted by $R(.)$. We use $S(P)$ as a short hand notation to represent support of the probability measure. We use $x\in S(P)$ to identify outcomes of $X$ with non zero probability. In our construction and similar to MPT, we do not consider $\theta$ as a random variable; only a numerical representation for the event $M_{\theta}$. For mapping $f$ with a vector space domain and range, the $y$-th entry of the output vector is represented by $f(x)[y]$.

Overview of MPT. Consider the probability space $(\Omega,\Sigma,P)$, the random variable $Z$ and some event $M_{\theta}\in\Sigma$, where $\theta$ is a vector in $\Re^{t}$. MPT shows that the following inequality holds,

$\displaystyle P(M_{\theta})\leq\underset{z\in R(Z)}{\min}\left\{\frac{P(z)}{P(z|M_{\theta})}\right\}.$

(1)

In MPT, probabilistic models are represented by events such as $M_{\theta}$, with a known conditional distribution. To regularize the probabilistic models, the upper bound of the model’s probability is maximized. The parameter vector $\theta$ is not necessarily considered as a random vector; therefore, explicit definition of a prior on the parameters is not necessary. As $\theta$ varies, the conditional distribution of the model over $Z$ changes, leading to change in $P(M_{\theta})$. Note that the probability upper-bound is equal to 1, only if $P(z)=P(z|M_{\theta})$, therefore, $-\log(P(M_{\theta}))$ can serve as a measure of deviation of the model from the prior.

$\alpha$-Parametrization of the MPT. The softmin family of lower bounds for the maximum probability bound is used as a smooth approximation for the min function. The softmin family parameterized by $\alpha>0$ is defined as

$\displaystyle P_{\alpha}(M_{\theta})\triangleq\left(\sum_{z\in R(Z)}\left(\frac{P(z)}{P(z|M_{\theta})}\right)^{-\alpha}\right)^{-\frac{1}{\alpha}}.$

(2)

The following inequality relates the softmin family and the probability upperbound in (1):

$\displaystyle P_{\alpha}(M_{\theta})\leq\underset{z\in R(Z)}{\min}\left\{\frac{P(z)}{P(z|M_{\theta})}\right\},$

(3)

where the equality holds as $\alpha\to+\infty$ :[33]. The proposed objective function by [33] is $\mathcal{L}(\theta)=P(M_{\theta},M^{*})$, the probability of intersection of the model and the event $M^{*}$ representing the true underlying model (the dataset in our case). $M^{*}$ in our work is modeled with its conditional distribution. The conditional distribution of $M^{*}$ is the empirical distribution of the dataset. Throughout the paper, $M^{*}$ is called the oracle.

In this paper, we consider the classification problem-OOD detection. We define random variables $X,Y,Z$ to represent the input, label and $Z=(X,Y)$ as shorthand for the concatenation of input and the label. We can write the intersection objective function in its conditional form and treat each term separately,

$\displaystyle\mathcal{L}=\ln P(M^{*}|M_{\theta})+\ln P(M_{\theta}).$

(4)

In practice, $M_{\theta}$ and $M^{*}$ are implicitly modeled by their conditional distribution over the random variable, $P(X,Y|M_{\theta})$. To use the explicit representation, we write the objective function in the marginal form. Prior to deriving the objective function, the relation of $M^{*}$ and $M_{\theta}$ needs to be assumed. The choice of modeling the relation of $M^{*}$ and $M_{\theta}$ dictates the optimization process. Similar to [33], the conditional independence of model and oracle is explored here, namely $M_{\theta}\perp M^{*}|Z$.

Conditional Independence. If we assume conditional independence between the model and the oracle, we can continue with the marginalization process of the objective function $\mathcal{L}^{\perp}(\theta)$:

	$\displaystyle\mathcal{L}^{\perp}(\theta)\triangleq\ln P(M^{*}|M_{\theta})+\ln P(M_{\theta}),$		(5)
	$\displaystyle=\ln\sum_{z\in R(Z)}P(M^{*}|z,M_{\theta})P(z|M_{\theta})+\ln P(M_{\theta}).$		(6)

Considering the conditional independence assumption, $P(M^{*}|z,M_{\theta})=P(M^{*}|z)=P(z|M^{*})P(M^{*})/P(z)$. Assuming that $P(z)$ is uniform, we can treat both $P(z)$ and $P(M^{*})$ as constants and drop them for ease of notation. As a result,

$\displaystyle\mathcal{L}^{\perp}(\theta)=\ln\sum_{z\in R(Z)}P(z|M^{*})P(z|M_{\theta})+\ln P(M_{\theta}).$

(7)

Finally, to $\alpha$-parametrize $P(M_{\theta})$, we can define the following lower bound objective function

$\displaystyle\mathcal{L}^{\perp}_{\alpha}\triangleq\ln\sum_{z\in R(Z)}P(z|M^{*})P(z|M_{\theta})+\ln P_{\alpha}(M_{\theta}).$

(8)

Alternatively, because of the concavity of the log function and the Jensen inequality, we can find another lower bound, the cross entropy loss with MPT regularization:

	$\displaystyle\mathcal{L}_{\alpha}^{\perp}(\theta)\geq\sum_{z\in R(Z)}P(z|M^{*})\ln P(z|M_{\theta})+\ln P_{\alpha}(M_{\theta}),$		(9)
	$\displaystyle\mathcal{L}_{\textrm{cross}}(\theta)\triangleq\sum_{z\in R(Z)}P(z|M^{*})\ln P(z|M_{\theta}),$		(10)

where $\mathcal{L}_{\textrm{cross}}(\theta)$ is the conventional cross entropy objective function. The summary of the inequality relations is thus

$\displaystyle\mathcal{L}^{\perp}(\theta)\geq\mathcal{L}_{\alpha}^{\perp}(\theta)\geq\mathcal{L}_{\textrm{cross}}(\theta)+\ln P_{\alpha}(M_{\theta}).$

(11)

EBMs add another dimension to the training of classification models. The models are trained on the input distribution without the need for the labels. Here, we incorporate the MPT view of machine learning and the objective functions described in previous section in the formulation of EBMs.

Consider the training set consisted of $N_{t}$ entries and the test set is consisted of $N_{v}$ entries, i.e. $D_{t}=\{z^{(i)}\triangleq(x^{(i)},y^{(i)})\}_{i=1}^{N_{t}}$. The training and test set could be understood as an empirical distribution of the underlying true model. We denote the oracle as $M^{*}$. We adapt the intersection objective function described in MPT view. The classification model parameterized by $\theta\in\Re^{t}$ is represented by $f(x;\theta)$ where $f:\Re^{n}\to\Re^{m}$. We construct the model’s probability distribution $P(x,y|M_{\theta})$ using $f$:

$\displaystyle P(x,y|M_{\theta})=e^{f(x,\theta)[y]}P(x,y)/\sum_{x^{\prime},y^{\prime}\in S(P)}e^{f(x^{\prime},\theta)[y^{\prime}]}P(x^{\prime},y^{\prime}).$

(12)

Having Eq (12), and the probability upperbound in (1), we can write the maximum probability of $M_{\theta}$ as

$\displaystyle P(M_{\theta})\leq\frac{\sum_{x^{\prime},y^{\prime}\in S(P)}e^{f(x^{\prime},\theta)[y^{\prime}]}P(x^{\prime},y^{\prime})}{\underset{x,y\in S(P)}{\max}\left\{e^{f(x,\theta)[y]}\right\}}.$

(13)

For $P(M_{\theta})$ to be nonzero, either $f$ needs to be bounded or the prior needs to have bounded support. We denote the normalizing constant in (12) by $\eta(\theta)$. Under independence assumption, when we add the log-likelihood term, $\eta(\theta)$ will cancel out. Therefore, in our objective functions only the support of the prior affects the regularization term. In above equation, we used the prior in the construction of our model. Without the prior, the probability distribution is not necessarily normalizable [2], which leads to instability of the optimization.

We used a 12 layer CNN to train on CIFAR10, CIFAR100 and MNIST datasets. We chose to use CReLU activation [40] instead of the ReLU because of its gradient norm preserving property. For all convolution layers we used 64 filters. The CNN does not have Batch Normalization [18] and Dropout  [15]. We avoided any other probabilistic layer to avoid their interference with the optimization. Including probabilistic layers in our experiments might interact with the objective function we proposed. Their effect is out of the scope of the current study. The purpose of using the described model was to experiment with a backbone and simple architecture. We avoided hyperparameter tuning as much as possible and tested all the variants with the same hyperparameter range. The learning rate was fixed to 0.01. We experimented with different batch sizes to include the effect of the bias of subgradients in our experiment. The batch sizes in our experiments were chosen from 64, 128, 256. Each hyperparameter setup was trained 5 times. The total number of networks trained are 360 for each dataset.

The experiment is designed to show the effects of using MPT with the proposed objective functions on the CNN and its energy based variant on the classification task. In MPT based objective functions, the $\alpha$ hyper parameter determines how much the model is regularized. A comparative study is embedded in our experiments. Setting $\alpha$ to 1, removes the regularization effect of MPT and is equivalent to the usual likelihood training methods. We have included $\alpha=1$ to compare MPT with the conventional cross entropy and energy based loss.

Our experiments are three fold, (i) we experiment with classification task to show the stability of training and generalization performance having the MPT regularization. (ii) we experiment the robustness of trained networks to the additive noise to input. (iii) We demonstrate the effect of MPT regularization on OOD task. We use the the popular metric, Area Under the Receiver Operating Characteristics (AUROC) to evaluate the OOD performance.

The empirical aspect of our research is mainly ablation studies and proof of concept for effectiveness of MPT. Our empirical evidence holds, at least, when the model architecture and the hyper parameters are not tuned to the dataset. It is hard to confidently assert that MPT works for the vast variations of models that currently exist in the machine learning scope. However, so far we have not identified any incompatible probabilistic model in our empirical evidence. Furthermore, since the proof for MPT does not rely on any assumptions, it is reasonable to expect that the framework is general enough.

Given that MPT is theoretically well grounded, and offers a black box view on priors, it is a promising direction for further investigation. Also, the empirical evidence, although naturally limited, is inline with what is expected from theory. For example, the generalization of the model using MPT can be tied to smoothness of model’s distribution on input/label, $P(x,y|M_{\theta})$. The test accuracy of a model is only dependant on the test data distribution and the model’s distribution. Maximizing the model probability forces the model’s distribution to be close to the prior distribution (e.g. Uniform). If the prior is smooth, maximizing the probability of the model enforces the model’s distribution to be smoother. Models assuming smoother conditional distributions are more invariant to changes of input/label distribution. Therefore, the smoother models are more robust to the difference of the train and test distribution. We intend to follow up on this thought for future research and verify whether probability of the models could be tied theoretically to their generalization performance.

MPT does not restrict the possibility of having explicit priors on parameters. MPT provides an upper bound for the prior density of parameters. Any explicit assumption made about the parameters can be combined with MPT to form a new prior. For example including the explicit prior $q(\theta)$ can be done by constructing the density, $p(\theta)=P_{\alpha}(M_{\theta})q(\theta)$ followed by normalization. So far the benefits and adverse effects of combining explicit priors is not clear to us. It is possible to speculate that the gradients of regularization may be affected by the gradient vanishing problem in the MPT framework. Having explicit priors could help with the gradient vanishing of regularization. We intend to test the effect of incorporating explicit priors to the MPT framework in our future experiments.

In this paper, we have incorporated Maximum Probability Theorem into training of Energy Based Models (EBMs) and demonstrated its black-box regularization property in classification and out of distribution detection problems. Obtained from six publicly available datasets, our experiments demonstrated that (1) learning input and the label jointly in EBMs with adequate regularization outperform CNNs with softmax layer. (2) MPT regularization, without any exception in our experiments, increases the stability of training, lowers the variance of test accuracy, and improves generalization performance. (3) Incorporating MPT regularization improves AUROC in the OOD tasks.