SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathcal {L} _ {\mathrm {p r e d}} (\mathbf {x}, \mathbf {c}, y, \mu^ {(0)}, \mu^ {(T)}) := \frac {\mathrm {C E} \big (y , f (\tilde {g} (\mathbf {x} , \mu^ {(0)} , \mathbf {c})) \big) + \gamma^ {T} \mathrm {C E} \big (y , f (\tilde {g} (\mathbf {x} , \mu^ {(T)} , \mathbf {c})) \big)}{1 + \gamma^ {T}}
$$
Here  $\gamma \in [1,\infty)$  is a scaling factor that penalises IntCEM more heavily for mispredicting  $y$  at the end of the trajectory, by a factor of  $\gamma^T > 1$ , than for mispredicting  $y$  at the start of the trajectory. In practice, we observe that a value in  $\gamma \in [1.1, 1.5]$  works well; we use  $\gamma = 1.1$  in our experiments unless specified otherwise. For an ablation study showing that IntCEM's intervention performance surpasses that of existing methods for a wide array of scaling factors  $\gamma$ , see Appendix A.6.
A key realisation is that expressing interventions as in Equation (1) yields an operator which is differentiable with respect to the concept intervention mask  $\mu$ . This trick allows us to backpropagate gradients from  $\mathcal{L}_{\mathrm{pred}}$  into  $\psi$  when the sampling operation  $\eta^{(t)} \sim p(\eta | \psi(\hat{\mathbf{c}}^{(t-1)}, \mu^{(t-1)}))$  is differentiable with respect to its parameters. In this work, this is achieved by relaxing our trajectory sampling using a differentiable Gumbel-Softmax [28] categorical sampler (see Appendix A.2 for details).
Concept Loss  $(\mathcal{L}_{\mathrm{concept}})$  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 {c o n c e p t}} (\mathbf {c}, \hat {\mathbf {p}}) := \frac {1}{k} \sum_ {i = 1} ^ {k} \mathrm {C E} \left(c _ {i}, \hat {p} _ {i}\right) = \frac {1}{k} \sum_ {i = 1} ^ {k} \left(- c _ {i} \log \hat {p} _ {i} - \left(1 - c _ {i}\right) \log \left(1 - \hat {p} _ {i}\right)\right)
$$
Putting everything together  $(\mathcal{L})$  We learn IntCEM's parameters  $\theta$  by minimising the following:
$$
\theta^ {*} = \underset {\theta} {\arg \min } \mathbb {E} _ {(\mathbf {x}, \mathbf {c}, y) \sim \mathcal {D}} \left[ \mathbb {E} _ {\mu \sim p (\mu), T \sim p (T)} \left[ \mathcal {L} (\mathbf {x}, \mathbf {c}, y, \mathcal {T} (\mathbf {x}, \mathbf {c}, \mu^ {(0)}, T)) \right] \right] \tag {2}
$$
The outer expectation can be optimised via stochastic gradient descent, while the inner expectation can be estimated using Monte Carlo samples from the user-selected priors  $(p(\mu), p(T))$ . For the sake of simplicity, we estimate the inner expectation using a single Monte Carlo sample per training step for both the initial mask  $\mu^{(0)}$  and the trajectory horizon  $T$ . In our experiments, we opt to use a Bernoulli prior for  $p(\mu)$ , where each concept is selected with probability  $p_{\mathrm{int}} = 0.25$ , and a discrete uniform prior  $\mathrm{Unif}(\{1, \dots, T_{\mathrm{max}}\})$  for  $p(T)$ , where  $T_{\mathrm{max}}$  is annealed within [2, 6]. Although domain-specific knowledge can be incorporated into these priors, we leave this for future work and show in Appendix A.6.3 that IntCEMs are receptive to interventions as we vary  $p_{\mathrm{int}}$  and  $T_{\mathrm{max}}$ .
In this section, we evaluate IntCEMs by exploring the following research questions:
- Unintervened Performance (Q1): In the absence of interventions, how does IntCEM's task and concept predictive performance compare to CEMs's and CBM's? Does IntCEM's updated loss function have detrimental effects on its unintervened performance?
- Intervention Performance (Q2A): Are IntCEMs more receptive to test-time interventions than state-of-the-art CBM variants?
- Effects of Policies during Deployment (Q2B): What is the impact on IntCEM's performance when employing different intervention policies at test-time?
- Benefits of End-to-End Learning of  $\psi$  (Q3): Is it beneficial to learn our intervention policy  $\psi$  at train-time?
Datasets and Tasks We consider five vision tasks: (1) MNIST-Add, a task inspired by the UMNIST dataset [29] where one is provided with 12 MNIST [30] images containing digits in  $\{0,\dots ,9\}$ , as well as their digit labels as concept annotations, and needs to predict if their sum is at least half of the maximum attainable, (2) MNIST-Add-Incomp, a concept-incomplete [17] extension of the MNIST-Add task where only 8/12 operands are provided as concept annotations, (3) CUB [31], a bird classification task whose set of concepts cover 112 bird features, e.g., "wing shape", selected by Koh et al. [4], (4) CUB-Incomp, a concept-incomplete extension of CUB where we provide only  $25\%$  of all concept groups, and (5) CelebA, a task on the Celebrity Attributes Dataset [32] whose task and concept labels follow those selected in [5]. For more details refer to Appendix A.3.
Models We evaluate IntCEMs against CEMs with the exact same architecture, averaging all metrics of interest over five different random initialisations. We also include CBMs that are trained jointly (Joint CBM), sequentially (Seq. CBM) and independently (Ind. CBM) as part of our evaluation to
Table 1: Task predictive performance (accuracy or AUC in %) and mean concept AUC (%) across all tasks and baselines. Values within one standard deviation from the highest baseline are emphasised. CelebA's task accuracies are relatively low across all baselines given its inherent difficulty (256 highly imbalanced classes). Nevertheless, we highlight these results match previous observations [5].
Dataset	IntCEM(λvoll = 5)	IntCEM(λvoll = 1)	IntCEM(λroll = 0.1)	CEM	Joint CBM-Sigmoid	Joint CBM-Logit	Ind. CBM	Seq. CBM
Task	MNIST-Add	90.74 ± 1.66	92.07 ± 0.20	92.05 ± 0.23	90.13 ± 0.41	79.25 ± 2.93	86.09 ± 0.60	87.66 ± 0.62	83.20 ± 1.08
MNIST-Add-Incomp	89.41 ± 0.16	89.43 ± 0.39	89.63 ± 0.13	86.61 ± 0.75	76.61 ± 1.13	84.03 ± 0.87	86.49 ± 0.67	85.23 ± 0.78
CUB	76.20 ± 0.98	78.03 ± 0.58	77.79 ± 0.19	79.00 ± 0.61	74.95 ± 0.56	78.16 ± 0.24	60.25 ± 3.20	53.41 ± 4.10
CUB-Incomp	70.39 ± 1.11	74.70 ± 0.22	74.66 ± 0.50	75.52 ± 0.37	59.73 ± 4.06	68.42 ± 7.91	44.82 ± 0.33	40.00 ± 1.59
CelebA	38.09 ± 0.26	32.89 ± 0.86	30.66 ± 0.41	31.04 ± 0.96	24.15 ± 0.43	23.87 ± 1.26	24.25 ± 0.26	24.53 ± 0.3
Concept	MNIST-Add	80.81 ± 3.28	85.26 ± 0.22	85.03 ± 0.27	85.51 ± 0.26	89.73 ± 0.15	89.21 ± 0.71	82.29 ± 0.66	82.29 ± 0.66
MNIST-Add-Incomp	87.07 ± 1.08	87.13 ± 1.98	87.36 ± 1.44	86.98 ± 0.01	91.15 ± 0.45	90.88 ± 0.25	87.58 ± 0.37	87.58 ± 0.37
CUB	89.63 ± 0.42	92.51 ± 0.08	93.39 ± 0.22	94.46 ± 0.04	93.76 ± 0.16	93.68 ± 0.09	89.85 ± 0.78	89.85 ± 0.78
CUB-Incomp	92.70 ± 0.19	94.42 ± 0.13	94.50 ± 0.15	94.65 ± 0.10	93.58 ± 0.09	92.87 ± 2.01	93.65 ± 0.25	93.65 ± 0.25
CelebA	88.25 ± 0.24	86.46 ± 0.48	85.35 ± 0.07	85.28 ± 0.18	81.14 ± 0.75	82.85 ± 0.44	82.85 ± 0.21	82.85 ± 0.21
be able to compare IntCEM's policy against state-of-art intervention policies developed for CBMs. Since previous work shows that interventions in joint CBMs are sensitive to its bottleneck activation function [4, 10], we include joint CBMs with a sigmoidal (Joint CBM-Sigmoid) and logit-based (Joint CBM-Logit) bottleneck. We note that we do not include other noteworthy concept-based models such as leakage-free CBMs [6], GlanceNets [33], Self-explaining Neural Networks (SENNs) [34], Concept Whitening (CW) [35], or Concept Model Extraction (CME) [16] as these either lack a well-defined mechanism to be intervened on (e.g., as in SENNs and CW) or represent more constrained versions of other baselines such as CEMs (e.g., leakage-free CBMs, GlanceNets, and CME).
Throughout our experiments, we strive for a fair comparison by using the same architectures for the concept encoder  $g$  and label predictor  $f$  of all baselines. Unless specified otherwise, for each task, we tune IntCEM's  $\lambda_{\mathrm{roll}}$  hyperparameter by varying it over  $\{5,1,0.1\}$  and selecting the model with the highest area under the curve defined by plotting the validation task accuracy vs number of interventions. For the CUB-based tasks and Ce1ebA, we use the same values of  $\lambda_{\mathrm{concept}}$  selected in [5]. Otherwise, we select this value from  $\{10,5,1\}$  using CEM's validation error. Following [5], we use embeddings with  $m = 16$  activations and a value of  $p_{\mathrm{int}} = 0.25$  for CEM's RandInt. For further details on model and training hyperparameters, see Appendix A.4.
4.1 Task and Concept Performance in the Absence of Interventions (Q1)
First, we examine how IntCEM's loss function affects its unintervened test performance compared to existing baselines as we vary  $\lambda_{\mathrm{roll}}$ . We summarise our results in Table 1, where, for all tasks and baselines, we show the downstream task predictive test accuracy, or the AUC when the task is binary (i.e., MNIST-based tasks), and the concept performance. We measure the latter via the mean AUC of predicting ground-truth test concept  $c_{i}$  from its corresponding predicted concept probability  $\hat{p}_i$ .
In the absence of concept interventions, IntCEM is as accurate as state-of-the-art baselines. Evaluating test task accuracy reveals that IntCEMs are more accurate (e.g., by  $\sim 7\%$  in CelebA) or just as accurate as other baselines across most tasks. Only in CUB do we observe a small drop in performance (less than  $\sim 1\%$ ) for the best-performing IntCEM variant. Notice that three of the five tasks we use, namely MNIST-Add-Incomp, CUB-Incomp, and CelebA, lack a complete set of training concept annotations. Therefore, our results suggest that IntCEM maintains a high task performance even with concept incompleteness – a highly desired property for real-world deployment. Finally, the performance improvements of IntCEMs over CEMs in MNIST, MNIST-Incomp, and CelebA suggest that our loss function may lead to more informative concept representations that enable more accurate label predictors to be trained.
IntCEM's concept performance is competitive with respect to traditional CEMs. The bottom half of Table 1 shows that IntCEM's mean concept AUC, particularly for low values of  $\lambda_{\mathrm{roll}}$ , is competitive against that of CEMs. The only exception is in the MNIST-based task, where Joint CBMs outperform CEMs and IntCEMs. Nevertheless, this drop is overshadowed by the poor task accuracy obtained by joint CBMs in these two tasks, highlighting that although CBMs may have better concept predictive accuracy than CEMs, their task performance is generally bounded by that of a CEM. These results demonstrate that IntCEM's explanations are as accurate as those in CEMs without sacrificing task accuracy, suggesting our proposed objective loss does not significantly affect IntCEM's performance in the absence of interventions. Finally, we observe a trade-off between incentivising the IntCEM to be highly accurate after a small number of interventions (i.e., high  $\lambda_{\mathrm{roll}}$ )
IntCEM
CEM
Figure 3: Task accuracy of all baseline models after receiving a varying number of randomly selected interventions. For our binary MNIST-based tasks, we show task AUC rather than accuracy. Here and elsewhere, we show the means and standard deviations (can be insignificant) over five random seeds.
Joint CBM-Logit
Joint CBM-Sigmoid
Sequential CBM
Independent CBM
and incentivising the IntCEM to generate more accurate concept explanations (i.e., low  $\lambda_{\mathrm{roll}}$ ), which can be calibrated by setting  $\lambda_{\mathrm{roll}}$  after considering how the IntCEM will be deployed.
4.2 Intervention Performance (Q2)
Given that a core appeal of CBM-like models is their receptiveness to interventions, we now explore how test-time interventions affect IntCEM's task performance compared to other models. We first study whether IntCEM's training procedure preconditions it to be more receptive to test-time interventions by evaluating the task accuracy of all baselines after receiving random concept interventions. Then, we investigate whether IntCEM's competitive advantage holds when exposing our baselines to state-of-the-art intervention policies. Below, we report our findings by showing each model's task accuracy as we intervene, following [4], on an increasing number of groups of mutually exclusive concepts (e.g., "white wing" and "black wing").
IntCEM is significantly more receptive to test-time interventions. Figure 3, shows that IntCEM's task accuracy after receiving concept interventions is significantly better than that of competing methods across all tasks, illuminating the benefits of IntCEM preconditioning for receptiveness. We observe this for both concept-complete tasks (i.e., CUB and MNIST-Add) as well as concept-incomplete tasks (i.e., rest). In particular, we see significant gains in CUB and CelebA; IntCEMs attain large performance improvements with only a handful of random interventions ( $\sim 10\%$  with  $\sim 25\%$  of concepts intervened) and surpass all baselines.
IntCEM's performance when intervened on with a randomly selected set of concepts can be better than the theoretically-proven optimal intervention performance achievable by CEMs alone. We explore IntCEM's receptiveness to different test-time intervention policies by computing its test performance while intervening following: (1) a Random intervention policy, where the next concept is selected, uniformly at random, from the set of unknown concepts, (2) the Uncertainty of Concept Prediction (UCP) [8] 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|$ ), (3) the Cooperative Policy (CooP) [9], where we select the concept that maximises a linear combination of its predicted uncertainty (akin to UCP) and the expected change in the predicted label's probability when intervening on that concept, (4) Concept Validation Accuracy (CVA) static policy, where concepts are selected using a fixed order starting with those whose validation errors are the highest as done by Koh et al. [4], (5) the Concept Validation Improvement (CVI) [9] static policy, where concepts are selected using a fixed order starting with those concepts that, on average, lead to the biggest improvements in validation accuracy when intervened on, and finally (6) an oracle Skyline policy, which selects  $c_*$  on every step, indicating an upper bound for all greedy policies. For details on each policy baseline, including how their hyperparameters are selected, see Appendix A.7.
We demonstrate in Figure 4 that IntCEM consistently outperforms CEMs regardless of the test-time intervention policy used, with CooP and UCP consistently yielding superior results. Further, we uncover that intervening on IntCEM with a Random policy at test-time outperforms the theoretical best performance attainable by a CEM (Skyline) under many interventions, while its best-performing policy outperforms CEM's Skyline after very few interventions. These results suggest that IntCEM's train-time conditioning not only results in IntCEMs being more receptive to test-time interventions
Figure 4: Task accuracy of IntCEMs and CEMs on CUB and CelebA when intervening with different test-time policies. We show similar improvements of IntCEMs over CBMs in Appendix A.7.
Figure 5: Task performance when intervening on IntCEMs following test-time policies  $\psi$ ,  $CooP$ , Random, and BC-Skyline. Our baseline "IntCEM no  $\psi$  " is an IntCEM whose test and train interventions are sampled from a Random policy rather than from  $\psi$  (i.e., a policy is not learnt in this baseline).
but may lead to significant changes in an IntCEM's theoretical-optimal intervention performance. We include analyses on our remaining datasets, revealing similar trends, in Appendix A.7.
4.3 Studying IntCEM's Intervention Policy (Q3)
Finally, we explore IntCEM's intervention policy by evaluating test-time interventions predicted by  $\psi$ . To understand whether there is a benefit of learning  $\psi$  end-to-end (i.e., in conjunction with the rest of model training), rather than learning it post-hoc after training an IntCEM, we compare it against a Behavioural Cloning (BC) [26] policy "BC-Skyline" trained on demonstrations of Skyline applied to a trained IntCEM. Furthermore, to explore how  $\psi$ 's trajectories at train-time help improve IntCEM's receptiveness to test-time interventions, we study an IntCEM trained by sampling trajectories uniformly at random ("IntCEM no  $\psi$ "). Our results in Figure 5 suggest that: (1) IntCEM's learnt policy  $\psi$  leads to test-time interventions that are as good or better than CooP's, yet avoid CooP's computational cost (up to  $\sim 2x$  faster as seen in Appendix A.5), (2) learning  $\psi$  end-to-end during training yields a better policy than one learnt through BC after the IntCEM is trained, as seen in the BC policy's lower performance; and (3) using a learnable intervention policy at train-time results in a significant boost of test-time performance. This suggests that part of IntCEM's receptiveness to interventions, even when using a random test-time policy, can be partially attributed to learning and using  $\psi$  during training. We show some qualitative examples of rollouts of  $\psi$  in the CUB dataset in Figure 6.
Concept Leakage Within IntCEMs Previous work [36, 6, 10] has shown that CBMs are prone to encoding unnecessary information in their learnt concept representations. This phenomenon, called concept leakage, may lead to less interpretable concept representations [36] and detrimental concept interventions in CBMs [10]. Given that interventions in IntCEMs involve swapping a concept's predicted embeddings rather than overwriting an activation in the bottleneck, such interventions are distinct from those in CBMs as they enable leaked information to be exploited after an intervention
True Label: Acadian Flycatcher
True Label: Fish Creek Predicted by Crops
Figure 6: Examples of interventions on an IntCEM following its policy  $\psi$  in CUB. We show the task label distribution  $p(\hat{y}|\hat{\mathbf{c}})$  for the most likely classes after each intervention. We highlight the correct label's probability using orange bars and show the selected concept by  $\psi$  above each panel.
is performed. This means our loss function may incentivise IntCEM to exploit this mechanism to improve a model's receptiveness to interventions. In Appendix A.10 we explore this hypothesis and find that we can detect more leakage in IntCEM's concept representations than in those learnt by CBMs and CEMs. This suggests that, contrary to common assumptions, leakage may be a healthy byproduct of more expressive concept representations in models that accommodate such expressivity. Nevertheless, we believe further work is needed to understand the consequences of this leakage.
Limitations and Future Work To the best of our knowledge, we are the first to frame learning to receive concept interventions as a joint optimisation problem where we simultaneously learn concepts, downstream labels, and policies. IntCEMs offer a principled way to prepare concept-based models for interventions at test-time, without sacrificing performance in the absence of interventions. Here, we focused on CEMs as our base model. However, our method can be extended to traditional CBMs (see Appendix A.8), and future work may consider extending it to more recent concept-based models such as post-hoc CBMs [37], label-free CBMs [38], and probabilistic CBMs [39].
Furthermore, we note some limitations. First, IntCEM requires hyperparameter tuning for  $\lambda_{\mathrm{roll}}$  and  $\lambda_{\mathrm{concept}}$ . While our ablations in Appendix A.6 suggest that an IntCEM's performance gains extend across multiple values of  $\lambda_{\mathrm{roll}}$ , users still need to tune such parameters for maximising IntCEM's utility. These challenges are exacerbated by the computational costs of training IntCEMs due to their train-time trajectory sampling (see Appendix A.5). However, such overhead gets amortised over time given  $\psi$ 's efficiency over competing policies. Further, we recognise our training procedure renders IntCEMs more receptive to any form of intervention, including adversarial interactions (see Appendix A.9). This raises a potential societal concern when deciding how our proposed model is deployed. Future work may explore such issues by incorporating error-correcting mechanisms or by considering intervention-time human uncertainty [29].
Finally, our evaluation was limited to two real-world datasets and one synthetic dataset, all of which have medium-to-small training set sizes. Therefore, future work may explore how to apply our proposed architecture to larger real-world datasets and may explore how to best deploy IntCEM's policy in practice via large user studies.
Conclusion A core feature of concept-based models is that they permit experts to interpret predictions in terms of high-level concepts and intervene on mispredicted concepts hoping to improve task performance. Counter-intuitively, such models are rarely trained with interventions in mind. In this work, we introduce a novel concept-interpretable architecture and training paradigm – Intervention-Aware Concept Embedding Models (IntCEMs) – designed explicitly for intervention receptiveness. IntCEMs simulate interventions during train-time, continually preparing the model for interventions they may receive when deployed, while learning an efficient intervention policy. Given the cost of querying experts for help, our work addresses the lack of a mechanism to leverage help when received, and demonstrates the value of studying models which, by design, know how to utilise help.
The authors would like to thank Naveen Raman, Andrei Margeloiu, and the NeurIPS 2023 reviewers for their insightful and thorough comments on earlier versions of this manuscript. MEZ acknowledges support from the Gates Cambridge Trust via a Gates Cambridge Scholarship. KMC acknowledges support from the Marshall Commission and the Cambridge Trust. AW acknowledges support from a Turing AI Fellowship under grant EP/V025279/1, The Alan Turing Institute, and the Leverhulme Trust via CFI. GM is funded by the Research Foundation-Flanders (FWO-Vlaanderen, GA No 1239422N). MJ is supported by the EPSRC grant EP/T019603/1.