Exploring Counterfactual Alignment Loss towards Human-centered AI

Current deep learning (DL) models, though have reached remarkable performance, are often biased towards human non-explainable features in the decision-making process Najafabadi et al. (2015); Geirhos et al. (2018). Such biases can raise concerns about trustworthiness, impeding these models from being deployed in safety-critic domains such as healthcare. To address this issue, apart from prediction accuracy, it is desirable for DL models to improve their alignment with human decision-making.

To address this issue, many studies attempted to learn human explainable features for prediction. Examples include Liu et al. (2021b); Wang et al. (2022) that incorporated human knowledge with specialized neural network architectures; and Zheng et al. (2017); Sreedevi et al. (2022) that proposed decision-making systems based on the human cognitive model. In particular, explanation-guided learning approaches Ross et al. (2017); Ismail et al. (2021); Gao et al. (2022); Fei (2022) proposed to align the attention map of DL models to image regions annotated by human experts, thereby enforcing an intrinsically human-centered model. However, the attention map these methods are based on is obtained in a data-driven manner and not necessarily provides a causal attribution for the model’s decision, which largely compromises their validity for alignment.

In this paper, we propose a novel human-centered framework based on counterfactual generation. Specifically, we first leverage the counterfactual generation to identify features that causally attribute to the model’s prediction. Then, to align these features with human annotations, we propose a novel CounterFactual Alignment (CF-Align) loss, which punishes the network once the identified causal factors do not align with human annotations. To optimize the proposed loss that involves counterfactual generation utilizing implicit function forms, we propose an implicit gradient solver rooted in the implicit function theorem. Returning to the lung cancer example, Fig. 1 shows that our method can learn explainable features; as a contrast, features utilized by the baseline DL model are difficult to interpret.

Contributions. To summarize, our contributions are:

1. 1.

   We propose a counterfactual alignment loss for human-centered AI.

2. 2.

   We propose to leverage the implicit function theorem for optimization.

3. 3.

   We achieve more accurate and explainable results than existing methods on real-world data.

Existing works on explainable AI (XAI) can be classified in various ways Speith (2022). One such classification distinguishes between human-centered and technology-centered approaches, with the former focusing on the needs of specific end-users and the latter on those of AI professionals. Another classification is based on whether the method is ante-hoc or post-hoc explainable, which respectively refers to instinct explainability and explanation after the model has been trained.

Our method belongs to human-centered, ante-hoc XAI. Therefore, we will provide a detailed discussion of these two branches of works below, leaving the comprehensive review of other XAI methods to Arrieta et al. (2020) and Ali et al. (2023).

