Learning to Receive Help:
Intervention-Aware Concept Embedding Models

The purpose of the first term in our objective function, namely $\mathcal{L}_{\text{roll}}(\mathbf{x},\mathbf{c},y,\mathbf{\mu})$, is to provide feedback to the concept selection policy $\psi$ regarding good choices for concepts to intervene next. Inspired by the “Skyline” optimal greedy policy proposed by coop, we design this loss term by first noticing that, given an intervention mask $\mu\in\{0,1\}^{k}$ with at least one unintervened concept in it, if we have an oracle that enables us to obtain any sample’s ground truth concepts $\mathbf{c}$ and task labels $y$, then we could tractably compute the optimal next concept $c_{*}$ to intervene on given by calculating

$$c_{*}=\operatorname*{arg\,max}_{\begin{subarray}{c}1\leq i\leq k\\ \mu_{i}\neq 1\end{subarray}}f\big{(}\tilde{g}(\mathbf{x},\mathbf{\mu},\mathbf{c})\big{)}_{y}$$

where $f(\cdot)_{y}$ represents the probability of class $y$ predicted by the label predictor $f$. Intuitively, $c_{*}$ corresponds to the concept whose ground-truth label and intervention would lead to the maximum probability in the ground-truth class $y$ after a single intervention. Notice that at test-time, it is impossible for one to determine $c_{*}$ as we do not know $y$ and the ground truth value of $c_{*}$ (this is what we are trying to request from an expert). Nevertheless, we can exploit the fact that these labels are available during training by providing meaningful feedback to $\psi$ through the following loss function:

$$\mathcal{L}_{\text{roll}}(\mathbf{x},\mathbf{c},y,\mathbf{\mu})=\text{CE}(\mathbf{\eta}^{*}(\mathbf{x},\mathbf{\mu},\mathbf{c},y),\psi(\tilde{g}(\mathbf{x},\mathbf{\mu},\mathbf{c}),\mathbf{\mu}))$$

In this term, we use $\text{CE}(\mathbf{p},\hat{\mathbf{p}})$ to mean the cross-entropy loss between ground truth distribution $\mathbf{p}$ and predicted distribution $\hat{\mathbf{p}}$ and use as a target distribution of the CE, i.e. $\mathbf{\eta}^{*}(\mathbf{x},\mathbf{\mu},\mathbf{c},y)$, a distribution that captures how much the probability of the ground-truth output class changes as one intervenes on different concepts. If we let $\mathbf{\mu}\vee\mathds{1}_{i}$ represent the action of adding an intervention on concept $i$ to the mask of intervened concepts $\mathbf{\mu}$, we can compute a meaningful target distribution by assigning $f\big{(}\tilde{g}(\mathbf{x},\mathbf{\mu}\vee\mathds{1}_{i},\mathbf{c})$ as a score to each concept $c_{i}$ and normalising the resulting vector of scores:

$$\mathbf{\eta}^{*}(\mathbf{x},\mathbf{\mu},\mathbf{c},y):=\text{Softmax}\Big{(}\Big{[}f\big{(}\tilde{g}(\mathbf{x},\mathbf{\mu}\vee\mathds{1}_{1},\mathbf{c})\big{)}_{y},\cdots,f\big{(}\tilde{g}(\mathbf{x},\mathbf{\mu}\vee\mathds{1}_{k},\mathbf{c})\big{)}_{y}\Big{]}\Big{)}$$

Intuitively, this can be thought of as a form of training-time behavioural cloning [behavioural_cloning, ross2011reduction], where $\psi$ is being trained to learn to predict $c_{*}$ given $\mathbf{\mu}$ and the current bottleneck $\tilde{g}(\mathbf{x},\mathbf{\mu},\mathbf{c})$.

As in traditional CEMs, we include a loss term that penalises our model for not being able to predict the ground-truth label $y$ from $\hat{\mathbf{c}}$. In contrast to how traditional CEMs or CBMs use such a loss, however, our task loss term $\mathcal{L}_{\text{task}}$ has two key components. The first component includes a cross-entropy term $\text{CE}\big{(}y,f(\tilde{g}(\mathbf{x},\mathbf{\mu},\mathbf{c}))$ that penalises our model for mispredicting $y$ after interventions $\mathbf{\mu}$. This term, however, is scaled by a factor of $\gamma^{\sum_{i=1}^{k}\mu_{i}}$ (where $\gamma\geq 1$ is a user-selected hyperparameter close to $1$) whose purpose is to penalise a task misprediction proportionally to the number of intervened concepts. The second component of our loss works by sampling a next concept intervention $\mathbf{\eta}\in\{0,1\}^{k}$ from the probability distribution learnt by $\psi(\tilde{g}(\mathbf{x},\mathbf{\mu},\mathbf{c}),\mathbf{\mu})$ and compute the task cross-entropy error $\text{CE}\big{(}y,f(\tilde{g}(\mathbf{x},\mathbf{\mu},\mathbf{c}))$ following this extra interventions. As done for when we intervene only on concepts $\mathbf{\mu}$, we scale this loss term by $\gamma^{1+\sum_{i=1}^{k}\mu_{i}}$ to take into account how many interventions we have performed up to that point. When we put everything together we obtain the following loss term:

$$\mathcal{L}_{\text{roll}}(\mathbf{x},\mathbf{c},y,\mathbf{\mu})=\text{CE}(\mathbf{\eta}^{*}(\mathbf{x},\mathbf{\mu},\mathbf{c},y),\psi(\tilde{g}(\mathbf{x},\mathbf{\mu},\mathbf{c}),\mathbf{\mu}))$$

Notice that in order for us to update the parameters in $\psi$ using feedback from the output class, we must be able to propagate gradients to the sampling operation performed to obtain $\mathbf{\eta}$. In this work, we achieve this by using a differentiable Gumbel-Softmax [gumbel_softmax] categorical sampler which enables us to sample from a categorical distribution whose parameters are being generated by $\psi(\cdot)$ in a differentiable manner. Finally, for the sake of simplicity and tractability, unless specified otherwise, in practice we estimate the inner expectation with a single Monte Carlo sample from $\psi(\cdot)$.

The last term in our loss incentivises accurate concept explanations for IntCEM’s predictions using the cross entropy loss averaged across all training concepts:

$$\mathcal{L}_{\text{concept}}(\mathbf{c},\hat{\mathbf{p}}):=\frac{1}{k}\sum_{i=1}^{k}\text{CE}(c_{i},\hat{p}_{i})=\frac{1}{k}\sum_{i=1}^{k}\big{(}-c_{i}\log{\hat{p}_{i}}-(1-c_{i})\log{(1-\hat{p}_{i})}\big{)}$$

Using our defined objective function, we learn IntCEM’s parameters $\theta$ by minimising, via gradient descent, the expected value of $\mathcal{L}$ in the training set $\mathcal{D}$ while we sample interventions masks $\mathbf{\mu}$ from a prior distribution $p(\mathbf{\mu})$:

$$\theta^{*}=\operatorname*{arg\,min}_{\theta}\mathbb{E}_{(\mathbf{x},\mathbf{c},y)\sim\mathcal{D}}\Big{[}\mathbb{E}_{\mathbf{\mu}\sim p(\mathbf{\mu})}\big{[}\mathcal{L}(\mathbf{x},\mathbf{c},y,\mathbf{\mu})\big{]}\Big{]}$$

(1)

Notice that the outer expectation can be estimated through Monte Carlo samples coming from our training set while the inner expectation can be estimated using $S_{\text{masks}}$ Monte Carlo samples of $p(\mathbf{\mu})$. It is important to notice that the concept mask’s prior distribution $p(\mathbf{\mu})$ may incorporate domain-specific knowledge by considering possible correlations between concepts. For example, if one can partition the set of $k$ training concepts into $k^{\prime}$ subsets of mutually exclusive concept groups (as is the case for popular datasets used in concept learning such as CUB [cub]), then a useful prior may be constructed by first sampling a mask of concept groups $\mathbf{\mu}^{\prime}$ from $\text{Bernoulli}(p_{\text{int}})$, for some user-selected hyperparameter $p_{\text{int}}\in[0,1]$, and then generating a concept-level mask $\mathbf{\mu}$ by setting $\mu_{i}$ be $1$ if its parent concept group was selected by $\mathbf{\mu}^{\prime}$. For the sake of simplicity, we follow such a prior in this work and leave exploring more interesting and domain-specific priors as future work. Similarly, for the sake of tractability and simplicity, unless specified otherwise we estimate the inner expectation in Eq. (1) using a single Monte Carlo sample of $p(\mathbf{\mu})$ (i.e., $S_{\text{masks}}$ = 1). TODO: We explore how our results change when we sample multiple masks per training sample in Appendix A.6.

We study how IntCEM performs as we intervene on more concepts by deploying our method on five vision tasks defined on three different datasets. First, inspired by the UMNIST dataset proposed by collins2023human, we define two simple tasks based on the MNIST dataset [mnist]. In our first task, which we call MNIST-Add, one is given a total of $12$ MNIST digits per sample (where each operand is confined to be within a fixed range of possible values) and needs to determine whether the sum of these digits is at least half of the maximum attainable sum. This binary task is provided as concept annotations the one-hot-encodings of each operand’s real value. The MNIST-based second task, which we call MNIST-Add-Incomp, is exactly the same as MNIST-Add with the exception that we provide as concept annotations the value of only 8/12 operands (the operands whose annotations are provided are selected at random before training commences). We explore these two variants of the same task to evaluate how IntCEM behaves when the set of concept annotations is both a complete and an incomplete description of the output task. Following this, we evaluate IntCEM on two tasks based on the CUB [cub] dataset. This dataset contains images of birds in the wild together with a downstream label representing the bird’s type and a set of binary annotations describing some of the attributes of a bird. The first task, which we simply call CUB, processes the CUB dataset as suggested by cbm and selects 112 of the bird’s attributes (grouped into 28 sets of mutually exclusive binary attributes) as concept annotations while using the bird type as the downstream task. The second CUB-based task, which we call CUB-Incomp, is constructed in the same line as MINST-Add-Incomp by providing only $25\%$ of all concept groups during training. Finally, we evaluate our method in a task defined on the Celebrity Attributes Dataset [celeba] (CelebA) using the same concept and task labels as cem. This results in a task with 256 classes and 6 binary concept annotations. For specific details of these datasets, including their size and properties, refer to Appendix A.1.

For the sake of obtaining a fair and complete evaluation of IntCEM, we compare its performance against that of a traditional CEM with the exact same architecture. Similarly, given that state-of-the-art intervention policies were designed with CBMs in mind rather than CEMs, we include in our evaluation traditional CBMs that are trained jointly (Joint CBM), sequentially (Sequential CBM) and independently (Independent CBM). Furthermore, given the strong empirical evidence that the activation function used in a joint CBM can have significant consequences on how interventions affect the CBM’s performance [cbm, metric_paper], we also include as part of our baselines both a joint CBM with a Sigmoidal bottleneck (Joint CBM-Sigmoid) and a CBM whose bottleneck correspond to the log probabilities of the predicted concepts (Joint CBM-Logit). Throughout all of our experiments, we attempt to get as close as a fair comparison across baselines by using the same concept encoder and label predictor architectures as in IntCEM for the concept encoders and label predictors of all other baselines. Across all tasks, unless specified otherwise, we tune IntCEM’s $\lambda_{\text{roll}}$ hyperparameter by selecting the model with the best validation loss after we vary $\lambda_{\text{roll}}$ over $\{5,1,0.1\}$. As for $\lambda_{concept}$, we fix this regulariser’s strength to the same value (selected from $\{10,5,1\}$ using the validation error of a traditional CEM) across all methods for a given task. The used embedding size $m$ and the initial intervention mask probability $p_{\text{int}}$ are set to $m=16$ and $p_{\text{int}}=0.25$ following the suggestions by cem. For an ablation study on how these hyperparameters affect our observed results, refer to Appendix A.4. Similarly, for a complete set of details on the architectures and hyperparameters used to train each baseline, refer to Appendix A.2.

We explore an IntCEM receptiveness to interventions by computing its test performance while intervening following one of these policies: (1) IntCEM’s learnt policy $\psi$, where we select the next intervention concept by choosing the unknown concept whose probability of selection, as given by $\psi(\cdot)$, is the highest, (2) a Random intervention policy, where the next concept is selected, uniformly at random, from the set of unknown concepts, (3) the Uncertainty of Concept Prediction (UCP) [closer_look_at_interventions] policy, where the next concept is selected by choosing the concept $c_{i}$ whose predicted probability $\hat{p}_{i}$ has the highest uncertainty (measured by $1/|\hat{p}_{i}-0.5|$), (4) the Cooperative Policy (CooP) [coop], where we select the concept that maximises a linear combination of their predicted uncertainty (akin to UCP) and the expected change in the predicted label’s probability when intervening on that concept, and finally (5) Skyline, an oracle intervention policy which selects $c_{*}$ on every step, indicating an upper bound for how good a greedy policy may be for a given task. Furthermore, as IntCEM learns a policy on the fly that attempts to imitate the Skylne policy, we include a Behavioural Cloning [behavioural_cloning] policy which we train to predict $c_{*}$ from pairs of training samples $(\tilde{g}(\mathbf{x},\mathbf{\mu},\mathbf{c}),\mathbf{\mu})$ after IntCEM has been trained. For details on how each of these policy baselines is used, including how their hyperparameters are selected, see Appendix A.5.

Our intervention policy study, shown in Figure 3, underscores three crucial properties of IntCEM. First, it shows that IntCEM’s learnt policy $\psi$ is able to capture meaningful properties about the intervention process as it is able to outperform a random intervention policy across all tasks (by more than $\sim 10\%$ in some instances as in CUB. Second, we observe there is a strong incentive to learn $\psi$ in conjunction with $f$ and $g$ as its equivalent behavioural cloning policy, which is learnt after that IntCEM has been trained, performs significantly worse than our learnt policy, only being slightly better than the random policy in CUB (TODO: In the current CelebA plot the behavioural cloning policy is worse than the learnt policy but this will be fixed once I update these plots to include the latest results). Finally, we observe that (i) using CooP to intervene on top of IntCEM rather than $\psi$ can lead to even further performance boosts from interventions. This, and the significant difference between the optimal policy and RandInt’s learnt policy, suggests that future work could develop stronger inductive biases for $\psi$ to improve its predictive capabilities (TODO: The difference with CooP is something that may change very soon with the new results I am getting. However, for the sake of being conservative, I will frame it now as something worth exploring in the future). For further results showing similar trends in the rest of the tasks used for this paper, as well as showing how static policies fare against the dynamic intervention policies shown here, refer to TODO: Appendix A.5.

Following our initial intervention evaluation of IntCEM using different intervention policies, we proceed to study how IntCEM, when intervened with its own learnt policy, fares against competing methods intervened with CooP. For the sake of being able to clearly study trends in each model’s receptiveness to interventions, we only show a comparison between IntCEM’s own policy and CooP in other baselines. We opt to present our results in this way as we empirically observe that CooP outperforms all other non-oracle policies for other methods (see Appendix A.5). Our results, seen in Figure 1(a), show that IntCEM’s task accuracy is significantly better than that of competing methods across all tasks. Specifically, we

Our test task accuracy evaluation, shown in the top half of Table 1, shows that throughout all but one of the tasks, IntCEMs are as accurate if not more accurate than all other baselines. It is only in the CUB task where we observe a small drop in performance (less than $\sim 0.5\%$) for the best-performing IntCEM variant with $\lambda_{\text{roll}}=0.1$. Notice that from the five tasks we use during evaluation, three of them, namely MNIST-Add-Incomp, CUB-Incomp, and CelebA, do not have a complete set of concept annotations at training time. Therefore, our results suggest that just as in traditional CEMs, our model is able to maintain a high task performance even when the set of concept annotations is not fully descriptive of the downstream task, a highly desired property to have for deploying these models in real-world tasks.

The second half of Table 1 shows that IntCEM’s mean concept AUC, particularly for low values of $\lambda_{\text{roll}}$, is competitive against that of CEMs, falling behind it by at most $\sim 1.2\%$ at its worst. In the MNIST-based tasks, we observe that both IntCEMs and CEMs lag behind traditional joint CBMs in their concept of predictive AUC. Nevertheless, we note that the gains seen in concept accuracy for the joint CBMs over IntCEM do not merit their observed drop in task performance seen compared to IntCEMs in these two tasks. These results suggest that IntCEM’s concept explanations are as accurate as those provided by state-of-the-art CEMs, leading a set of predicted concepts by IntCEM to constitute a faithful explanation of its downstream prediction.

One interesting effect we observe in both our task and concept accuracy results is that IntCEM is particularly sensitive to its $\lambda_{\text{roll}}$ parameter, as seen in particular for CUB. As we will explore in more detail in the next section, these observed differences seem to indicate there exists a tradeoff in IntCEM’s objective function between incentivising the model to be highly accurate without any interventions and incentivising the model to be highly accurate after a small number of interventions. This tradeoff can be calibrated depending on the intended application of the IntCEM: if one expects to deploy the model in a setting where interventions will be expected in the average case, then using a larger value of $\lambda_{\text{roll}}$ during training may condition the model to be very receptive to interventions, albeit its initial performance may be slightly lower than an equivalent CEM. On the other hand, if one expects to intervene on a trained IntCEM very infrequently, a smaller value of $\lambda_{\text{roll}}$ may lead to an IntCEM with state-of-the-art performance without interventions yet whose initial gain in performance once we intervene on it will not be as high as a model trained with a higher value of $\lambda_{\text{roll}}$. While this hyperparameter adds complexity to IntCEM’s training process and hyperparameter section process, in practice we found that trying out values in $\lambda_{\text{roll}}\in\{0.1,1,5\}$ results in highly accurate CEMs.