Capsa: A Unified Framework for Quantifying Risk in Deep Neural Networks
The modern pervasiveness of large-scale deep neural networks (NNs) is driven
by their extraordinary performance on complex problems but is also plagued by
their sudden, unexpected, and often catastrophic failures, particularly on
challenging scenarios. Existing algorithms that provide risk-awareness to NNs
are complex and ad-hoc. Specifically, these methods require significant
engineering changes, are often developed only for particular settings, and are
not easily composable. Here we present capsa, a framework for extending models
with risk-awareness. Capsa provides a methodology for quantifying multiple
forms of risk and composing different algorithms together to quantify different
risk metrics in parallel. We validate capsa by implementing state-of-the-art
uncertainty estimation algorithms within the capsa framework and benchmarking
them on complex perception datasets. We demonstrate capsa's ability to easily
compose aleatoric uncertainty, epistemic uncertainty, and bias estimation
together in a single procedure, and show how this approach provides a
comprehensive awareness of NN risk.
Sadhana Lolla, Iaroslav Elistratov, Alejandro Perez, Elaheh Ahmadi, Daniela Rus, Alexander Amini
# _Capsa_: A Unified Framework for Quantifying Risk in Deep Neural Networks

###### Abstract

The modern pervasiveness of large-scale deep neural networks (NNs) is driven by their extraordinary performance on complex problems but is also plagued by their sudden, unexpected, and often catastrophic failures, particularly on challenging scenarios. Existing algorithms that provide risk-awareness to NNs are complex and ad-hoc. Specifically, these methods require significant engineering changes, are often developed only for particular settings, and are not easily composable. Here we present capsa, a framework for extending models with risk-awareness. Capsa provides a methodology for quantifying multiple forms of risk and composing different algorithms together to quantify different risk metrics in parallel. We validate capsa by implementing state-of-the-art uncertainty estimation algorithms within the capsa framework and benchmarking them on complex perception datasets. We demonstrate capsa's ability to easily compose aleatoric uncertainty, epistemic uncertainty, and bias estimation together in a single procedure, and show how this approach provides a comprehensive awareness of NN risk.

## 1 Introduction

Neural networks (NNs) continue to push the boundaries of modern artificial intelligence (AI) systems across a wide range of complex real-world domains, from robotics and autonomy (Bojarski et al., 2016; Hawke et al., 2020; Codevilla et al., 2018), to healthcare and medical decision making (Ching et al., 2018; Topol, 2019). While their performance in these domains remains unmatched, modern NNs still encounter sudden, unexpected, and inexplicable failures that are often catastrophic - especially in safety-critical environments. These failures are largely due to systemic issues that propagate throughout the entire modern AI lifecycle, from imbalances (He & Garcia, 2009; Buda et al., 2018) and noise (Beigman & Klebanov, 2009) in data that lead to algorithmic bias (Bolukbasi et al., 2016; Caliskan et al., 2017; Buolamwini & Gebru, 2018; Chen et al., 2018; Obermeyer et al., 2019; Seyyed-Kalantari et al., 2021) to predictive uncertainty (Kendall & Gal, 2017; Kompa et al., 2021; Nado et al., 2021; Amini et al., 2020) that plagues model performance on unseen or out-of-distribution data. In order to realize the widespread adoption of AI in society, NN models must not only identify these potential failure modes, but also effectively use this awareness to obtain unified and calibrated measures of risk and uncertainty. There is thus a critical need for unified systems that can estimate quantitative risk metrics for any NN model, and in turn integrate this awareness back into the learning lifecycle to improve robustness, generalization, and safety.

Existing algorithmic approaches to risk quantification narrowly estimate a singular form of risk in AI models, often in the context of a limited number of data modalities (Nix & Weigend, 1994;
Kendall & Gal, 2017; Lakshminarayanan et al., 2017; Buolamwini & Gebru, 2018; Zhang et al., 2018; Gilitschenski et al., 2019). These methods present critical limitations as a result of their reductionist, ad hoc, and narrow focus on single metrics of risk or uncertainty. However, generalizable methods that provide a larger holistic awareness of risk have yet to be realized and deployed (Nado et al., 2021; Tran et al., 2022). This is in part due to the significant engineering changes required to integrate an individual risk algorithm into a larger machine learning system (Tran et al., 2016; Dillon et al., 2017; Bingham et al., 2019; Shi et al., 2017), which in turn can impact the quality and reproducibility of results. The lack of a unified approach for composing different risk estimation algorithms or risk-aware models limits the scope and capability of each algorithm independently, and further limits the robustness of the system as a whole. A general, model-agnostic framework for extending NN systems with holistic risk-awareness, covering both uncertainty and bias, would advance the ability and robustness of end-to-end systems.

To address these fundamental challenges, we present capsa - an algorithmic framework for wrapping any arbitrary NN model with state-of-the-art risk-awareness capabilities. By decomposing the algorithmic stages of risk estimation into their core building blocks, we unify different algorithms and estimation metrics under a common data-centric paradigm. Additionally, because capsa allows the underlying NN to be aware of a variety of risk metrics in parallel, we achieve improved performance and quality in risk estimation through principled redundancy, and open the door to achieving a unified composition and hierarchical understanding of NN risk.

In summary, the key contributions of this paper are:

1. Capsa, a flexible, and easy-to-use framework for equipping any given neural network with calibrated awareness of different forms of risk - including bias, label noise, and predictive uncertainty;
2. An algorithm for decomposing different types of risk and their estimation methods into modular components that can in turn be integrated and composed together to achieve greater accuracy, robustness, and efficiency; and
3. Empirical validation of capsa on a range of dataset complexities and modalities, along with the application of capsa for mitigation of algorithmic bias, identification of label noise, and detection of anomalies and out-of-distribution data.

We refer readers to Capsa Pro(Amini et al., 2023) for information on the software library with the full functionality described in this publication.

## 2 Background and Methodology

### Preliminaries

We consider the problem of supervised learning, where we are given a labeled dataset of \(n\) input, output pairs, \(\{x,y\}_{i=1}^{n}\). Our goal is to learn a model, \(f\), parameterized by weights, \(\mathbf{W}\), that minimizes the average loss over the entire dataset: \(\sum_{i}\mathcal{L}(f_{\mathbf{W}}(x),y)\). While traditionally, the model,

Figure 1: Capsa unifies state-of-the-art algorithms for quantifying neural network risks ranging from (A) under-representation bias, (B) epistemic (model) uncertainty, and (C) aleatoric uncertainty (label noise). Capsa converts existing models into risk-aware variants, capable of identifying risks efficiently during training and deployment.
a neural network, outputs predictions in the form of \(\hat{y}=f_{\mathbf{W}}(x)\), we now introduce a risk-aware transformation operation, \(\Phi\), which transforms a model, \(f\), into a risk-aware variant, such that

\[g=\Phi_{\mathbf{\theta}}(f_{\mathbf{W}}),\] \[\hat{y},R=g(x),\]

where \(R\) are the estimated "risk" measures from a set of metrics, \(\mathbf{\theta}\). The goal of this paper is to propose a common transformation backbone for \(\Phi_{\mathbf{\theta}}(\cdot)\), which automatically transforms an arbitrary model, \(f\), to be aware of risks, \(\mathbf{\theta}\).

All measures of risk aim to capture, on some level, the reliability of a given prediction. This can stem from the data source (aleatoric uncertainty, or representation bias) or the predictive capacity of the model itself (epistemic uncertainty). Within capsa, we define various risk metrics to identify and measure these sources of risk. We propose the idea of _wrappers_, which are instantiations of \(\Phi_{\theta}\), for a singular risk metric, \(\theta\). Wrappers are given an arbitrary neural network and, while preserving the structure and function of the network, add and modify the relevant components in the model. This allows them to serve as a drop-in replacement that is able to estimate the risk metric, \(\theta\). Wrappers can be further composed using a set of metrics, \(\mathbf{\theta}\), that are faster and more accurate than individual metrics.

### Capsa: The Wrapping Algorithm

While risk estimation algorithms can take a variety of forms and are often developed in ad hoc settings, we present a unified algorithm for building \(\Phi_{\mathbf{\theta}}\) in order to wrap an arbitrary neural network model. There are four main components: (1) constructing the shared feature extractor, (2) applying modifications to the existing model needed to capture the uncertainty, (3) creating additional models and augmentations if necessary, and (4) modifying the loss functions.

The feature extractor, which defaults to the model until its last layer, can be leveraged as a shared backbone by multiple wrappers at once to predict multiple compositions of risk. This results in a fast, efficient method of reusing the main body of the model, which does not require training multiple models and risk estimation methods from scratch. Next, capsa modifies the existing network according to metric-specific modifications; for example, this could entail modifying every weight in the model to be drawn from a distribution (to convert to a Bayesian neural network (Blundell et al., 2015)) or adding stochastic dropout layers (Gal and Ghahramani, 2016). Depending on the metric, capsa also adds new layers or augmentations to the model that predict new outputs. Note that these are not modifications to the model, but rather augmentations for the given metric; for example, new layers to output \(\sigma\)(Nix and Weigend, 1994), or extra model copies when ensembling (Lakshminarayanan et al., 2017). Lastly, we modify the loss function to capture any remaining metric-specific changes that need to be made. This entails combining the user-specified and metric-specific loss functions (e.g., KL-divergence (Kingma and Welling, 2013), negative log-likelihood (Nix and Weigend, 1994), etc). All of the following modifications are integrated together into a custom metric-specific forward pass, and train step. These are used to capture variations in the forward and backward passes during training and inference.

Figure 2: **Overview of Capsa architecture**. (A) Capsa converts arbitrary NN models into risk-aware variants, that can simultaneously predict both their output along with a list of user-specified risk metrics. (B) Each risk metric forms the basis of a singular model wrapper which is constructed through metric-specific modifications to the model architecture and loss function.
### Risk Metrics and Background

In this section, we outline three high-level categories of risk which we quantitatively define and estimate.

**Representation Bias -** The representation bias of a dataset uncovers imbalance in the space of features and captures whether certain combinations are more prevalent than others. Note that this is fundamentally different from label imbalance, which only captures distributional imbalance in the labels. For example, in driving datasets, it has been demonstrated that the combination of straight roads, sunlight, and absence of traffic is higher than any other feature combinations. This indicates that these samples are overrepresented (Amini et al., 2018). Similar combinations have been identified for facial detection (Buolamwini and Gebru, 2018; Amini et al., 2019), medicine (Puyol-Anton et al., 2021; Soleimany et al., 2021), and clinical trials (Xu et al., 2022). Uncovering feature representation bias is a computationally expensive process as these features are (1) often unlabeled, and (2) extremely high-dimensional (e.g., images, videos, language, etc). However, they can be estimated by learning the density distribution of the data. We accomplish this by estimating densities in feature space. For high-dimensional feature spaces we estimate a low-dimensional embedding using a variational autoencoder (Kingma and Welling, 2013) or by using the features from the penultimate layer of the model. Bias is then the imbalance between parts of the density space estimated either discretely (using a discretely-binned histogram) or continuously (using a kernel distribution (Rosenblatt, 1956)).

**Aleatoric Uncertainty-** Aleatoric uncertainty captures noise in the data, e.g., mislabeled datapoints, ambiguous labels, classes with low separation, etc. We model aleatoric uncertainty using Mean and Variance Estimation (MVE) (Nix and Weigend, 1994). In the regression case, we pass the outputs of the model's feature extractor to another layer that predicts the standard deviation of the output. We train using NLL, and use the predicted variance as an estimate of the aleatoric uncertainty. We apply a modification to the algorithm to generalize to the classification case in Alg. 1. We assume the classification logits are drawn from a normal distribution and stochastically sample from them using the reparametrization strategy. We average stochastic samples and backpropogate using cross entropy loss through logits and their inferred uncertainties.

**Epistemic Uncertainty-** Epistemic uncertainty measures uncertainty in the model's predictive process - this captures scenarios such as examples that are "hard" to learn, examples whose features are underrepresented, and out-of-distribution data. We provide a unified approach for a variety of epistemic uncertainty methods ranging from Bayesian neural networks (Blundell et al., 2015), ensembling (Lakshminarayanan et al., 2017), and reconstruction-based (Kingma and Welling, 2013) approaches. Below, we outline three metrics and how they fit into capsa's unified risk estimation framework.

A _Bayesian neural network_ can be approximated by stochastically sampling, during inference, from a neural network with probabilistic layers (Blundell et al., 2015; Gal and Ghahramani, 2016). Adding dropout layers (Srivastava et al., 2014) to a model is one of the simplest ways to capture epistemic uncertainty (Gal and Ghahramani, 2016). To calculate the uncertainty, we run \(T\) forward passes, which is equivalent to Monte Carlo sampling. Computing the first and second moments from the \(T\) stochastic samples yields a prediction and uncertainty estimate, respectively.