Human-centered XAI aims to develop AI models that behave in a human-understandable manner. To achieve this goal, existing methods have proposed incorporating various human priors, such as knowledge Liu et al. (2021b); Wang et al. (2022), memory Fu et al. (2014), and cognitive models (Zheng et al. (2017); Sreedevi et al. (2022), into the design of AI models. By doing so, their models could behave in a certain way that is similar to humans, thereby gaining their trust. However, most of these methods require building model-specific architectures, which makes them difficult to adapt to different settings. In contrast, we propose a model-agnostic loss for human-centered learning, which can be easily adapted to various task settings and neural network architectures.

Ante-hoc XAI refers to the design of AI models that are transparent and intrinsically explainable. Examples of such models include white-box models, such as linear regression and decision trees; hybrid models that combine a white-box model with a black-box model for improved performance Nauta et al. (2021); Wang & Lin (2021); joint prediction-explanation models that were trained to provide both predictions and explanations Hind et al. (2019); Rieger et al. (2020); and methods that achieved explanation through architecture adjustment Zhang et al. (2018); Chen et al. (2019).

Of particular relevance to our work are explanation-guided learning approaches Ross et al. (2017); Ismail et al. (2021); Gao et al. (2022); Fei (2022), which proposed to align model attention or gradient with human-annotated areas. Nevertheless, these attention regions might not align with the causal factors influencing the decision-making process, potentially undermining their appropriateness for alignment purposes. In contrast, our method is based on the counterfactual generation that intrinsically corresponds to causal attributions, which thus ensures the alignment of the decision process with that of domain experts.

Counterfactual explanation. Our work is closely related to the counterfactual explanation methods Verma et al. (2020); Balasubramanian et al. (2020); Lang et al. (2021), which aimed at answering what could the outcome have changed to had input to a model had been changed in a particular way. To implement, they proposed to minimize alterations that changed the prediction. Such a modified region can be taken as the causal factor to determine the model’s prediction Parafita & Vitrià (2019). Compared to the attention-based methods, this method can causally attribute the model’s decision Parafita & Vitrià (2019). However, it is important to note that existing counterfactual explanation methods were designed to explain trained black-box models, while our method aims to design an intrinsically explainable model (i.e., ante-hoc explainability).

In this section, we introduce the problem setting, followed by an overview of the counterfactual explanation that our method builds upon.

Problem setup & notations. We consider the classification scenario, where the system includes an image $\boldsymbol{x}\in\mathcal{X}\in\mathbb{R}^{p}$ and a label $y\in\mathcal{Y}$ from an expert annotator. In addition to $y$, we assume the annotator also provides an explanation $\boldsymbol{e}$ that explains his/her decision for each image.

In practice, the explanation mainly refers to the region of central interest. For example, in a radiology report, the radiologist often explains his/her decision by annotating disease-related/lesion areas. Motivated by this, we assume that for each sample $(\boldsymbol{x},y)$, the explanation can be formed as the mask $\boldsymbol{r}\in[0,1]^{p}$, where $r_{i}=1$ meaning that the feature $i$ belongs to the region of interest. In this regard, the training data we collect can be denoted as $\{\boldsymbol{x}_{i},y_{i},\boldsymbol{r}_{i}\}_{i=1}^{n}$. With this data, our objective is to learn a classification model $f_{\boldsymbol{\theta}}:\mathcal{X}\mapsto\mathcal{Y}$ that i) predicts $y$ accurately, and ii) makes its prediction based on the region $\boldsymbol{r}$.

In this section, we propose to align the decision basis of the model with regions provided by human experts. To this end, we propose a regularity term dubbed as the CounterFactual Alignment loss (CF-Align loss). The goal of this loss is to enforce the network’s decision basis (identified with Eq. (1)) to align with the region of interest $\boldsymbol{r}$. In this regard, the classifier will be forced to make predictions based on human-explainable features. The overall pipeline of our method is illustrated in Fig. 2.

Formally speaking, for a prediction model $f_{\boldsymbol{\theta}}$, we define the CF-Align loss as:

where $\boldsymbol{x}^{*}(\boldsymbol{\theta})$ obtained from Eq. 1 is a function of the model parameter $\boldsymbol{\theta}$, $\boldsymbol{r}_{i}$ is a binary mask that represents the region of interest, and $\odot$ denotes the element-wise Hadamard product between matrices.

By combining with this loss to the cross-entropy loss $\mathcal{L}_{cls}(\boldsymbol{\theta})$, the overall training objective for the prediction model $f_{\boldsymbol{\theta}}$ is:

where the hyper-parameter $\alpha$ is the trade-off between classification accuracy and alignment to experts’ annotation boundary. A higher value of $\alpha$ could potentially lead to improved alignment, but it might also compromise prediction accuracy because the model could leverage other features that may not be understandable to domain knowledge for making predictions. Specifically, it was well established that deep learning methods mainly focused on textual features Geirhos et al. (2018). Besides, existing works in medical imaging have also shown that microscopic features (e.g., contour, curvature) Wang et al. (2021) or contextual features Liu et al. (2021a) can improve the prediction power. Moreover, apart from the lack of explainability, such features may fall outside the realm of causal relationships grasped by domain experts and may pertain to spurious correlations. These correlations could potentially lead to non-robustness when faced with perturbations or out-of-distribution scenarios.

Optimization. To optimize Eq. (3), we need to compute the gradient $\nabla_{\boldsymbol{\theta}}\mathcal{L}_{a}(\boldsymbol{\theta})$, which involves the term $\nabla_{\boldsymbol{\theta}}\boldsymbol{x}^{*}(\boldsymbol{\theta})$. The main challenge lies in that the functional form of $\boldsymbol{x}^{*}(\boldsymbol{\theta})$ is not explicit, making it hard to derive the gradient using chain rules.

