Deep Discriminative to Kernel Density Graph for In- and Out-of-distribution Calibrated Inference

Consider a supervised learning problem with independent and identically distributed training samples $\{(\mathbf{x}_{i},y_{i})\}_{i=1}^{n}$ such that $(\mathbf{X},Y)\sim P_{X,Y}$, where $\mathbf{X}\sim P_{X}$ is a $\mathcal{X}\subseteq\mathbb{R}^{D}$ valued input and $Y\sim P_{Y}$ is a $\mathcal{Y}=\{1,\cdots,K\}$ valued class label. Let $\mathcal{S}$ be the high density region of the marginal, $P_{X}$, thus $\mathcal{S}\subsetneq\mathcal{X}$. Here the goal is to learn a confidence score, $\mathbf{g}:\mathbb{R}^{D}\rightarrow[0,1]^{K}$, $\mathbf{g}(\mathbf{x})=[g_{1}(\mathbf{x}),g_{2}(\mathbf{x}),\dots,g_{K}(\mathbf{x})]$ such that,

where $P_{Y|X}(y|\mathbf{x})$ is the posterior probability for class $y$ given by the Bayes formula:

Here $P_{X|Y}(\mathbf{x}|y)$ is the class conditional density which we will refer as $f_{y}(\mathbf{x})$ hereafter for brevity.

Deep discriminative networks partition the feature space $\mathbb{R}^{d}$ into a union of $p$ affine polytopes $Q_{r}$ such that $\bigcup_{r=1}^{p}Q_{r}=\mathbb{R}^{d}$, and learn an affine function over each polytope [4, 21]. Mathematically, the unnormalized class-conditional density for the label $y$ estimated by these deep discriminative models at a particular point $\mathbf{x}$ can be expressed as:

For example, in the case of a decision tree, $\mathbf{a}_{r}=\mathbf{0}$, i.e., decision tree assumes uniform distribution for the class-conditional densities over the leaf nodes. Among these polytopes, the ones that lie on the boundary of the training data extend to the whole feature space and hence encompass all the OOD samples. Since the posterior probability for a class is determined by the affine activation over each of these polytopes, the algorithms tend to be overconfident when making predictions on the OOD inputs. Moreover, there exist some polytopes that are not populated with training data. These unpopulated polytopes serve to interpolate between the training sample points. If we replace the affine activation function of the populated polytopes with Gaussian kernels and prune the unpopulated ones, the tail of the kernel will help interpolate between the training sample points while assigning lower likelihood to the low density or unpopulated polytope regions of the feature space. This results in better confidence calibration for the proposed modified approach.

We will call the above discriminative approaches as the ‘parent approach’ hereafter. Consider the collection of polytope indices $\mathcal{P}$ from the parent approach which are populated by the training data. We replace the affine functions over the populated polytopes with Gaussian kernels $\mathcal{G}(\cdot;\hat{\mu}_{r},\hat{\Sigma}_{r})$. For a particular inference point $\mathbf{x}$, we consider the Gaussian kernel with the minimum distance from the center of the kernel to the corresponding point:

where $\|\cdot\|$ denotes a distance. As we will show later, the type of distance metric considered in Equation 4 highly impacts the performance of the proposed model. In short, we modify Equation 3 from the parent ReLU-net or random forest to estimate the class-conditional density (unnormalized):

where $n_{y}$ is the total number of samples with label $y$ and $n_{ry}$ is the number of samples from class $y$ that end up in polytope $Q_{r}$. We add a small constant to the class conditional density $\tilde{f}_{y}$:

Note that in Equation 6, $\frac{b}{\log(n)}\rightarrow 0$ as the total training points, $n\rightarrow\infty$. The intuition behind the added constant will be clarified further later in Proposition 4.3. The confidence score $\hat{g}_{y}(\mathbf{x})$ for class $y$ given a test point $\mathbf{x}$ is estimated using the Bayes rule as:

where $\hat{P}_{Y}(y)$ is the empirical prior probability of class $y$ estimated from the training data. We estimate the class for a particular inference point $\mathbf{x}$ as:

We fit Gaussian kernel parameters to the samples that end up in the $r$-th polytope. We set the kernel center along the $d$-th dimension:

where $x^{d}_{i}$ is the value of $\mathbf{x}_{i}$ along the $d$-th dimension. We set the kernel variance along the $d$-th dimension:

where $\lambda$ is a small constant that prevents $\hat{\sigma}^{d}_{r}$ from being $0$. We constrain our estimated Gaussian kernels to have diagonal covariance.

For a high dimensional dataset with low training sample size, the polytopes are sparsely populated with training samples. For improving the estimate of the ratio $\frac{n_{ry}}{n_{y}}$ in Equation 5, we incorporate the samples from other polytopes $Q_{s}$ based on the similarity $w_{rs}$ between $Q_{r}$ and $Q_{s}$ as:

As $n\to\infty$, the estimated weights $w_{rs}$ should satisfy the condition:

For simplicity, we will describe the estimation procedure for $w_{rs}$ in the next sections. Note that if we satisfy Condition 12, then we have $\frac{\hat{n}_{ry}}{\hat{n}_{y}}\to\frac{n_{ry}}{n_{y}}$ as $n\to\infty$. Therefore, we modify Equation 5 as:

where $\hat{r}^{*}_{\mathbf{x}}=\operatornamewithlimits{argmin}_{r}\|\hat{\mu}_{r}-\mathbf{x}\|$. Now we use $\hat{f}_{y}(\mathbf{x})$ estimated using (13) in Equation (6), (7) and (8), respectively. Below, we describe how we estimate $w_{rs}$ for KDF and KDN .