An _ensemble_ of \(N\) models, each a randomly initialized stochastic sample, is a common approach used to accurately estimate epistemic uncertainty (Lakshminarayanan et al., 2017). However, this comes with significant computational costs. To reduce the cost of training ensembles, capsa automates the construction and management of the training procedure for all members and parallelizes their computation.

_Variational autoencoders (VAEs)_ are typically used to learn a robust, low-dimensional representation of the latent space. They can be used as to estimate epistemic uncertainty by using the reconstruction loss \(MSE(\hat{x},x)\). In cases of out-of-distribution data, samples that are hard to learn, or underrepresented samples, we expect that the VAE will have high reconstruction loss, since the mapping to
the latent space will be less accurate. Conversely, when the model is very familiar with the features, or the data is in distribution, we expect the latent space mapping to be robust and the reconstruction loss to be low. To construct the VAE for any given model in capsa, we use the feature extractor as the encoder, and reverse the feature extractor automatically when possible to create a decoder.

### Metric Composability

Using capsa, we compose multiple risk metrics to create more robust estimates (e.g., by combining multiple metrics, or alternatively by capturing different measures of risk independently). By using the feature extractor as a shared common backbone, we can optimize for multiple objectives, ensemble multiple metrics, and obtain different types of uncertainty estimates simultaneously.

We propose a novel composability algorithm within capsa to automate this process. Again, we leverage our shared feature extractor as the common backbone for all metrics and incorporate all model modifications. Then, we apply the new model augmentations either in series or in parallel, depending on the use case (i.e., we can ensemble a metric in series to average the metric over multiple joint trials, or we can apply ensembling in parallel to estimate a independent measure of risk). Lastly, the model is jointly optimized using all of the relevant loss functions by computing the gradient of each one with regard to the shared backbone's weights and stepping into the direction of the accumulated gradient.

## 3 Experimental Results

In the following section, we analyze the risk metrics obtained by wrapping various models with capsa on several datasets. We show that capsa provides accurate, scalable, composable risk metrics that are efficient and can be used to quantify bias, aleatoric, and epistemic uncertainty using multiple methods.

### Representation Bias

Using capsa's bias and epistemic wrapper capabilities, we analyzed the Celeb-A (Liu et al., 2015) dataset. The task for the neural network was to detect faces from this dataset against non-face images (collated from various negatives in the ImageNet dataset). Fig. 3A quantifies an accuracy vs bias tradeoff that neural networks exhibit, where they tend to perform better on overrepresented training features. We used a VAE as a feature extractor to demonstrate that the it can be used for single-shot bias and epistemic uncertainty estimation without any added computation cost.

Fig. 3B qualitatively inspects the different percentiles of bias ranging from underrepresentation (left) to overrepresentation (right). We found that the underrepresented samples in the dataset commonly contained darker skin tones, darker lighting, and faces not looking at the camera. As the percentile

Figure 3: **Bias and Epistemic Uncertainty on Faces** (A) Under-represented and over-represented faces in the Celeb-A dataset found by capsa using the VAE and HistogramBias wrappers. As the percentile bias of the data increases, the skin tone gets lighter, lighting gets brighter, and hair color gets lighter, and (B) accuracy on these datapoints increases. We also determine the points with the highest epistemic uncertainty, which have artifacts such as sunglasses, hats, colored lighting, etc.
of the bias gets higher, we see that the dataset is biased towards lighter skin tones, hair colors, and a more uniform facial direction. With our approach, we highlight a critical difference between bias and epistemic estimation methods in Fig. 3C. The samples estimated to have the highest epistemic uncertainty were not necessarily only underrepresented, but also contain features that obscure the predictive power of the model (e.g., faces with colored lighting, covering masks, and artifacts such as sunglasses and hats).

### Aleatoric Uncertainty

Next, we experiment on capsa's ability to successfully detects label noise in datasets using aleatoric uncertainty estimation. An example of this can be shown in Fashion-MNIST, which contains two very similar classes: "tshirt/top" and "shirt". The methods presented in capsa identify samples in Fashion-MNIST with high aleatoric uncertainty, which are light sleeveless tops with similar necklines with minimal visual differences. Short-sleeved shirts with round necklines are also classified as either category. Compared to randomly selected samples from these two classes, the samples considered noisy by capsa are visually indistinguishable, and difficult for humans (and models) to categorize.

