Regularizing Explanations in Bayesian Convolutional Neural Networks

TL;DR

The paper tackles overfitting to spurious features and lack of interpretable uncertainty in neural networks by marrying explainable AI with Bayesian inference. It introduces an explanation-regularization approach that integrates explanation feedback into Bayesian CNN training via an activation-based likelihood term, preserving $ELBO$ optimization while guiding the model to focus on relevant features. Empirical results on four datasets show improved predictive performance when spurious cues are present or uncertainty is high, and sharper, more localized explanations compared to data augmentation alone. The work advances practical, uncertainty-aware models with correct explanations, though it relies on human-provided explanation feedback and invites future work on adaptive, scalable feedback collection.

Abstract

Neural networks are powerful function approximators with tremendous potential in learning complex distributions. However, they are prone to overfitting on spurious patterns. Bayesian inference provides a principled way to regularize neural networks and give well-calibrated uncertainty estimates. It allows us to specify prior knowledge on weights. However, specifying domain knowledge via distributions over weights is infeasible. Furthermore, it is unable to correct models when they focus on spurious or irrelevant features. New methods within explainable artificial intelligence allow us to regularize explanations in the form of feature importance to add domain knowledge and correct the models' focus. Nevertheless, they are incompatible with Bayesian neural networks, as they require us to modify the loss function. We propose a new explanation regularization method that is compatible with Bayesian inference. Consequently, we can quantify uncertainty and, at the same time, have correct explanations. We test our method using four different datasets. The results show that our method improves predictive performance when models overfit on spurious features or are uncertain of which features to focus on. Moreover, our method performs better than augmenting training data with samples where spurious features are removed through masking. We provide code, data, trained weights, and hyperparameters.

Figures (3)

• Figure 1: Method Overview. a) During training, a nn gets an input sample $\mathsf{X}_i\in\mathbb{R}^{(w\times h\times c)}$ from the training dataset and tries to match the prediction $\hat{y}_i$ with the ground truth label $y_i$. Our method provides the nn with additional evidence in the form of explanation feedback $\mathbf{E}_i\in\{0,1\}^{(w\times h)}$. A value of $1$ in $\mathbf{E}_i$ indicates a region in the input space as irrelevant to the prediction, while $0$ indicates that we do not have any concern. The explanation feedback is used to regularize the model's focus to give correct explanation and add domain knowledge. b) A new input sample $\mathsf{X}_j$ from the test dataset is fed to the model and an explanation is generated. Without explanation regularization, the nn uses the patch to make the prediction. With our method, the nn no longer looks at the patch in the image. The skin images are from the isic dataset Codella_SkinLesionAnalysis_2019Tschandl_HAM10000DatasetLarge_2018Rieger_InterpretationsareUseful_2020.
• Figure 2: Finding Activations. Given an explanation feedback $\mathbf{E}_i\in\{0,1\}^{(w\times h)}$ for the sample $\mathsf{X}_i\in\mathbb{R}^{(w\times h\times c)}$, we find activations to add the explanation feedback. A value of $1$ in $\mathbf{E}_i$ indicates irrelevant regions in the input. A value of $0$ denotes features that no preference is given. First, $\mathbf{E}_i$ is downsized to the size of feature maps of the last convolutional layer using the function $f(\cdot)$. Afterward, since the height and widths are the same, we simply overlay the explanation feedback with the feature maps to find activations to target. Specifically, we inject this information via the likelihood term of \ref{['eq:lrt_elbo']}. The skin image is from the isic dataset Codella_SkinLesionAnalysis_2019Tschandl_HAM10000DatasetLarge_2018Rieger_InterpretationsareUseful_2020.
• Figure 3: Examples of saliency maps on samples randomly drawn from the test dataset. More examples can be found in the link given on the first page.