Consider $T$ number of decision trees in a random forest trained on $n$ $iid$ training samples $\{(\mathbf{x}_{i},y_{i})\}_{i=1}^{n}$. Each tree $t$ partitions the feature space into $p_{t}$ polytopes resulting in a set of polytopes: $\{\{Q_{t,r}\}_{r=1}^{p_{t}}\}_{t=1}^{T}$. The intersection of these polytopes gives a new set of polytopes $\{Q_{r}\}_{r=1}^{p}$ for the forest. For any two points $\mathbf{x}\in Q_{r}$ and $\mathbf{x}^{\prime}\in Q_{s}$, we define the kernel $\mathcal{K}(r,s)$ as:

where $t_{rs}$ is the total number of trees, $\mathbf{x}$ and $\mathbf{x}^{\prime}$ end up in the same leaf node. Here, $0\leq\mathcal{K}(r,s)\leq 1$.

If the two samples end up in the same leaf in all the trees, i.e., $\mathcal{K}(r,s)=1$, they belong to the same polytope, i.e. $r=s$. In short, $\mathcal{K}(r,s)$ is the fraction of total trees where the two samples follow the same path from the root to a leaf node. We exponentiate $\mathcal{K}(r,s)$ so that Condition 12 is satisfied:

We choose $k$ using grid search on a hold-out dataset.

Consider a fully connected $L$ layer ReLU-net trained on $n$ $iid$ training samples $\{(\mathbf{x}_{i},y_{i})\}_{i=1}^{n}$. We have the set of all nodes denoted by $\mathcal{N}_{l}$ at a particular layer $l$. We can randomly pick a node $n_{l}\in\mathcal{N}_{l}$ at each layer $l$, and construct a sequence of nodes starting at the input layer and ending at the output layer which we call an activation path: $m=\{n_{l}\in\mathcal{N}_{l}\}_{l=1}^{L}$. Note that there are $N=\Pi_{l=1}^{L}{|\mathcal{N}_{l}|}$ possible activation paths for a sample in the ReLU-net. We index each path by a unique identifier number $z\in\mathbb{N}$ and construct a sequence of activation paths as: $\mathcal{M}=\{m_{z}\}_{z=1,\cdots,N}$. Therefore, $\mathcal{M}$ contains all possible activation pathways from the input to the output of the network.

While pushing a training sample $\mathbf{x}_{i}$ through the network, we define the activation from a ReLU unit at any node as ‘$1$’ when it has positive output and ‘$0$’ otherwise. Therefore, the activation indicates on which side of the affine function at each node the sample falls. The activation for all nodes in an activation path $m_{z}$ for a particular sample creates an activation mode $a_{z}\in\{0,1\}^{L}$. If we evaluate the activation mode for all activation paths in $\mathcal{M}$ while pushing a sample through the network, we get a sequence of activation modes: $\mathcal{A}_{r}=\{a_{z}^{r}\}_{z=1}^{N}$. Here $r$ is the index of the polytope where the sample falls in.

If the two sequences of activation modes for two different training samples are identical, they belong to the same polytope. In other words, if $\mathcal{A}_{r}=\mathcal{A}_{s}$, then $Q_{r}=Q_{s}$. This statement holds because the above samples will lie on the same side of the affine function at each node in different layers of the network. Now, we define the kernel $\mathcal{K}(r,s)$ as:

Note that $0\leq\mathcal{K}(r,s)\leq 1$. In short, $\mathcal{K}(r,s)$ is the fraction of total activation paths which are identically activated for two samples in two different polytopes $r$ and $s$. We exponentiate the kernel using Equation 15. Pseudocodes outlining the two algorithms are provided in Appendix D.

Consider $\mathcal{P}_{n}=\{Q_{1},Q_{2},\cdots,Q_{p}\}$ as a partition of $\mathbb{R}^{d}$ given by a random forest or a ReLU-net after being trained on $n$ training samples. We measure distance between two points $\mathbf{x}\in Q_{r},\mathbf{x}^{\prime}\in Q_{s}$ using the kernel introduced in Equation 14 and Equation 16, and call it ‘Geodesic’ distance [22]:

We use Geodesic distance to find the nearest polytope to the inference point. As Geodesic distance cannot distinguish between points within the same polytope, it has a resolution similar to the size of the polytope. For discriminating between two points within the same polytope, we fit a Gaussian kernel within the polytope (described above). As $h_{n}\to 0$, the resolution for Geodesic distance improves. In Section 5, we will empirically show that using Geodesic distance scales better with higher dimension compared to that of Euclidean distance.

Given $n$ training samples $\{(\mathbf{x}_{i},y_{i})\}_{i=1}^{n}$, we define the distance of an inference point $\mathbf{x}$ from the training points as: $d_{\mathbf{x}}=\min_{i=1,\cdots,n}\|\mathbf{x}-\mathbf{x}_{i}\|$, where $\|\cdot\|$ denotes Euclidean distance.

In vision data, each image pixel contains local information about the neighboring pixels. To extract the local information, we use convolutional or vision transformer encoders at the front-end. More precisely, we have a front-end encoder, $h_{e}:\mathbb{R}^{D}\mapsto\mathbb{R}^{m}$ and typically, $m<<D$. After the encoder there is a few fully connected dense layers for discriminating among the $K$ class labels, $h_{f}:\mathbb{R}^{m}\mapsto\mathbb{R}^{K}$. Note that the $m$-dimensional embedding outputs from the encoder are partitioned into polytopes by the dense layers (see Equation (3)) and we fit a KDN on the embedding outputs. The above approach results in extraction of better inductive bias by KDN from the parent model and makes KDN more scalable with larger parent models and training sample size.