### Epistemic Uncertainty

In this section, we benchmark capsa's epistemic methods on toy datasets. We demonstrate how capsa's ability to compose multiple methods (e.g., dropout and VAEs) can achieve more robust, efficient performance. We combine aleatoric methods with epistemic methods (i.e., ensembling the MVE metric) to strengthen aleatoric methods, since they are averaged across multiple runs. We can also treat the ensemble of MVEs as a mixture of normals. Similarly, to combine VAE and dropout, we use a weighted sum of their variances or we run the VAE \(N\) times with dropout layers and treat multiple runs as \(N\) normals.

#### 3.3.1 Cubic Dataset and UCI Benchmarking

We compose various epistemic and aleatoric methods on a cubic dataset with injected aleatoric noise and a lack of data in some parts of the test set. We train models on \(y=x+\epsilon\), where \(\epsilon\sim\mathcal{N}(1.5,0.9)\). Training data is within \([-4,4]\) and test within \([-6,6]\). Fig. 5 demonstrates that composed metrics can successfully detect regions with no data, as well as the aleatoric uncertainty in the center.

Additionally, we benchmark raw epistemic uncertainty methods on real-world regression datasets, and evaluate VAEs, ensembles, and dropout uncertainty on these datasets based on Root Mean Squared Error (RMSE) and negative log-likelihood (NLL) in Tab. 1. More composability results, as well as training times for all methods, are available in the appendix in Tab. 4 and Tab. 3.

Figure 4: **Fashion MNIST Aleatoric Uncertainty** (A) Randomly selected samples from two classes of fashion-mnist. These samples are visually distinguishable, and have a low aleatoric uncertainty, as opposed to (B), which shows samples with highest estimated aleatoric noise. It is not clear what features distinguish these shirts from tshirts/tops, as they have similar necklines, sleeve lengths, and cuts.

Figure 5: **Risk metrics on cubic regression. A regression dataset \(y=x+\epsilon\), where \(\epsilon\) is drawn from a Normal centered at \(x=1.5\). Models are trained on \(x\in[-4,4]\) and tested on \(x\in[-6,6]\). Composing using MVE results in a single metric that can seamlessly detect epistemic and aleatoric uncertainty without any modifications to the model construction or training procedure.**
#### 3.3.2 Depth Estimation

In this section, we transition to more complex models and datasets and demonstrate how capsa can be used as a large-scale risk and uncertainty benchmarking framework for existing methods. To that end, we train a U-Net style model on the task of monocular end-to-end depth estimation (see Tab. 2). Importantly, capsa works "out of the box" without requiring any modifications since it is a highly configurable, model-agnostic framework with modularity as one of the core of its design principles.

Specifically, we take a U-Net style model whose final layer outputs a single \(H\times W\) activation map and wrap it with capsa. We then train the wrapped model on NYU Depth V2 dataset (Nathan Silberman & Fergus, 2012) (27k RGB-to-depth image pairs of indoor scenes) and evaluate on a disjoint test-set of scenes. Additionally, we use outdoor driving images from ApolloScapes (Liao et al., 2020) as OOD data points.

We see that when we wrap the model with an aleatoric method in Fig. 10, we can successfully detect label noise or mislabeled data. The model exhibits increased aleatoric uncertainty on object boundaries. Indeed, we see that the ground truth has noisy labels particularly along the edges of objects which could be due to sensor noise or motion noise.

With dropout (Fig. 11) or ensemble (Fig. 12) wrappers, we capture uncertainty in the model's prediction. We see that increased epistemic uncertainty roughly corresponds to the semantically and visually challenging pixels where the model returns erroneous output.

