Identifying Drivers of Predictive Aleatoric Uncertainty

## Abstract

Abstract Explainability and uncertainty quantification are two pillars of trustable artificial intelligence. However, the reasoning behind uncertainty estimates is generally left unexplained. Identifying the drivers of uncertainty complements explanations of point predictions in recognizing potential biases and model limitations. It additionally facilitates the detection of oversimplification in the uncertainty estimation process. Explanations of uncertainty enhance communication of and trust in decisions. They allow for verifying whether the main drivers of model uncertainty are relevant and may impact model usage in certain applications. So far, the subject of explaining uncertainties has been rarely studied. The few exceptions in existing literature are tailored to Bayesian neural networks or rely heavily on technically intricate approaches, such as auxiliary generative models, thereby hindering their broad adoption. We propose variance feature attribution, a simple and scalable solution to explain predictive aleatoric uncertainties. First, we estimate uncertainty as predictive variance by adapting a neural network, for example, by equipping it with a Gaussian output distribution. We achieve this by adding a variance output neuron and can thereby rely on pre-trained point prediction models and fine-tune them for meaningful variance estimation. Second, we apply out-of-the-box explainers on the variance output of these models to explain the uncertainty estimation. This two-step method can be easily applied to any neural network with model-agnostic or model-specific explainers. We evaluate our approach in a synthetic setting where the data-generating process is known. We show that our method can explain uncertainty influences more reliably and faster than the established literature baseline CLUE, while the uncertainty estimation stage does not impede the accuracy of the model.
As an illustrative application, we fine-tune a state-of-the-art age regression model to estimate uncertainty and generate attributions for age prediction uncertainty. Our exemplary explanations highlight reasonable potential sources of uncertainty, such as laugh lines and frowning. Variance feature attribution provides accurate explanations for uncertainty estimates with little modifications to the model architecture and low computational overhead.

## 1 Introduction

Researchers have recognized the importance of uncertainty quantification and explainability of machine learning (ML) predictions to ensure the successful real-world adoption of ML-based systems in safety-critical applications (Abdar et al., ; Vilone & Longo, ). These dimensions serve as key indicators of a model’s trustworthiness, reliability, and fairness, which are crucial for its broad acceptance and use (Lambert et al., ; Lötsch et al., ).
Predictive uncertainty in ML refers to the degree of confidence associated with a model’s predictions (Chua et al., ). It can be decomposed into an epistemic and aleatoric component (Kendall & Gal, ). Epistemic uncertainty arises from the scarcity of data in specific areas of the input space, for example, because a particular condition may be underrepresented. Principally, it can be reduced by acquiring additional examples of this condition. Aleatoric uncertainty refers to inherent randomness or variability in the data, representing uncertainty that cannot be reduced by including additional training examples. Such uncertainty can arise due to measurement errors or certain variables relevant to the observed process not being collected. Uncertainty estimation is critical in risk management. It allows taking conservative action, relying on the model only when it exhibits a high degree of confidence in its predictions, and avoiding usage outside its area of competence (Kompa et al., ).

Explainability encompasses methods that enhance the transparency of ML models by highlighting how features influence the model’s output or by rendering the internal computations of black-box models more interpretable. Explainability methods enable understanding whether a model has learned relevant patterns from the input data and can reveal interesting previously unknown associations (Samek et al., ; Schwalbe & Finzel, ). Uncertainty quantification and explainability ensure accountable, informed and, therefore, responsible decision making and help mitigate biases and risks (Bhatt et al., ; Zhou et al., ; McGrath et al., ).

Figure: Figure 1: Overview of the variance feature attribution pipeline. (A) An original point prediction model with an output layer with weight matrix ${\bm{W}}\in\mathbb{R}^{d\times 1}$ and a scalar bias. We equip this model with a Gaussian distribution resulting in (B), a model with output weight matrix ${\bm{W}}\in\mathbb{R}^{d\times 2}$ and bias $b\in\mathbb{R}^{2}$. The two outputs are the mean $\hat{\mu}$ and the variance parameters $\hat{\sigma}^{2}$ of the predictive distribution. (C) From there, we can explain the variance using any suitable explainability method resulting in input-specific feature attributions that can be used to understand the drivers of the model’s uncertainty.
Refer to caption: /html/2312.07252/assets/x1.png

In the majority of applications, the focus of model explainability predominantly lies in interpreting point predictions (Vilone & Longo, ). There is a significant gap in understanding and explaining the drivers of uncertainty estimates. In practice, when a machine learning algorithm is deployed and yields a substantial uncertainty estimate for a specific instance, the possible courses of action involve abstaining from employing the model if alternatives are available or accepting an increased level of risk. With explainable uncertainties, however, users gain the capability to identify the factors contributing to elevated uncertainty levels. This understanding allows domain experts to judge their relevance in a given scenario. Additionally, it provides valuable insights concerning necessary modifications to augment the model’s predictive certainty and performance. In cases where abstaining from model usage is still necessary, factors influencing the decision can be understood and communicated. For example, if such an uncertainty factor is a feature indicating a person’s minority status, it could suggest a potential bias of the model. The bias would be undetectable by naive explanations if the feature only influences the uncertainty estimation but not the mean prediction.

Explanations can be categorized as either local or global (Schwalbe & Finzel, ; Adadi & Berrada, ). Local explanations provide insights into how a model makes predictions for a specific instance. A local explanation of the model’s uncertainty could foster more transparent discussions about ML-assisted decisions and risks, leading to increased trust. Global explanations provide an overview of a model’s behavior across the entire input space. In the realm of uncertainty explanations, they can be used to detect general drivers of uncertainty and certainty or to partition the input space into easy- vs. hard-to-predict regions. This knowledge can be utilized to formulate hypotheses to improve the model or to detect unintended shortcuts in the uncertainty estimation process, such as spurious correlations or biases.