To compute this implicit gradient, we employ the Implicit Function Theorem (IFT). The idea is to note the fact that $\nabla_{\boldsymbol{x}}\arg\min_{\boldsymbol{x}}g(\boldsymbol{x},\boldsymbol{\theta})\equiv\boldsymbol{0}$ for any differential function $g$, which allows us to write $\nabla_{\boldsymbol{\theta}}\boldsymbol{x}^{*}(\boldsymbol{\theta})$ as the product of the inverse Hessian $H_{g}[\boldsymbol{x}]^{-1}$ and the mixed derivative $\nabla_{\boldsymbol{\theta}}(\nabla_{\boldsymbol{x}}g)$. Formally, we have:

By setting $g(\boldsymbol{x},\boldsymbol{\theta}):=\mathrm{CE}(f_{\boldsymbol{\theta}}(\boldsymbol{x}),y^{*})$, we can compute the implicit gradient $\nabla_{\boldsymbol{\theta}}\boldsymbol{x}^{*}(\boldsymbol{\theta})$ and thereby the gradient $\nabla_{\boldsymbol{\theta}}\mathcal{L}_{a}$ with the chain rule. Besides, it is important to note that our proposed CF-align loss is model-agnostic. Therefore, it can be easily plugged into the training of various neural networks for human-centered learning.

Extension to attributes $\boldsymbol{a}$. In addition to the class label $y$, Eq. 2 also applies when the annotated regions $\boldsymbol{r}$ explains its decision to attributes annotations $\boldsymbol{a}$, which will subsequently affect the label $y$. For instance, in the context of lung cancer diagnosis, $\boldsymbol{a}$ pertains to clinical attributes related to nodules, which include factors such as size, margin, spiculation, and more. During the diagnostic process, clinicians commonly annotate the nodule’s region and provide assessments of these attributes, which serve as the basis for their diagnostic decisions.

If these attributes are provided during training, then we can leverage them into our framework. In this case, the overall classifier can be written as $g_{\boldsymbol{\gamma}}\circ f_{\theta}$, where $f_{\boldsymbol{\theta}}:\mathcal{X}\mapsto\mathcal{A}$ and $g_{\boldsymbol{\gamma}}:\mathcal{A}\mapsto\mathcal{Y}$. To determine which attributes to modify, we first employ Zhao et al. (2023) to identify key attributes $\boldsymbol{\bar{a}}$ that have causal influences on $y$, and then generate the counterfactual image by:

To train $(\boldsymbol{\theta},\boldsymbol{\gamma})$, we can optimize the objective function Eq. 3 such that the CE loss is replaced with $\mathcal{L}_{cls}(\boldsymbol{\theta},\boldsymbol{\gamma})$ to account for the classifier $g_{\boldsymbol{\gamma}}$ and that the $\boldsymbol{x}^{*}(\boldsymbol{\theta})$ is generated through Eq. 4. For a new sample $\boldsymbol{x}$ to predict, we first use $f_{\boldsymbol{\theta(x)}}$ to predict attributes $\boldsymbol{a}$ based on explainable features, then predict $y$ based on attributes $\boldsymbol{a}$.

To demonstrate the practicability of enhancing the explainability within our method, we apply our model to the pulmonary nodule benign/malignant classification. It’s important to clarify that our objective is not to attain state-of-the-art performance on this particular task but rather to leverage it as an illustrative example to emphasize our key point.

In this paper, we present a counterfactual-based framework aimed at aligning the model with human-explainable features for prediction. To this end, a counterfactual alignment loss is introduced to ensure that the model only modifies regions within human annotations during counterfactual generation. To optimize this loss, we leverage the implicit function theorem to compute the gradient of the alignment loss, which involves implicit forms in the counterfactual generation process. Our framework’s effectiveness is demonstrated by its ability to guide the prediction model in leveraging nodule-related features for the diagnosis of pulmonary nodule malignancy.