\begin{table}
\begin{tabular}{c c c c|c c c} \hline \hline  & \multicolumn{2}{c|}{**RMSE**} & \multicolumn{2}{c}{**NLL**} & \multicolumn{2}{c}{**OOD AUC**} \\ \hline Boston & 2.449 +/- 0.134 & 2.323 +/- 0.117 & 2.589 +/- 0.113 & 2.282 +/- 0.03 & 2.497 +/- 0.047 & 2.253 +/- 0.056 \\ Power-Plant & 4.327 +/- 0.030 & 4.286 +/- 0.0120 & 4.221 +/- 0.028 & 2.892 +/- 0.00 & 2.964 +/- 0.000 & 2.841 +/- 0.005 \\ Yacht & 1.540 +/- 0.133 & 1.418 +/- 0.222 & 1.393 +/- 0.0965 & 2.399 +/- 0.03 & 2.637 +/- 0.131 & 1.035 +/- 0.116 \\ Concrete & 6.628 +/- 0.286 & 6.382 +/- 0.101 & 6.456 +/- 0.846 & 3.427 +/- 0.042 & 3.361 +/- 0.016 & 3.139 +/- 0.115 \\ Naval & 0.004 +/- 0.000 & 0.004 +/- 0.000 & 0.004 +/- 0.000 & 1.453 +/- 0.667 & -2.482 +/- 0.229 & -3.542 +/- 0.015 \\ Energy & 1.661 +/- 0.090 & 1.377 +/- 0.091 & 1.349 +/- 0.175 & 2.120 +/- 0.022 & 1.999 +/- 0.113 & 1.395 +/- 0.066 \\ Kin8nm & 0.088 +/- 0.001 & 0.0826 +/- 0.001 & 0.072 +/- 0.000 & -0.972 +/- 0.01 & -0.913 +/- 0.00 & -1.26 +/- 0.008 \\ Protein & 4.559 +/- 0.031 & 4.361 +/- 0.0156 & 4.295 +/- 0.029 & 4.452 +/- 0.012 & 3.345 +/- 0.011 & 2.723 +/- 0.023 \\ \hline \hline \end{tabular}
\end{table}
Table 1: **Regression benchmarking on the UCI datasets**

\begin{table}
\begin{tabular}{c c c c} \hline \hline  & **Test Loss** & **NLL** & **OOD AUC** \\ \hline Base & 0.0027 \(\pm\) 0.0002 & – & – \\ \hline VAE & 0.0027 \(\pm\) 0.0001 & – & 0.885 \(\pm\) 0.0361 \\ Dropout & 0.0027 \(\pm\) 0.0001 & 0.1397 \(\pm\) 0.0123 & 0.9986 \(\pm\) 0.0026 \\ Ensembles & **0.0023 \(\pm\) T-05** & 0.0613 \(\pm\) 0.0217 & 0.9989 \(\pm\) 0.0018 \\ MVE & 0.0036 \(\pm\) 0.0010 & **0.0532 \(\pm\) 0.0224** & 0.9798 \(\pm\) 0.0118 \\ \hline Dropout + MVE & 0.0027 \(\pm\) 0.0001 & 0.1291 \(\pm\) 0.0146 & 0.9986 \(\pm\) 0.0026 \\ VAE + Dropout & 0.0027 \(\pm\) 0.0001 & **0.0932 \(\pm\) 0.0201** & **0.9988 \(\pm\) 0.0024** \\ VAE + MVE & 0.0034 \(\pm\) 0.0012 & 0.1744 \(\pm\) 0.0156 & 0.9823 \(\pm\) 0.0102 \\ \hline \hline \end{tabular}
\end{table}
Table 2: **Depth regression results.** VAE + dropout outperforms all other epistemic methods and is more efficient.

Figure 6: **Risk estimation on monocular depth prediction**. (A) Example pixel-wise depth predictions and uncertainty. Model uncertainty calibration for individual metrics (B) and composed metrics (C). OOD detection assessed via AUC-ROC (D) and a full p.d.f. histogram (E).
## 4 Applications

The benefits of seamlessly and efficiently integrating a variety of risk estimation methods into arbitrary neural models extends far beyond benchmarking and unifying these algorithms. In this section, we outline critically important applications that are possible with the estimation abilities in capsa.

### Debiasing Facial Recognition Systems

Using the bias tools provided capsa, one application is to not only estimate and identify imbalance in the dataset (which we show also leads to performance bias) but to actively reduce the performance bias by adaptively re-sampling datapoints depending on their estimated representation bias during the course of training. As shown in Fig. 7, using capsa, we pinpoint exactly which samples need under/oversampling, and therefore can intelligently resample from the dataset during training. The benefits of this are twofold - we can improve sample efficiency by training on less data if some data is redundant, and we can also oversample from areas of the dataset where our latent representation is more sparse.

By composing multiple risk metrics together (in this case, VAEs and histogram bias) we can achieve even greater robustness during training, more sample efficiency, and combine epistemic uncertainty and bias to reduce risk.

### Detecting mislabeled examples

Another application of capsa is cleaning mislabeled or noisy datasets. We previously described capsa's ability to find noisy labels with high accuracy in Sec. 3.2. In the following experiment, we replaced a random collection of the 7s in the MNIST dataset with 8s. As shown in Fig. 8(A), the samples with high aleatoric uncertainty are dominated by the mislabeled examples, and also include a naturally mislabeled sample. We further test capsa's sensitivity to mislabeled datasets by artificially corrupting our labels with varying levels of probability \(p\). In Fig. 8B, as \(p\) increases, the average aleatoric uncertainty per class also increases. These experiments highlight capsa's capability to serve as the backbone of a dataset quality controller and cleaner, due to its high-fidelity aleatoric noise detection.

### Anomaly and Adversarial Noise Detection

Another application of the "uncertainty estimation" functionality provided by capsa is anomaly detection. The core idea behind this approach is that a model's epistemic uncertainty on out-of-distribution (OOD) data is naturally higher than the same model's epistemic uncertainty on in-distribution (ID) data. Thus, given a risk aware model, we visualize density histograms of per image uncertainty estimates provided by a model on both ID (unseen test-set for NYU Depth V2 dataset) and OOD data (ApolloScapes) (see Fig. 6E). At this point, OOD detection is possible by a simple thresholding. We use AUC-ROC to quantitatively assess the separation of the two density histograms, a higher AUC indicates a better quality of the separation (see Fig. 6D).

Figure 7: **Debiasing Facial Recognition Systems** (A) Facial datasets are overwhelmingly biased towards light-skinned females. (B) The feature combinations present in dark-skinned males make up only 1.49% of the dataset, and those present in dark-skinned female faces only take up 8.18% of the dataset. Since capsa identifies underrepresented datapoints, we can implement smart sampling schemes that increase the representation of these feature combinations.
It is critical for a model to recognize that it is presented with an unreasonable input (e.g., OOD). This capability could be used for autonomous vehicles that yield control to humans when model performance is expected to be poor. For example, in Fig. 9A we see that depth estimates degrade as images drift from the distribution. We are able to detect these shifts and can use this information to avoid incorrect predictions.

Further, the approach described above could be used to detect adversarial attacks (perturbations). In Fig. 9A we see that even though the perturbed images are not immediately distinguishable to a human eye, the method described above successfully detects the altered input images.

Another way of interpreting the adversarial perturbations is as a way of gradually turning ID datapoints into OOD. Such a granular control allows for enhanced model introspection. In Fig. 9 we see that as the epsilon of the perturbation increases, the density histograms of per image uncertainty estimates on both the ID and perturbed images become more disentangled (B) and thus the quality of separation increases (C).

## 5 Conclusions

In this paper, we present a unified, model-agnostic framework for risk estimation, which allows for seamless and efficient integration of uncertainty estimates in a couple lines of code. Our approach opens new avenues for greater reproducibility and benchmarking of risk and uncertainty estimation methods. We validate the scalability and the convenience of the framework on a variety of datasets. We showcase how our method can compose different algorithms together to quantify different risk metrics efficiently in parallel. We demonstrate how the obtained uncertainty estimates can be used for downstream tasks. We further show how the framework yields interpretable risk estimation results that can provide a deeper insight into decision boundaries of NNs. We refer readers Capsa Pro(Amini et al., 2023) for details regarding our comprehensive implementation of the functionality described in this publication. Capsa has the goal of accelerating and unifying advances in the areas of uncertainty estimation and trustworthy AI. In the future, we plan to extend our approach to

Figure 8: **Mislabeled Examples in the MNIST dataset** (A) If we purposefully inject label noise into the MNIST dataset by labeling 20% of the 7s in the dataset as 8, the mislabeled items have the highest aleatoric uncertainty. We also find a naturally mislabeled sample in the dataset. (B) As the percentage of mislabeled items increases, the average measured aleatoric uncertainty per class also increases.

Figure 9: **Robustness under adversarial noise** Across increasing levels of adversarial perturbations: (A) Pixel-wise depth predictions and uncertainty visualizations, (B) Density histograms of per image uncertainty, (C) OOD detection assessed via AUC-ROC, (D) Calibration curves.
other data modalities including irregular types (graphs) and temporal data (sequences), as well as to support other model types and other risk metrics.