There is little prior work on explaining uncertainties, and existing literature mainly focuses on classification and generally relies on Bayesian neural networks (BNNs) or technical intricacies such as auxiliary generative models (Antoran et al., ; Perez et al., ; Ley et al., ; Wang et al., ). BNNs assign probability distributions to network weights to capture uncertainty (MacKay, ). However, due to their computational complexity and altered training process, BNNs have not been as widely adopted as classical neural networks. This hinders the adoption of BNN uncertainty explanation approaches (Lakshminarayanan et al., ).

We propose a simple and scalable solution for explaining uncertainties that can be readily integrated into existing ML pipelines (see [Figure 1](#S1.F1)). We propose extending existing pre-trained point prediction models to additionally estimate parameters of the spread of a given probability distribution. In this work, we predict parameters of a Gaussian distribution similar to a classic heteroscedastic regression model. We approximate the error distribution by a Gaussian as its variance parameter can be interpreted as a measure of the aleatoric uncertainty of the model. We can then use any state-of-the-art explainability method to explain the variance estimate provided by this probabilistic model. By highlighting the features contributing to the variance output, we can identify the input features that contribute to the model’s uncertainty. While we focus on heteroscedastic regression in this work, the approach can be generalized to other target spaces, for example, for classification. To evaluate our method, we construct a synthetic problem where we can control how data is generated and how features contribute to the aleatoric uncertainty present in the data. We train a simple feed-forward neural network on the regression task and test how many injected uncertainty-driving features we can recover. We compare our method to the established literature baseline Counterfactual Latent Uncertainty Explanations (CLUE) (Antoran et al., ). Furthermore, we illustrate a potential application of our method to an age regression task on images. We extend a pre-trained state-of-the-art transformer model and analyze the feature-attribution maps obtained for the uncertainty estimate.

### 1.1 Related Work

Uncertainty quantification and explainability are rich areas of research (Abdar et al., ; Vilone & Longo, ).
Yet, few researchers have recognized the importance of explaining the sources of uncertainty in predictions.
Yang & Li () have developed an explainable uncertainty quantification approach for the prediction of molecular properties. They employ message-passing neural networks and generate unique uncertainty distributions for each atom of a molecule. However, this approach is inherently specialized for graph-based representations of molecules.
Counterfactual Latent Uncertainty Explanations (CLUE) (Antoran et al., ) and related approaches (Perez et al., ; Ley et al., ) derive counterfactual explanations by optimizing for an adversarial input that is close to the original input but minimizes uncertainty. The adversarial input is constrained to the data manifold with a deep generative model of the input data to prevent out-of-distribution explanations. This necessitates an optimization process for each explanation and the training of an auxiliary generative model, rendering CLUE and its extensions computationally demanding and difficult to implement. While CLUE is applicable to other probabilistic ML methods, the authors focus on explaining BNNs.

Wang et al. () have developed a gradient-based uncertainty attribution method for image classification with BNNs. They modify the backpropagation to attain complete, non-negative pixel attribution and to prevent vanishing gradient issues.
To develop a method for detecting deterioration in model performance and explaining its cause, Mougan & Nielsen () use traditional ML methods and bootstrapping to obtain estimates of epistemic uncertainty. To explain the sources of uncertainty, they train a model and obtain uncertainty estimates on a test set transformed with an artificial distribution shift. In a second step, they train another model that predicts the uncertainty estimates obtained in the first step. Subsequently, Shapley values are estimated for the second model to extract the feature attributions for the uncertainty. Mehdiyev et al. () also focus on traditional ML methods and employ quantile regression forests in operations research to obtain prediction intervals that quantify uncertainty. They extract feature attributions for the uncertainty by estimating Shapley values directly for these prediction intervals as output.
Watson et al. () take an information-theoretic approach and explain uncertainty by extending Shapley values to the predictive distribution’s higher-order moments in a conditional entropy setting for classification.

## 2 Methods

### 2.1 Deep Heteroscedastic Regression

We want to explain uncertainties in the context of deep distributional neural networks that go beyond mere point predictions, i.e., they estimate the parameters of a specified output distribution. For this work, we focus on deep heteroscedastic regression with a Gaussian output, where we capture the mean and variance of the target, thereby directly modeling input dependence of the output noise.
Here, we consider a regression setting with $n$ independent training examples $\{({\bm{x}}_{i},y_{i})\}_{i=1}^{n}$ with input feature vector ${\bm{x}}_{i}\in\mathbb{R}^{k}$ and targets $y_{i}\in\mathbb{R}$, $i=1,\ldots,n$.
Instead of providing a full picture of the conditional distribution of the target, deep regression models usually only estimate its conditional mean by optimizing the mean squared error (MSE) or comparable loss functions.
As a result, uncertainty information, which can be essential for decision-making and risk assessment, is not captured (Li et al., ).
In contrast, we assume a heteroscedastic Gaussian as the conditional distribution

$$ $y\mid{{\bm{x}}}\sim\mathcal{N}\left(\mu_{{\bm{x}}},\sigma^{2}_{{\bm{x}}}\right)$ (1) $$

and represent its mean $\mu_{{\bm{x}}}$ and variance $\sigma^{2}_{{\bm{x}}}$ using a neural network $f_{\theta}:\mathbb{R}^{k}\to\mathbb{R}\times\mathbb{R}^{+}$ with weights ${\bm{\theta}}$. $k$ is the dimension of the input, and two output neurons produce the mean and variance estimates

$$ $f_{\bm{\theta}}\left({{\bm{x}}}\right)=(\hat{\mu}_{{\bm{x}}},\hat{\sigma}^{2}_{{\bm{x}}})\text{,}$ (2) $$

respectively. We can then optimize the Gaussian negative log-likelihood

$$ $\mathcal{L}\propto\sum_{i=1}^{n}\left(\log(\hat{\sigma}^{2}_{{\bm{x}}_{i}})+\frac{(y_{i}-\hat{\mu}_{{\bm{x}}_{i}})^{2}}{\hat{\sigma}^{2}_{{\bm{x}}_{i}}}\right)$ (3) $$

and interpret the predicted variance as a measure of the aleatoric uncertainty of the model.
However, naively optimizing this criterion with overparametrized models such as deep neural networks can be unstable. These models frequently overfit by excessively shrinking the variance estimate or underfit by predicting only a mean estimate of the target and fitting the variance to the overall target variation in the data (Kuprashevich & Tolstykh, ; Wong-Toi et al., ; Nix & Weigend, ; Seitzer et al., ). In practice, these convergence difficulties can be mitigated by initially training the model using solely the MSE $\sum_{i=1}^{n}\left(y_{i}-\hat{\mu}_{{\bm{x}}_{i}}\right)^{2}$ and subsequently switching to the Gaussian loss as in [Equation 3](#S2.E3) (Sluijterman et al., ). Additionally, instead of applying separate networks, mean and variance are estimated with shared hidden representations (Stirn et al., ).

### 2.2 Extending Pre-trained Models to Output Uncertainty Estimations

The two-stage training process with mean warmup, as described above, naturally fits into the framework of transfer learning: The MSE-based initial training can be seen as a pre-training phase. The model is subsequently fine-tuned, switching to the Gaussian negative log-likelihood loss to capture predictive uncertainty ([Equation 3](#S2.E3)).
The variance estimate is constrained to be positive with an exponential transformation and trained alongside the mean estimate using the Gaussian negative log-likelihood. In a scenario where pre-training with the MSE-loss and fine-tuning with the Gaussian negative log-likelihood is conducted jointly, it is possible to directly construct the neural network with the additional output neuron and exclude it when optimizing the MSE. However, in practice, it might be interesting to extend existing pre-trained models to capture uncertainty, where, for example, the model size and associated training costs make full re-training unfeasible. In this case, we can readily extend such pre-trained regression models by concatenating a column of randomly initialized weights to the weight matrix of the output layer to attain a variance estimate, i.e., an additional output neuron.
We define the parameter matrix connecting the activations of the last hidden layer to the output as:

$$ ${\bm{W}}=\begin{bmatrix}w_{1}\\ \vdots\\ w_{d}\end{bmatrix}\text{.}$ (4) $$

Therefore, the final scalar output of a point prediction neural network can be calculated as:

$$ ${\bm{x}}{\bm{W}}+b=\hat{y}=\hat{\mu}_{{\bm{x}}}\text{,}$ (5) $$

where ${\bm{x}}=\begin{bmatrix}x_{1}&\cdots&x_{d}\end{bmatrix}$ is the activation vector of the last hidden layer and $b$ is the output layer’s bias term. To adjust the network to produce both parameters of a Gaussian distribution $(\mu_{{\bm{x}}},\sigma^{2}_{{\bm{x}}})$ and not just a point estimate $(\mu_{{\bm{x}}})$, we inject a parameter column in the last layer’s weight matrix that will result in an additional output being produced by the neural network:

$$ ${\bm{x}}\times\begin{bmatrix}&w^{*}_{1}\\ {\bm{W}}&\vdots\\ &w^{*}_{d}\end{bmatrix}+\begin{bmatrix}b&b^{*}\end{bmatrix}=\begin{bmatrix}\hat{y}_{1}&\hat{y}_{2}\end{bmatrix}=\begin{bmatrix}\hat{\mu}_{{\bm{x}}}&\hat{\sigma}^{2}_{{\bm{x}}}\end{bmatrix}\text{,}$ (6) $$

where ${\bm{w}}^{*}$ is the randomly initialized new weight vector. Additionally, the bias term needs to be extended by $b^{*}$ to account for the increased number of output neurons. Alternatively, it is equally straightforward to achieve an additional variance output neuron by adding a completely separate multi-layer variance head. We enforce the positivity of the variance with a suitable activation, such as the exponential function.

### 2.3 Post-hoc Explanation of Predictive Variance

Classical explainability methods focus on explaining the predicted class in a classification scenario or the point prediction in a regression task. We similarly explain the variance output in a heteroscedastic regression model. In these models, variance is an additional output to which we can apply any existing, appropriate explainability method. Variance feature attribution should not be understood as a new feature attribution method but as an extension of the scope on which such methods are applied. The choice of a suitable method follows the same criteria as for the explanation of the mean prediction or point predictions, and relevant literature should be consulted (Vilone & Longo, ; Samek et al., ; Schwalbe & Finzel, ).
In principle, we can formulate our uncertainty explanation for any parametrized output distribution for which an explicit formulation of the uncertainty is available. In the case of a Gaussian output distribution, the application is most direct since the variance, as a parameter of the Gaussian, is an explicit output of the neural network. Furthermore, unlike distributions such as the Poisson or the exponential distribution, the variance is uncoupled from the mean output so that factors influencing the mean output are disentangled from uncertainty drivers.

In this work, we employ a model-agnostic and a model-specific explainability method to explain uncertainty.
Model-specific methods are limited in the type of models that they can explain but may offer advantages such as lower computational complexity. In contrast, model-agnostic methods can be applied to any machine learning model (Adadi & Berrada, ).
For our synthetic evaluation, we rely on KernelSHAP (Lundberg & Lee, ), a model-agnostic, local explainability method. KernelSHAP approximates Shapley values using a weighted linear surrogate model with an appropriate weighting kernel.
For the age detection experiment, we use the CAM-based approach HiResCAM (Zhou et al., ; Draelos & Carin, ), a model-specific, local explainability technique that extracts an explanation by weighting the last feature map with the gradient of the output with respect to this last layer feature map. It was originally developed for convolutional neural networks and classification but can be applied to vision transformers and regression. We follow Chefer et al. () who apply a similar CAM-based approach to the $[CLS]$ token of the last attention layer of the transformer model. Their approach can be readily adapted for transformers with class attention.

### 2.4 Synthetic Benchmark Pipeline

Evaluating explainability methods on real-world data is challenging due to the subjective nature of interpreting explanations based on expert prior knowledge. To address this, we employ synthetic data with a known data-generating process. Thereby, we can introduce controlled sources of heteroscedastic, aleatoric uncertainty, which we subsequently aim to detect through our method, variance feature attribution.
Specifically, we sample a synthetic ground truth using a linear system

$$ ${\bm{\mu}}={\bm{V}}\bm{\beta}$ (7) $$

with a design matrix ${\bm{V}}\in\mathbb{R}^{n\times p}$ with entries ${\bm{V}}_{ij}\sim\mathcal{N}(0,1)$, and ground truth coefficients $\bm{\beta}\in\mathbb{R}^{p}$ with $\bm{\beta}_{i}\overset{\text{iid}}{\sim}\text{Uniform}([-1,1])$.
We introduce heteroscedastic noise sources with an absolute-value transformed polynomial model for the heteroscedastic noise standard deviation:

$$ $\bm{\sigma}=\mid\phi(\bm{U})\bm{\gamma}+\bm{\delta}\mid,$ (8) $$

whereby $\bm{U}\in\mathbb{R}^{n\times{p^{\prime}}}$ is a design matrix with entries $\bm{U}_{ij}\overset{\text{iid}}{\sim}\mathcal{N}(0,1)$, $\phi(u_{1},u_{2},\ldots,u_{p^{\prime}})\to(u_{1},\ldots,u_{p^{\prime}},u_{1}^{2},u_{1}u_{2},\ldots,u_{p^{\prime}}^{2})$ is a second degree polynomial feature map, and $\bm{\delta}\sim\mathcal{N}\left(\bm{0},\sigma^{2}_{\delta}{\bm{I}}\right)$ is the uncertainty model error. $\bm{\gamma}\in\mathbb{R}^{\binom{p^{\prime}+2}{2}-1}$ are ground truth noise coefficients with entries sampled from $\bm{\gamma}_{i}\sim\text{Uniform}([-1,-0.5]\cup[0.5,1])$ to avoid negligible effects.

We can then sample the target ${\bm{y}}\in\mathbb{R}^{n}$ with

$$ ${\bm{y}}\sim\mathcal{N}\left({\bm{\mu}},\alpha\cdot\text{diag}(\bm{\sigma}^{2})+\sigma_{\epsilon}^{2}{\bm{I}}\right).$ (9) $$

$\alpha\in\mathbb{R}^{+}$ determines the overall strength of the heteroscedastic uncertainty and $\sigma_{\epsilon}^{2}\in\mathbb{R}^{n}$ regulates the homoscedastic noise.

For our experiments, we set $\alpha=2.0$, $\sigma_{\epsilon}^{2}=0.02$, and $\sigma^{2}_{\delta}=0.05$ to get non-negligible, feature-dependent noise that can be modeled.
We choose $p=70$ and $p^{\prime}=5$ so that the uncertainty sources have to be detected among a larger set of features that do not influence the uncertainty. We sample $n=41,500$ data points and concatenate both design matrices to attain the input ${\bm{X}}_{(n\times 75)}=\left[{\bm{U}}_{(n\times 5)},{\bm{V}}_{(n\times 70)}\right]$ which we split into a train set, validation set, and test set of 32,000, 8,000, and 1,500 instances, respectively. We fit a deep neural network that has four hidden layers with 64, 64, 64, and 32 units. The network has two output neurons for the mean and variance prediction. During training, we apply dropout regularization with a dropout probability of 0.1 to the first two layers. We use the Adam optimizer and a batch size of 64. We commence the training using the MSE and subsequently optimize the model using the Gaussian negative log-likelihood as the loss function. This is akin to acquiring a pre-trained model and then refining it through fine-tuning. In both training phases, we stop training when improvement on the validation set ceases and choose the model weights that resulted in the smallest validation loss.

To interpret the uncertainties, we apply the post-hoc feature attribution method KernelSHAP from the SHAP package (Lundberg & Lee, ). We attain a feature importance measure as the mean absolute estimated Shapley values over the test dataset. We can then scrutinize if the noise sources $u_{i}$ are rediscovered as the most important features for the uncertainty estimation.

We compare our approach to Counterfactual Latent Uncertainty Explanations (CLUE), for which we have to train a variational autoencoder on the train data and apply the optimization as detailed in (Antoran et al., ). We calculate the CLUE feature importances as the mean of the absolute differences between each latent explanation to each original input over the test dataset.

### 2.5 Application to Age Detection

To showcase the applicability of our approach in a non-synthetic setting, we turn to the problem of age detection. Age detection is a relevant task in computer vision and finds application in a wide range of areas, from security to retail. We apply MiVOLO (Kuprashevich & Tolstykh, ), a recent state-of-the-art transformer-based model that achieves best-of-its-class performance in multiple age detection benchmarks (Kuprashevich & Tolstykh, ; Lin et al., ; Zhang et al., ). MiVOLO was designed explicitly for age detection, alongside the related problem of gender detection, and tackles both problems simultaneously to leverage synergies between the tasks. MiVOLO builds on VOLO (Yuan et al., ), a transformer architecture based on outlook attention developed specifically for computer vision tasks. In contrast to VOLO, MiVOLO has two separate inputs. The first is the image patch with a person’s face, and the second is the image patch of the body with the face removed. However, for simplicity, we use a version of the model that only uses the face input, bringing it closer to the original VOLO. The model produces a single three-dimensional output tensor, representing the prediction of the male and female sexes and the person’s age.

We use a pre-trained version of MiVOLO and, following our procedure introduced in [Section 2.2](#S2.SS2), extend the parameter matrices of the MiVOLO head, auxiliary head, and their respective bias terms. We initialize the extension of the parameter matrices using a Gaussian distribution following Glorot & Bengio () and set the new bias terms to zero. The new dimension of the output vector is four, including the age prediction variance. After we have equipped the model with a Gaussian distribution, we train it using the IMBD-clean dataset (Lin et al., ), using the annotations and pre-processing by Kuprashevich & Tolstykh (). We fine-tune the model with the Gaussian negative log-likelihood as our loss function. Consequently, we turn the previous point prediction for age into the mean of our estimated Gaussian distribution. The loss term for the gender detection remains a BinaryCrossEntropy loss. We use an Adam optimizer with a learning rate of $1e\text{-}5$, a weight decay of $1e\text{-}2$, and a batch size of 176. We use the validation set only for model selection. To detect and visualize the drivers of uncertainty in the images, we use HiResCAM as described in [Section 2.3](#S2.SS3).

## 3 Results

### 3.1 Detection of Noise Sources in Synthetic Datasets

We examine the capability of our method to identify the drivers of uncertainty, which are features that correlate with the magnitude of the heteroscedastic noise. For the synthetic dataset, we know the data-generating process and, therefore, the ground truth noise sources and aim to rediscover them using variance feature attribution.
An analysis of the quality of the uncertainty estimate can be found in the [Section A.1](#A1.SS1). The results indicate that the uncertainty estimates generated by our model are meaningful; however, for them to be used to guide decision-making in practical applications, we argue that they should be explainable and ascertain if our method can detect the features influencing the heteroscedastic noise in the data.
We use variance feature attribution with estimated Shapley values to delineate each feature’s contribution to the model’s uncertainty on 200 random test examples. The SHAP summary, depicted in [Figure 2](#S3.F2), offers an aggregated view of feature impacts, including their directionality. The five ground truth noise features are accurately identified as the top five contributors to model uncertainty ([Figure 2](#S3.F2) A).

Figure: Figure 2: Explanations for uncertainty and mean predictions for the synthetic dataset. We display SHAP summaries for the 10 most important features of model uncertainty (A) or mean prediction (B) ordered by the mean of their absolute estimated Shapley values. We employ a neural network trained on synthetic data with heteroscedastic noise that non-linearly depends on 5 of 75 input features. By interpreting the variance output, we identify the key factors driving the aleatoric uncertainty of the model. The attribution of the mean output offers complementary information but disregards uncertainty features.
Refer to caption: /html/2312.07252/assets/x2.png

In practice, explaining uncertainties is particularly relevant in instances with very high or very low uncertainty. In addition to random instances, we, therefore, apply variance feature attribution to the 200 highest and lowest uncertainty instances of the 1,500 test samples. We depict the resulting feature importance, measured as mean absolute estimated Shapley values, in [Figure 3](#S3.F3) A, C, E.
For comparison, we explain the same instances for the same model using CLUE (Antoran et al., ) as an alternative explainer of model uncertainty (see [Figure 3](#S3.F3) B, D, F). We calculate CLUE feature importance as the mean absolute difference between CLUE uncertainty counterfactuals and the original sample. We find that variance feature attribution and CLUE both effectively identify uncertainty drivers for high-uncertainty instances. Variance feature attribution also delivers stable performance for random and low uncertainty examples while CLUE’s detection capability deteriorates. This suggests that, unlike CLUE, variance feature attribution can elucidate the factors contributing to certainty in addition to those relevant to uncertainty. To compare the runtimes, we use an Intel(R) Core(TM) i5-12600K machine with an NVIDIA RTX 3060 Ti. We apply the post-hoc explanation of both methods to the same trained model 15 times. A run takes $81\pm 34$ seconds for variance feature attribution and $1012\pm 33$ seconds for CLUE. Notably, the uncertainty estimation training stage only marginally reduces the accuracy of the regression model in our experiment. The MSE is 0.79 after the MSE pre-training stage and 0.80 after the uncertainty-aware fine-tuning stage. Naturally, the ability of variance feature attribution to detect drivers of uncertainty depends on the train dataset size and magnitude of the noise signal. We analyze this in [Section A.2](#A1.SS2).

Figure: Figure 3: Top 20 variance feature attribution importances and CLUE feature importances on a synthetic dataset. A, B: From a test set of 1,500 samples with 70 mean and 5 noise features, we explain the 200 instances with the highest predicted uncertainty. Our variance attribution method (A) and CLUE (B) both faithfully assign high importance to the ground truth uncertainty sources. C, D: For 200 random instances, variance attribution delivers accurate explanations (C), while CLUE is unreliable (D). E, F: Variance feature attribution also can attribute (un-)certainty for the 200 lowest uncertainty predictions (E), while CLUE performance deteriorates (F).
Refer to caption: /html/2312.07252/assets/x3.png

### 3.2 Finding Potential Drivers of Uncertainty in Age Detection

We investigate the application of variance feature attribution to a non-synthetic age detection task using the MiVOLO model (Kuprashevich & Tolstykh, ) and the IMDB-clean dataset (Lin et al., ). Similar to our synthetic experiment, we also evaluate the quality of the uncertainty estimate, and the results can be found in the [Section A.3](#A1.SS3). The results similarly suggest that the uncertainty estimates are relevant.

Applying our method to the MiVOLO model using HiResCAM as the explanation method reveals reasonable potential explanations for the predictive uncertainty (see [Figure 4](#S3.F4)). The explanations mainly focus on areas around the eyes, mouth, nose, and forehead. These areas seem to be highlighted especially strongly when the person in the image shows emotions such as joy and anger that lead to distortions of these facial areas that the model might confuse with age-induced wrinkles. Similar highlights in the explanations are present throughout a majority of images in the test dataset.

Figure: Figure 4: Examples of uncertainty explanations for images from MiVOLO’s (Kuprashevich & Tolstykh, ) IMDB-clean-based (Lin et al., ) test set. Variance feature attribution using HiResCAM (Draelos & Carin, ) mainly highlights facial distortions such as laugh lines and frowning. The corresponding input images are depicted in [Section A.3](#A1.SS3).
Refer to caption: /html/2312.07252/assets/mivolo_data_dataset_images_00_nm0085400_rm708744960_1948-8-30_2008.jpg.png

## 4 Discussion and Conclusion

In this work, we have proposed variance feature attribution, a simple strategy for explaining predictive aleatoric uncertainties. This approach requires minimal modifications to existing neural network regressors, making it a practical solution for broad adoption. We use neural networks with Gaussian output distribution to estimate uncertainty and apply appropriate explanation methods to the variance output. Thereby, we are able to provide explanations for the sources of uncertainty estimates. In synthetic experiments, the resulting explanations can be generated faster while matching CLUE baseline performance or outperforming it in situations with low and medium uncertainty. Further, estimating the uncertainty does not considerably impede prediction accuracy. Generally, the Gaussian assumption may be unsuitable in some applications. In principle, we can extend this framework to predict parameters of various other distributions, such as more heavy-tailed or mixture distributions, providing the opportunity for explainable uncertainties tailored to specific problem domains.
Since conventional evaluation metrics are not directly applicable, we have introduced an evaluation protocol designed to assess uncertainty explainers involving the use of synthetic data with a known ground-truth noise profile. We aim to expand this approach using more diverse data-generating processes or by fusing real data with synthetic noise sources.
Generally, we inherit practical limitations of uncertainty estimation, such as overconfidence, as well as challenges related to explainability, such as issues with faithfulness and consistency. Nonetheless, significant advancements are occurring in both of these domains, which can be seamlessly integrated into our method. Future work might involve the study of synergies of the explanations of point and uncertainty predictions. For example, in the context of explainable active learning (Ghai et al., ), a shared visualization of both explainability modes could be beneficial.

#### Reproducibility Statement

Our research findings can be reproduced using the code and data available at [https://anonymous.4open.science/r/vfa-variance-feature-attribution-0E63](https://anonymous.4open.science/r/vfa-variance-feature-attribution-0E63), including scripts to create the figures and information on the requirements. For the synthetic experiment, this includes the pipeline that can be used to create additional datasets to probe the method further. Details on how to run the synthetic experiments and how to download the age detection data and reproduce our explanations are in the repository’s README.

## Appendix A Appendix

### A.1 Uncertainty Evaluation in the Synthetic Experiment

Figure: Figure 5: (A) Uncertainty calibration is illustrated by binning predictions into 20 probability level intervals and plotting the expected against the empirical frequency of observing the true target in each interval. The dashed line indicates theoretical uncertainties with perfect calibration (Kuleshov et al., ). (B) The MSE of the trained deep heteroscedastic regression model. The test set is iteratively limited to the quantile of the test data with the lowest predicted uncertainty (red), to instances with the lowest ground truth standard deviation of the heteroscedastic noise (green), or as a baseline based on instances with the lowest distance from the mean (blue).
Refer to caption: /html/2312.07252/assets/x4.png

We assess the model’s ability to learn meaningful uncertainties by adapting a calibration curve for regression uncertainties as described by Kuleshov et al. ().
Precisely, we calculate the cumulative probability of the observations with respect to the predicted distribution

$$ $F_{i}=CDF_{Gaussian}\left(\frac{(y_{i}-\hat{\mu}(x_{i}))}{\sqrt{\hat{\sigma}^{2}(x_{i})}}\right)$ (10) $$

and calculate the empirical cumulative probability distribution at 20 probability levels $\tilde{p}_{j}=0.05j;j=1,2,...,20$:

$$ $\hat{p}_{j}=\frac{\mid\{y_{i}\mid F_{i}<\tilde{p}_{j},i=0,1,2,...,n\}\mid}{n}.$ (11) $$

In [Figure 5](#A1.F5) (A), we plot the empirical frequency $\hat{p}_{j}$ against the probability levels $\tilde{p}_{j}$ and expect $\hat{p}_{j}\to\tilde{p}_{j}$ when $n\to\infty$ for perfectly calibrated models. The uncertainty estimates are overall well calibrated.

We further look at the reduction in test MSE we can achieve by limiting the test set to quantiles of the test set with the lowest uncertainty. This mimics a scenario where users would opt out of utilizing the prediction for high-uncertainty instances. As seen in [Figure 5](#A1.F5) (B), limiting the test set to low uncertainty predictions starkly reduces the test error. We compare this to a baseline where instances are included based on them having a low absolute distance from the mean prediction. The reduction is close to the theoretical optimum, where instances are removed based on the ground truth standard deviation of the heteroscedastic noise.

### A.2 Uncertainty Driver Detection: Noise Magnitude and Dataset Size

Like in every supervised machine learning task, a model’s ability to learn and detect the signal, which in this case is the noise characteristic, depends on the signal-to-noise ratio and the size of the available dataset. We investigate the precision of noise feature detection across varying dataset sizes and signal strengths. For uncertainty estimation, the scaling factor $\alpha$ of [Equation 9](#S2.E9) used for the heteroscedastic noiser serves as an indicator of signal strength. Fig. [6](#A1.F6) displays the precision of the detection of heteroscedastic noise features for various dataset sizes and noise scalers. We compute the precision as the ratio of noise features that appear among the top five most important features. As the dataset size and the magnitude of the representable uncertainty grow, the capability to identify the origins of uncertainty improves.

Figure: Figure 6: Precision of the detection of heteroscedastic uncertainty features vs. the dataset size and intensity of the heteroscedastic noise effect in the synthetic dataset. We use our proposed method in single trial runs. Small effects can be detected in large datasets.
Refer to caption: /html/2312.07252/assets/heatmap_precision.png

### A.3 Uncertainty Evaluation and Input Images of the Age Detection Experiment

We verify that our model has learned to estimate uncertainty via the variance of the predicted distribution. We use the same approach as in our synthetic experiment (see [Section A.1](#A1.SS1)) and iteratively remove the sample with the highest predicted variance from our testing set and compare the development of the MSE when removing samples based on their predicted age (see [Figure 7](#A1.F7)). We show all exemplary explanations and their corresponding input in [Figure 8](#A1.F8).

Figure: Figure 7: The test set MSE of the age detection MiVOLO model, fine-tuned for uncertainty estimation. The test set is iteratively limited to the quantile of the test data with the lowest predicted uncertainty (red). As a baseline, we remove instances according to the predicted age (blue), which consistently produces a set with higher MSE.
Refer to caption: /html/2312.07252/assets/age_detection_p_plot.png

Figure: Figure 8: Input images and uncertainty explanations for images from MiVOLO’s (Kuprashevich & Tolstykh, ) IMDB-clean-based (Lin et al., ) in our age detection experiment. Variance feature attribution using HiResCAM (Draelos & Carin, ) mainly highlights facial distortions such as laugh lines and frowning.
Refer to caption: /html/2312.07252/assets/input_mivolo_data_dataset_images_00_nm0085400_rm708744960_1948-8-30_2008.jpg.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: - Empirical evaluation is done against only two baselines (CLUE, InfoSHAP); are these the best performing existing methods?
soundness: 3
strengths: - The paper proposes a novel method for explaining predictive aleatoric uncertainty by explaining the variance output in a heteroscedastic regression model. 
- The paper empirically demonstrates that the proposed method VFA-SHAP identifies noise features driving the model's aleatoric uncertainty.
summary: The paper proposes a novel method for explaining predictive aleatoric uncertainty by explaining the variance output in a heteroscedastic regression model. Their method identifies which input features contribute to model uncertainty (i.e., the variance output). Further, they propose benchmarks for evaluation of uncertainty explainers. Lastly, they empirically evaluate their method against existing approaches and demonstrate that the proposed method identifies the noise features driving the model's aleatoric uncertainty.
weaknesses: - A limitation of this method is its computational complexity; can the authors elaborate on the computational cost of their method, relative to existing approaches?
- For the real-world dataset evaluation, the authors show the output explanations of VFA. It is hard to understand the performance of the method on this dataset, since the ground-truth explanations are unknown (which I understand is not available) and only the performance of VFA is presented. It would be useful to show the performance of the other baselines on this real-world task to at least evaluate the performance of this method w.r.t. them.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. Line 286-287: Could you please elaborate on the need for prior knowledge since the metrics were introduced in Section 2.3?

2. For Figure 4, could you please provide the explanation of the prediction as well? The question is how different the decision explanation and the uncertainty explanation maps are.

3. If XGBoost for InfoSHAP was regularized, it might not choose the uncertainty features for Figure 2. Could you please comment on whether the model contains the uncertainty features?
soundness: 2
strengths: The idea of understanding reasons behind the aleatoric uncertainty is interesting. The paper picks a specific controllable case and studies it.
Multiple evaluation metrics are introduced and both synthetic and real-world datasets are considered for evaluation.
In total, the paper compares five different methods, including two baselines.
summary: The paper combines deep heteroscedastic regression with model explainability methods to identify potential reasons for aleatoric uncertainty. A framework is built, where to estimate uncertainty, neural regression networks with Gaussian output distribution are used. Explainability methods such as Kernel SHAP, Integrated Gradients, and Layer-Wise Relevance Propagation are then used to explain variance. Based on the synthetic and real-world datasets, the methods are compared to InfoSHAP and CLUE based on a set of metrics. The SHAP-based variance feature attribution method consistently outperforms other approaches.
weaknesses: 1. The paper introduction is built around understanding aleatoric uncertainty; however, it is not clear from the paper what "understanding uncertainty" is. In general, one would expect that understanding uncertainty means identifying the sources of it. For the case of heteroscedastic regression, the explanations are built by identifying features that contribute to output uncertainty. More discussion is needed on why this method is effective and scalable. The worry is that the methods can be misleading in multiple other cases, for example, when the sources of uncertainty are outputs and not inputs.

2. While applying model explanation methods to estimate uncertainty is an interesting idea, the explainability methods were designed for providing reasons for decision making; therefore, depending on the context, these two problems can be different in nature. More discussion is needed on laying a foundation for why the methods are effective for uncertainty estimation.

3. When the number of uncertainty sources is low (e.g., 5 out of 70) and there is no feature correlation, as in the synthetic dataset, the VFA methods identify the uncertain features. However, it is not clear if this approach will work as well in more realistic scenarios, when the number of uncertainty sources is higher and there are correlations between features.

4. For the synthetic dataset, the ground-truth features are known; for non-synthetic cases, the correlation between the squared residuals and the uncertainty estimates is considered. However, by its nature, this is a post-hoc method that can favor some methods over others. Can you please comment on the reliability of the ground-truth computation?

5. Could you please provide errors in the performance estimates, for example by performing cross-validation over different train and test folds.

Minor:

Line 11: “are two pillars” -> "two of the pillars," as there are many pillars of trustworthy AI that are not mentioned, including interpretability, fairness, robustness, reliability, and so on.

Line 14: Understanding reasons for decision making or uncertainty does not necessarily create trust and can create distrust.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: Are there quantitative metrics or other methods you can compare with for the age detection explainability task? I agree that the explanations focusing on the eyes, mouth, nose, and forehead is promising. 

You present your approach as “straightforward and scalable.” Is there an easy benchmark you can do to demonstrate usability (lower compute) against BNNs? Or even an anecdote of the computational improvement (can cite from elsewhere).

Lakshminarayanan et al., 2017 also proposed ensembles as an alternative non-Bayesian approach for uncertainty quantification. Why is your approach (network learned variance parameter) better?
soundness: 3
strengths: The method is mathematically straightforward and simple to grasp. Figure 1 is a nice conceptual overview. The methods and metrics are well explained (and therefore replicable), and the provided formulas and definitions are appropriately conceptual for understanding the method. In general, the paper organization is very good.

The experimental results, especially on the synthetic data where the noise process is known, are strong compared to the other methods. Results figures are also well-presented; I appreciated the other experimental conditions such as simple vs. complex noise model, or mixed features given in the Appendix.
summary: The authors propose a method for explaining predictive aleatoric uncertainties (uncertainties due to the inherent variability of the training data) in heteroscedastic regression settings (where aleatoric uncertainty is interesting). Their method generates variances as well as point predictions from the regression NN, then explains the variances using existing XAI methods. They compare to CLUE and InfoSHAP, and also provide a benchmark using synthetic data, real data, and synthetically augmented data.
weaknesses: My main concern is about originality: how to quantify a contribution that is simply combining existing concepts from explainability and uncertainty quantification, albeit in an appealingly simple and well-presented way. Is good presentation, fairly complete experiments, and mathematical appeal sufficient to call this a major contribution? I’ll say yes because of the performance improvement in identifying noise features and the usability/accessibility of the approach. 

There may be a few more experiments that would strengthen this contribution--see Questions.

---

confidence: 3
contribution: 1
initial_rating: 3
presentation: 3
questions: How good are the regression predictions of the different methods? I understand this is not the primary goal, but it might be relevant (especially if MSE for VFA-methods is substantially higher than for InfoSHAP/CLUE)
soundness: 2
strengths: - The setup of the method (VFA) and experiments are easy to follow. It is fairly clear what the metrics are measuring and that relevance-accuracy, faithfulness and robustness are desirable properties of an explanation.
- Applying various explainability methods for VFA gives a nice overview of advisable ways to do VFA.
- The language is consistently appropriate and has little to no grammatical/spelling errors.
- The results on the tabular data are persuasive and show strongly that VFA-based methods outcompete InfoSHAP and CLUE for this kind of data, with this type of noise on these metrics.
summary: The presented work addressed explainable uncertainty. Previous work on this exists, but no comparative analysis has been done. The authors propose a dataset and a set of metrics where they apply RRA and RMA for uncertainty. They do this by inserting heteroscedastic uncertainty on a tabular dataset. The authors show that methods based on fine-tuning a Gaussian NLL and then running explainability on the variance-head (VFA) outperforms two existing methods (CLUE, InfoSHAP). Authors also show some results where they apply VFA to age regression.
weaknesses: - The proposed benchmark only concerns regression with tabular data. I suspect the metrics do not extend well to high dimensional problems such as the age regression example. Additionally, the method (VFA) may not extend very well to classification (though formulations exist, they have limitations), which gives a limitation to the scope of the findings. From your results the conclusion of the comparative analysis is limited to regression on tabular data (where faithfulness is specifically even only for continuous features).
Since the main contribution is the comparative analysis / benchmark, the choices for how the dataset is made and how the metrics work are important, but these are not sufficiently elaborated. Specifically, it needs to be very clear why the heteroscedastic noise is introduced in these ways (1-S, 50-C) and not in some other ways, and how that might impact the results. The metrics are sufficiently understandable, but a more thorough explanation of their properties and potential limitations will make them more persuasive. The choice of datasets used should also be further elaborated (why these datasets and not others / more?).
- The results from the Faithfulness analysis are not in the main body of the paper. Particularly because these results conflict with other results, they should be in the main body. 
- The results on face-age regression are not particularly informative. The established criteria (Relevance Rank, Faithfulness and Robustness) are not applied, and the explanations are not comparative between methods and not showing a difference between mean-variance explanations. Additionally, since the “reasonableness” of explanation is determined post-hoc, many possible results could be considered reasonable. If the model would’ve identified uncertainty at the glasses/hair we might conclude that it’s finding occlusion. The suggestion that the highlighted areas relate to emotion are unfounded. High uncertainty is shown in the 3rd, 5th and last face, even though they do not show substantial emotion (other than smiling for a picture), while the second face has clear emotion but a relatively low uncertainty. I’d find the findings from 3.3 as they are uninformative. I’d consider removing them, or expanding substantially by creating a comparative analysis similar to what is done with the tabular data. 
- While the authors give some reasoning that aleatoric uncertainty explanations are important (Lines 60-65), the reasoning is not particularly strong or supported by some evidence. There is no evidence that this might be useful in a decision support system for domain experts. The example of minority status presenting a bias in the model seems shaky. A common problem where I might expect this is that facial recognition works poorer on black faces (so high uncertainty attribution to African facial features, or all of the skin tone), but this should be due to epistemic uncertainty not aleatoric uncertainty.  

## Additional Feedback (minor comments)

- It should be considered a substantial limitation that the prediction of uncertainty only (theoretically) reflects aleatoric uncertainty and would neglect epistemic uncertainty. This means that some factors that may drive uncertainty in a general sense are ignored. I think this is something that should be discussed in the limitations section of the paper.
- On line 74 “Local or Global” explanations are explained by their purpose, not by how they’re created. A simple definition of local=per sample, global=per dataset would help a reader not familiar with this concept. 
Section 1.1 reads like a list of papers, but does not read as a narrative. After reading the rest it makes sense that this section introduce CLUE and InfoSHAP, but this is not clear when reading the paragraph first. Structuring 1.1 into a narrative (with paragraphs) and extracting the meaningful parts will make it easier to read.
- A visualizations of the synthetic data in 2.4.1 and even the synthetic noises would greatly improve interpretability, even if simplified in 2 dimensions. This could be put in an appendix. 
- The results from Figure 2 are interesting, but quite hard to interpret. It’s not clear what is “good” behavior for the SHAP values. (large spread, or positive values?), and the plots in C are missing x/y axes labels and a legend (red=uncertainty related feature). The corresponding a-o letters with long caption are also difficult to connect. Perhaps C can be explained by rows=high/mid/low uncertainty samples, columns=Explainable Uncertainty-method, and then discussing the results without referring to the individual subplots. 
- There’s a typo in bottom right Table 2 in the appendix. 
- The definition for epistemic uncertainty (Lines 39-44) is not entirely accurate. This is not the main point of the paper, but might lead to confusion elsewhere. This definition ignores contributions of model misspecification and approximation error, and does not do good justice to covariate shift. The definition for aleatoric could also be more specific, “randomness in the data” could mean many things. A possible definition for aleatoric uncertainty is stochasticity in the true relationship between X and Y.