SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mateo Espinosa Zarlenga
University of Cambridge me466@cam.ac.uk
Katherine M. Collins
University of Cambridge kmc61@cam.ac.uk
Krishnamurthy (Dj) Dvijotham
Google DeepMind dvij@google.com
University of Cambridge Alan Turing Institute aw665@cam.ac.uk
Zohreh Shams
University of Cambridge  
zs315@cam.ac.uk
Mateja Jamnik
University of Cambridge mateja.jamnik@cl.cam.ac.uk
Concept Bottleneck Models (CBMs) tackle the opacity of neural architectures by constructing and explaining their predictions using a set of high-level concepts. A special property of these models is that they permit concept interventions, wherein users can correct mispredicted concepts and thus improve the model's performance. Recent work, however, has shown that intervention efficacy can be highly dependent on the order in which concepts are intervened on and on the model's architecture and training hyperparameters. We argue that this is rooted in a CBM's lack of train-time incentives for the model to be appropriately receptive to concept interventions. To address this, we propose Intervention-aware Concept Embedding models (IntCEMs), a novel CBM-based architecture and training paradigm that improves a model's receptiveness to test-time interventions. Our model learns a concept intervention policy in an end-to-end fashion from where it can sample meaningful intervention trajectories at train-time. This conditions IntCEMs to effectively select and receive concept interventions when deployed at test-time. Our experiments show that IntCEMs significantly outperform state-of-the-art concept-interpretable models when provided with test-time concept interventions, demonstrating the effectiveness of our approach.
It is important to know how to ask for help, but also important to know how to receive help. Knowing how to react to feedback positively allows for bias correction [1] and efficient learning [2] while being instrumental for mass collaborations [3]. Nevertheless, although the uptake of feedback is ubiquitous in real-world decision-making, the same cannot be said about modern artificial intelligence (AI) systems, where deployment tends to be in isolation from experts whose feedback could be queried.
Progress in this aspect has recently come from Explainable AI, where interpretable Deep Neural Networks (DNNs) that can benefit from expert feedback at test-time have been proposed [4-7]. In particular, Concept Bottleneck Models (CBMs) [4], a family of interpretable DNNs, enable expert-model interactions by generating, as part of their inference process, explanations for their predictions using high-level "concepts". Such explanations allow human experts to better understand a CBM's prediction based on interpretable units of information (e.g., input is a "cat" because it has "paws" and "whiskers") rather than low-level features (e.g., input pixels). This design enables humans to provide feedback to the CBM at test-time via concept interventions [4] (Figure 1), a process specific to CBM-like architectures in which an expert analyses a CBM's inferred concepts
Figure 1: When intervening on a CBM, a human expert analyses the predicted concepts and corrects mispredicted values (e.g., the mispredicted concept "legs"), allowing the CBM to update its prediction.
and corrects mispredictions. Concept interventions allow a CBM to update its final prediction after incorporating new expert feedback, and have been shown to lead to significant performance boosts [4, 5, 8, 9]. Nevertheless, and in contrast to these results, recent works have found that concept interventions can potentially increase a CBM's test-time error depending on its choice of architecture [10] and its set of train-time annotations [5]. This discrepancy suggests that a CBM's receptiveness to interventions is neither guaranteed nor fully understood.
In this paper, we posit that the sensitivity of a CBM's receptiveness to interventions is an artefact of its training objective function. Specifically, we argue that CBMs lack an explicit incentive to be receptive to interventions as they are neither exposed to interventions at train-time nor optimised to perform well under interventions. To address this, we propose a novel intervention-aware objective function and architecture based on Concept Embedding Models (CEMs) [5], a generalisation of CBMs that represent concepts as high-dimensional embeddings. Our architecture, which we call Intervention-aware Concept Embedding Model (IntCEM), has two key distinctive features. First, it learns an intervention policy end-to-end which imposes a prior over which concepts an expert could be queried to intervene on next to maximise intervention receptiveness. Second, its loss function includes an explicit regulariser that penalises IntCEM for mispredicting the task label after a series of train-time interventions sampled from its learned policy. We highlight the following contributions:
- We introduce IntCEM, the first concept-interpretable neural architecture that learns not only to explain its predictions using concepts but also learns an intervention policy dictating which concepts an expert should intervene on next to maximise test performance.
- We show that IntCEM significantly outperforms all baselines in the presence of interventions while maintaining competitiveness when deployed without any interventions.
- We demonstrate that, by rewriting concept interventions as a differentiable operation, IntCEM learns an efficient dynamic intervention policy that selects performance-boosting concepts to intervene on at test-time.
2 Background and Previous Work
Concept Bottleneck Models Methods in Concept-based Explainable AI [11-17] advocate for explaining DNNs by constructing explanations for their predictions using high-level concepts captured within their latent space. Within these methods, Concept Bottleneck Models (CBMs) [4] provide a framework for developing concept-based interpretable DNNs.
Given a training set  $\mathcal{D} := \left\{(\mathbf{x}^{(i)}, \mathbf{c}^{(i)}, y^{(i)})\right\}_{i=1}^{N}$ , where each sample  $\mathbf{x} \in \mathbb{R}^n$  (e.g., an image) is annotated with a task label  $y \in \{1, \dots, L\}$  (e.g., "cat" or "dog") and  $k$  binary concepts  $\mathbf{c} \in \{0, 1\}^k$  (e.g., "has whiskers"), a CBM is a pair of functions  $(g, f)$  whose composition  $f(g(\mathbf{x}))$  predicts a probability over output classes given  $\mathbf{x}$ . The first function  $g: \mathbb{R}^n \to [0, 1]^k$ , called the concept encoder, learns a mapping between features  $\mathbf{x}$  and concept activations  $\hat{\mathbf{c}} = g(\mathbf{x}) \in [0, 1]^k$ , where ideally  $\hat{c}_i$  is close to 1 when concept  $c_i$  is active in  $\mathbf{x}$ , and 0 otherwise. The second function  $f: [0, 1]^k \to [0, 1]^L$ , called the label predictor, learns a mapping from the predicted concepts  $\hat{\mathbf{c}}$  to a distribution  $\hat{\mathbf{y}} = f(\hat{\mathbf{c}})$  over  $L$  task classes. When implemented as DNNs,  $g$ 's and  $f$ 's parameters can be learnt (i) jointly, by minimising a combination of the concept predictive loss  $\mathcal{L}_{\mathrm{concept}}$  and the task cross-entropy loss  $\mathcal{L}_{\mathrm{task}}$ , (ii) sequentially, by first training  $g$  to minimise  $\mathcal{L}_{\mathrm{concept}}$  and then training
$f$  to minimise  $\mathcal{L}_{\mathrm{task}}$  from  $g$ 's outputs, or (iii) independently, where  $g$  and  $f$  are independently trained to minimise their respective losses. Because at prediction-time  $f$  has access only to the "bottleneck" of concepts activations  $\hat{\mathbf{c}}$ , the composition  $(f \circ g)$  yields a concept-interpretable model where  $\hat{c}_i$  can be interpreted as the probability  $\hat{p}_i$  that concept  $c_i$  is active in  $\mathbf{x}$ .
Concept Embedding Models Because  $f$  operates only on the concept scores predicted by  $g$ , if the set of training concepts is not predictive of the downstream task, then a CBM will be forced to choose between being highly accurate at predicting concepts or at predicting task labels. Concept Embedding Models (CEMs) [5] are a generalisation of CBMs that address this trade-off using high-dimensional concept representations in their bottlenecks. This design allows information of concepts not provided during training to flow into the label predictor via two  $m$ -dimensional vector representations (i.e., embeddings) for each training concept  $c_i$ :  $\mathbf{c}_i^+ \in \mathbb{R}^m$ , an embedding representing  $c_i$  when it is active (i.e.,  $c_i = 1$ ), and  $\mathbf{c}_i^- \in \mathbb{R}^m$ , an embedding representing  $c_i$  when it is inactive (i.e.,  $c_i = 0$ ).
When processing sample  $\mathbf{x}$ , for each training concept  $c_{i}$  a CEM constructs  $\mathbf{c}_{i}^{+}$  and  $\mathbf{c}_{i}^{-}$  by feeding  $\mathbf{x}$  into two learnable models  $\phi_{i}^{+},\phi_{i}^{-}:\mathbb{R}^{n}\to \mathbb{R}^{m}$  implemented as DNNs with a shared preprocessing module. These embeddings are then passed to a learnable scoring model  $s:\mathbb{R}^{2m}\rightarrow [0,1]$ , shared across all concepts, that predicts the probability  $\hat{p}_i = s([\hat{\mathbf{c}}_i^+, \hat{\mathbf{c}}_i^-])$  of concept  $c_{i}$  being active. With these probabilities, one can define a CEM's concept encoder  $g(\mathbf{x}) = \hat{\mathbf{c}} := [\hat{\mathbf{c}}_1,\dots ,\hat{\mathbf{c}}_k] \in \mathbb{R}^{km}$  by mixing the positive and negative embeddings of each concept to generate a final concept embedding  $\hat{\mathbf{c}}_i := \hat{p}_i\hat{\mathbf{c}}_i^+ + (1 - \hat{p}_i)\hat{\mathbf{c}}_i^-$ . Finally, a CEM can predict a label  $\hat{\mathbf{y}}$  for sample  $\mathbf{x}$  by passing  $\hat{\mathbf{c}}$  (i.e., its "bottleneck") to a learnable label predictor  $f(\hat{\mathbf{c}})$  whose output can be explained with the concept probabilities  $\hat{\mathbf{p}} := [\hat{p}_1,\dots ,\hat{p}_k]^T$ . This generalisation of CBMs has been empirically shown to outperform CBMs especially when the set of concept annotations is incomplete [5].
Concept Interventions A key property of both CBMs and CEMs is that they allow concept interventions, whereby experts can correct mispredicted concept probabilities  $\hat{\mathbf{p}}$  at test-time. By updating the predicted probability of concept  $c_{i}$  (i.e.,  $\hat{p}_i$ ) so that it matches the ground truth concept label (i.e., setting  $\hat{p}_i\coloneqq c_i$ ), these models can update their predicted bottleneck  $\hat{\mathbf{c}}$  and propagate that change into their label predictor  $g(\hat{\mathbf{x}})$ , leading to a potential update in the model's output prediction.
Let  $\mu \in \{0,1\}^k$  be a mask with  $\mu_{i} = 1$  if we will intervene on concept  $c_{i}$ , and  $\mu_{i} = 0$  otherwise. Here, assume that we are given the ground truth values  $\tilde{\mathbf{c}}$  of all the concepts we will intervene on. For notational simplicity, and without loss of generality, we can extend  $\tilde{\mathbf{c}}$  to be in  $[0,1]^k$  by setting  $\tilde{c}_i = c_i$  if  $\mu_{i} = 1$ , and  $\tilde{c}_i = 0.5$  otherwise (where 0.5 expresses full uncertainty but, as it will become clear next, the specific value is of no importance). Thus, this vector contains the expert-provided ground truth concept labels for concepts we are intervening on, while it assigns an arbitrary value (e.g., 0.5) to all other concepts. We define an intervention as a process where, for all concepts  $c_{i}$ , the activation(s)  $\hat{\mathbf{c}}_i$  corresponding to concept  $c_{i}$  in a bottleneck  $\hat{\mathbf{c}}$  are updated as follows:
$$
\hat {\mathbf {c}} _ {i} := \left(\mu_ {i} \tilde {c} _ {i} + (1 - \mu_ {i}) \hat {p} _ {i}\right) \hat {\mathbf {c}} _ {i} ^ {+} + \left(1 - \left(\mu_ {i} \tilde {c} _ {i} + (1 - \mu_ {i}) \hat {p} _ {i}\right)\right) \hat {\mathbf {c}} _ {i} ^ {-} \tag {1}
$$
where  $\hat{\mathbf{c}}_i^+$  and  $\hat{\mathbf{c}}_i^-$  are  $c_{i}$ 's positive and negative concept embeddings (for CBMs  $\hat{\mathbf{c}}_i^+ = [1]$  and  $\hat{\mathbf{c}}_i^- = [0]$ ). This update forces the mixing coefficient between the positive and negative concept embeddings to be the ground truth concept value  $c_{i}$  when we intervene on concept  $c_{i}$  (i.e.,  $\mu_{i} = 1$ ), while maintaining it as  $\hat{p}_i$  if we are not intervening on concept  $c_{i}$ . The latter fact holds as  $(\mu_i\tilde{c}_i + (1 - \mu_i)\hat{p}_i) = \hat{p}_i$  when  $\mu_{i} = 0$ , regardless of  $\tilde{c}_i$ . Because the predicted positive and negative concept embeddings, as well as the predicted concept probabilities, are all functions of  $\mathbf{x}$ , for simplicity we use  $\tilde{g} (\mathbf{x},\mu ,\tilde{\mathbf{c}})$  to represent the updated bottleneck of CBM  $(g,f)$  when intervening with mask  $\mu$  and concept values  $\tilde{\mathbf{c}}$ .
We note that although we call the operation defined in Equation 1 a "concept intervention" to follow the term used in the concept learning literature, it should be distinguished from interventions in causal models [18-20]. Instead, concept interventions in this context are connected, yet not identical, to previous work in active feature acquisition in which one has access to an expert at test-time to request the ground truth value of a small number of input features [21-23].
Intervention Policies In practice, it may not be sensible to ask a human to intervene on every concept [24]; instead, we may want a preferential ordering of which concepts to query first (e.g., some concepts may be highly informative of the output task, yet hard for the model to predict). "Intervention policies" have been developed to produce a sensible ordering for concept interventions [8, 25, 9].
These orderings can be produced independently of a particular instance (i.e., a static policy), or selected on a per-instance basis (a dynamic policy). The latter takes as an input a mask of previously intervened concepts  $\mu$  and a set of predicted concept probabilities  $\hat{\mathbf{p}}$ , and determines which concept should be requested next from an expert, where the retrieved concept label is assumed to be "ground truth". Two such policies are Cooperative Prediction (CooP) [9] and Expected Change in Target Prediction (ECTP) [8], which prioritise concept  $c_{i}$  if intervening on  $c_{i}$  next leads to the largest change in the probability of the currently predicted class in expectation (with the expectation taken over the distribution given by  $p(c_{i} = 1) = \hat{p}_{i}$ ). These policies, however, may become costly as the number of concepts increases and share the limitation that, by maximising the expected change in probability in the predicted class, they do not guarantee that the probability mass shifts towards the correct class.
3 Intervention-Aware Concept Embedding Models
Although there has been significant work on understanding how CBMs react to interventions [5, 6, 9, 25], the fact that these models are trained without any expert intervention – yet are expected to be highly receptive to expert feedback at test-time – is often overlooked. To the best of our knowledge, the only exception comes from Espinosa Zarlenga et al. [5], wherein the authors introduced RandInt, a procedure that randomly intervenes on the model during train-time. However, RandInt assumes that a model should be equally penalised for mispredicting  $y$  regardless of how many interventions have been performed. This leads to CEMs lacking the incentive to perform better when given more feedback. Moreover, RandInt assumes that all concepts are equally likely to be intervened on. This inadequately predisposes a model to assume that concepts will be randomly intervened on at test-time.
To address these limitations, we propose Intervention-Aware Concept Embedding Models (IntCEMs), a new CEM-based architecture and training framework designed for inducing receptiveness to interventions. IntCEMs are composed of two core components: 1) an end-to-end learnable concept intervention policy  $\psi (\hat{\mathbf{c}},\mu)$ , and 2) a novel objective function that penalises IntCEM if it mispredicts its output label after following an intervention trajectory sampled from  $\psi$  (with a heavier penalty if the trajectory is longer). By learning the intervention policy  $\psi$  in conjunction with the concept encoder and label predictor, IntCEM can simulate dynamic intervention steps at train-time while it learns a policy which enables it to "ask" for help about specific concepts at test-time.
3.1 Architecture Description and Inference Procedure
An IntCEM is a tuple of parametric functions  $(g, f, \psi)$  where: (i)  $g$  is a concept encoder, mapping features  $\mathbf{x}$  to a bottleneck  $\hat{\mathbf{c}}$ , (ii)  $f$  is a label predictor mapping  $\hat{\mathbf{c}}$  to a distribution over  $L$  classes, and (iii)  $\psi : \mathbb{R}^{km} \times \{0,1\}^k \to [0,1]^k$  is a concept intervention policy mapping  $\hat{\mathbf{c}}$  and a mask of previously intervened concepts  $\mu$  to a probability distribution over which concepts to intervene on next. IntCEM's concept encoder  $g$  works by (1) passing an input  $\mathbf{x}$  to a learnable "backbone"  $\zeta(\mathbf{x}) := \left(\left(\hat{\mathbf{c}}_i^+, \hat{\mathbf{c}}_i^-\right)\right)_{i=1}^k, \hat{\mathbf{p}}$ , which predicts a pair of embeddings  $\left(\hat{\mathbf{c}}_i^+, \hat{\mathbf{c}}_i^-\right)$  and a concept probability  $\hat{p}_i$  for all concepts, and (2) constructs a bottleneck  $\hat{\mathbf{c}}$  by concatenating all the mixed vectors  $\left\{\hat{p}_1\hat{\mathbf{c}}_1^+ + (1 - \hat{p}_1)\hat{\mathbf{c}}_1^-, \dots, \hat{p}_k\hat{\mathbf{c}}_k^+ + (1 - \hat{p}_k)\hat{\mathbf{c}}_k^-\right\}$ . While multiple instantiations are possible for  $\zeta$ , in this work we parameterise it using the same backbone as in a CEM. This implies that  $\zeta$  is formed by concept embedding generating models  $\left\{\left(\phi_i^+, \phi_i^-\right)\right\}_{i=1}^k$  and scoring model  $s$ . This yields a model which, at inference, is almost identical to a CEM, except that IntCEM also outputs a probability distribution  $\psi(\hat{\mathbf{c}}, \mu)$  placing a high density on concepts that may yield significant intervention boosts. At train-time, these models differ significantly, as we discuss next.
3.2 IntCEM's Training Procedure
The crux of IntCEM's training procedure, shown in Figure 2, lies in inducing receptiveness to test-time interventions by exposing our model to dynamic train-time interventions. Specifically, at the beginning of each training step, we sample an initial mask of intervened concepts  $\mu^{(0)}$  from a prior distribution  $p(\mu)$  as well as the number  $T\in \{1,\dots ,k\}$  of interventions to perform on top of  $\mu^{(0)}$  from prior  $T\sim p(T)$ . Next, we generate a trajectory of  $T$  new interventions  $\{\eta^{(t)}\in \{0,1\} ^k\sim p(\eta ;\omega^{(t)})\}_{t = 1}^T$  by sampling the  $t^{\mathrm{th}}$  intervention  $\eta^{(t)}\sim p(\eta ;\omega^{(t)})$  (represented as a one-hot encoding) from a categorical distribution with parameters  $\omega^{(t)}$ . A key component of this work is that parameters  $\omega^{(t)}$  are predicted using IntCEM's policy  $\psi \bigl (\hat{\mathbf{e}}^{(t - 1)},\mu^{(t - 1)}\bigr)$  evaluated on the
Figure 2: Given concept embeddings and probabilities  $\zeta(\mathbf{x}) = \left(\left(\hat{\mathbf{c}}_i^+, \hat{\mathbf{c}}_i^-\right)\right)_{i=1}^{k}, \hat{\mathbf{p}}$ , an IntCEM training step (1) samples an intervention mask  $\mu^{(0)} \sim p(\mu)$  and horizon  $T \sim p(T)$ , (2) generates an intervention trajectory  $\left\{(\mu^{(t-1)}, \eta^{(t)})\right\}_{t=1}^{T}$  from the learnable intervention policy  $\psi$ , (3) and predicts the task label at the start and end of the trajectory. Our loss incentives (i) good initial concept predictions  $(\mathcal{L}_{\mathrm{concept}})$ , (ii) performance-boosting intervention trajectories  $(\mathcal{L}_{\mathrm{roll}})$ , (iii) and low task loss before and after interventions  $(\mathcal{L}_{\mathrm{pred}})$ , with a heavier penalty  $\gamma^T > 1$  for mispredicting at the end of the trajectory. In this figure, dashed orange arrows indicate recursive steps in our sampling process.
previous bottleneck  $\hat{\mathbf{c}}^{(t - 1)}\coloneqq \tilde{g} (\mathbf{x},\mu^{(t - 1)},\mathbf{c})$  and intervention mask  $\mu^{(t - 1)}$ , where we recursively define  $\mu^{(t)}$  as  $\mu^{(t)}\coloneqq \mu^{(t - 1)} + \eta^{(t)}$  for  $t > 0$ . This allows IntCEM to calibrate its train-time trajectory of interventions to reinforce supervision on the embeddings corresponding to currently mispredicted concepts, leading to IntCEM learning more robust representations of inherently harder concepts.
We design our objective function such that it incentivises learning performance-boosting intervention trajectories and achieving high task performance as more interventions are provided:
$$
\mathcal {L} (\mathbf {x}, \mathbf {c}, y, \mathcal {T}) := \lambda_ {\mathrm {r o l l}} \mathcal {L} _ {\mathrm {r o l l}} (\mathbf {x}, \mathbf {c}, y, \mathcal {T}) + \mathcal {L} _ {\mathrm {p r e d}} (\mathbf {x}, \mathbf {c}, y, \boldsymbol {\mu} ^ {(0)}, \boldsymbol {\mu} ^ {(T)}) + \lambda_ {\mathrm {c o n c e p t}} \mathcal {L} _ {\mathrm {c o n c e p t}} (\mathbf {c}, \hat {\mathbf {p}})
$$
where  $\mathcal{T} := \{(\mu^{(t-1)},\eta^{(t)})\}_{t=1}^T$  is the intervention trajectory while  $\lambda_{\mathrm{roll}}$  and  $\lambda_{\mathrm{concept}}$  are user-defined hyperparameters.  $\lambda_{\mathrm{roll}}$  encodes a user's tradeoff for number of interventions versus task performance gains (e.g., high  $\lambda_{\mathrm{roll}}$  encourages rapid task performance gains from few interventions, whereas low  $\lambda_{\mathrm{roll}}$  is more suitable in settings where many expert queries can be made).  $\lambda_{\mathrm{concept}}$  expresses a user's preference between accurate explanations and accurate task predictions before any interventions.
Rollout Loss  $(\mathcal{L}_{\mathrm{roll}})$  The purpose of  $\mathcal{L}_{\mathrm{roll}}$  is to learn an intervention policy that prioritises the selection of concepts for which the associated concept embeddings fails to represent the ground truth concept labels, i.e., select concepts which benefit from human intervention. We achieve this through a train-time "behavioural cloning" [26, 27] approach where  $\psi$  is trained to learn to predict the action taken by an optimal greedy policy at each step. For this, we take advantage of the presence of ground truth concepts and task labels during training and incentivise  $\psi$  to mimic "Skyline", an oracle optimal policy proposed by Chauhan et al. [9]. Given a bottleneck  $\hat{\mathbf{c}}$  and ground-truth labels  $(y,\mathbf{c})$ , the Skyline policy selects  $c_*(\mathbf{x},\mu ,\mathbf{c},y)\coloneqq \arg \max_{1\leq i\leq k}f\big(\tilde{g} (\mathbf{x},\mu \vee \mathbb{1}_i,\mathbf{c})\big)_y$  as its next concept, where  $f(\cdot)_y$  is the probability of class  $y$  predicted by  $f$  and  $\mu \vee \mathbb{1}_i$  represents the action of adding an intervention on concept  $i$  to  $\mu$ . Intuitively,  $c_*$  corresponds to the concept whose ground-truth label and intervention would yield the highest probability over the ground-truth class  $y$ . We use such a demonstration to provide feedback to  $\psi$  through the following loss:
$$
\mathcal {L} _ {\text {r o l l}} (\mathbf {x}, \mathbf {c}, y, \mathcal {T}) := \frac {1}{T} \sum_ {t = 1} ^ {T} \mathrm {C E} \left(c _ {*} (\mathbf {x}, \mu^ {(t - 1)}, \mathbf {c}, y), \psi \left(\tilde {g} (\mathbf {x}, \mu^ {(t - 1)}, \mathbf {c}), \mu^ {(t - 1)}\right)\right)
$$
where  $\mathrm{CE}(\mathbf{p},\hat{\mathbf{p}})$  is the cross-entropy loss between ground truth distribution  $\mathbf{p}$  and predicted distribution  $\hat{\mathbf{p}}$ , penalising  $\psi$  for not selecting the same concept as Skyline throughout its trajectory.
Task Prediction Loss  $(\mathcal{L}_{\mathrm{pred}})$  We penalise IntCEM for mispredicting  $y$ , both before and after our intervention trajectory, by imposing a higher penalty when it mispredicts  $y$  at the end of the trajectory: