Are They the Same Picture? Adapting Concept Bottleneck Models for Human-AI Collaboration in Image Retrieval

1 Introduction

Recent advances in AI have shown great promise in important, often high-risk domains like healthcare [Peiffer-Smadja et al., 2020; Rajpurkar et al., 2020], wildlife conservation [Kulits et al., 2021; Beery et al., 2019], and reducing misinformation [Mendes et al., 2022]. However, these advances are imperfect, and can lead to harm when deployed. For example, [Beede et al., 2020] reports errors when deploying AI models to detect diabetic retinopathy due to challenging realworld factors like lighting, leading to potential human harms.

Researchers have proposed human-AI collaboration as a promising approach to mitigate the shortcomings of AI models in these domains [De-Arteaga et al., 2020]. For example, prior work in health AI has sought to achieve better accuracy via decision-support tools for clinicians [Peiffer-Smadja et al., 2020] that assist humans in decision-making. Human-inthe-loop participatory systems have also been proposed to accurately and robustly categorize wildlife images [Kulits et al., 2021; Miao et al., 2021; Bondi et al., 2022] and to aid factcheckers [Nguyen et al., 2018; Mendes et al., 2022]. Concept Bottleneck Models (CBMs) are another recent and promising method to facilitate collaboration [Koh et al., 2020]. They allow humans to interact with AI models by viewing and manipulating intermediate, high-level concepts (e.g., whether a bird has a blue wing), which are then used for prediction.

While recent works using CBMs have significantly improved performance in classification [Zarlenga et al., 2023], we note that there exist many non-classification application areas that also require human-AI collaboration. For example, ElephantBook [Kulits et al., 2021] is a state-of-the-art elephant re-identification platform - meant to determine which individual elephant is depicted in each image - that adopts a semi-automated human-in-the-loop approach. This approach retrieves and presents to a user the most likely matches to individuals in a known population database for each new elephant sighting. They combine a neural network-based visual similarity ranking with a similarity ranking based on domain expert input via a weighted average (see Figure 1a). However, the weighted average approach requires hyperparameter tuning (α) depending on each user's expertise. Enabling humans to leverage concepts that they already use for retrieval to instead interact with the neural network in the embedding space has the potential to significantly reduce the amount of tuning required for each user while still including human inputs to improve team performance, as illustrated in Figure 1b).

In this work, we aim to adapt CBMs to facilitate close human-AI collaboration at deployment time for image retrieval tasks, such as ElephantBook. We specifically aim to answer the following research questions: (RQ1) How do representations generated by CBMs compare to corresponding traditional models? (RQ2) How can we augment CBMs to enable human intervention in image retrieval and classification? (RQ3) How can we train these models to incorporate varying levels of expertise? To answer these questions, we analyze how CBMs compare with traditional models for image retrieval, introduce CHAIR (CBM-Enabled Human-AI Collaboration for Image Retrieval), a novel CBMarchitecture that allows humans to intervene and encode concepts to improve retrieval, and perform extensive evaluations on how these interventions enable embedding-level human-AI collaboration.

Figure 1: Our proposed method allows human-AI collaboration in image retrieval by enabling humans to edit embeddings by correcting them through high-level concepts. We also enable flexible intervention to lower the expertise needed to participate.

2 Background andMotivation

Concept Bottleneck Models: The key promise of CBMs is two-fold: CBMs a) predict high-level intermediate concepts, which are then used to predict the final class label, thus improving interpretability, and b) enable humans to intervene and correct these intermediate concepts to improve classification performance, thus providing intervenability. This is achieved in two steps: i) Concept Bottleneck, which predicts the high-level, understandable concepts, and ii) Classifier, which predicts the final class based on the predicted concepts (see Figure 2). Human-AI collaboration is made possible here by allowing humans to correct the model by modifying these intermediate concepts. This is shown to increase the overall classification performance on various tasks [Koh et al., 2020].

Image retrieval: Image retrieval is a key component in visual tasks like image re-identification [Wang et al., 2020], wildlife conservation [Kulits et al., 2021], remote sensing [Liu et al., 2020], and visual recommendation systems [Shankar et al., 2017]. Image retrieval requires systems to fetch the most relevant images from a database given a query image, where relevance is defined depending on the application. Traditionally, applications using deep learning techniques to perform image retrieval leverage (latent) embeddings generated by neural network trained for classification. For a given query image, a query embedding is generated using the same neural network to find the top-k nearest embeddings using some distance function [Wan et al., 2014], typically cosine distance. We refer to the images and embeddings that are searched over as gallery images and embeddings, respectively. Following previous literature, we utilize the Recall@k metric and RecallAccuracy@k metrics to measure the performance of an image retrieval technique. For a given image and label pair (xi, yi), let the top-k retrieved images be y′ i ∈ Rk . Then, the Recall@k and RecallAccuracy@k are defined as follows:

The Recall@k metric evaluates if there exists at least one accurate image in the top-k images that were retrieved, whereas the RecallAccuracy@k evaluates the number of accurate images retrieved in the top-K images.

Motivation of our work: Current methods in image retrieval provide little room for human-AI collaboration. Platforms like ElephantBook, a practical, state-of-the-art computer vision system, still require a human-in-the-loop to achieve reliable performance, thus highlighting the importance of human-AI collaboration. Our contributions in this paper leverage CBMs to enable human-AI collaboration in image retrieval tasks. Previous research in CBMs until now focus largely on improving performance by proposing different architectures [Espinosa Zarlenga et al., 2022; Kim et al., 2023; Marconato et al., 2022], modify loss functions [Sheth and Ebrahimi Kahou, 2024], mitigate leakage [Havasi et al., 2022], improve intervenability [Zarlenga et al., 2023; Marcinkeviˇcs et al., 2024] or circumvent requirements of labels [Oikarinen et al., 2023]. In this work, our novelty lies in expanding intervenability capabilities of CBMs to tasks like image retrieval, which require capturing corrected concepts into the embedding. All the previously mentioned works are orthogonal to the aim of this paper and are complementary in improving the performance of CBMs.

3 Do CBMs AlreadyWork for Image Retrieval?

This section aims to establish the difficulties faced in adopting single-agent (only AI or human) approaches and show how human-AI teams, specifically through CBMs, can potentially resolve some of these issues.

Neural networks - irremediable by humans: The key drawback of using neural networks is that when these models make mistakes or inaccurate inferences, their architecture and design allow little room for intermediate human oversight, especially for tasks like image retrieval where the inaccuracies are not obvious until the results are inspected.

Expert coding - costly and high barrier of entry: Handcrafted features and coding of retrieval systems require advanced domain knowledge and experience, which raises the barrier of entry for inexperienced users to utilize such systems. Furthermore, these codes often require entire feature sets to achieve perfect retrieval. Any ambiguity in identifying these codes, for example, due to ambiguity in or occlusion of concept-relevant parts of the image, is marked as Wildcards in such systems. Since these codes typically lie on a spectrum of easily identifiable to years of expertise needed, we can leverage neural networks to alleviate some of these difficulties while having a variable level of human oversight and help reduce wildcard entries with a collaborative human-AI retrieval system.

Are CBMs the solution? Potentially: Our contributions stem from the observation that high-level concepts are analogous to the codes developed by domain experts to help humans retrieve relevant images. However, adopting these concepts (in the case of CBMs) or expert codes (in the case of humans) to enable collaborative embedding generation for image retrieval is not supported by existing CBM architectures (see Table 1).

CBMs typically contain a linear layer as a classifier, and the predictions change as a function of predicted concepts and human interventions. Therefore, a straightforward way of achieving intervenability for retrieval would be to extract the latent embeddings for comparison with other images after the concept predictions by adding another linear layer before predicting the final class. This enables capturing human intervention in the embedding space. However, this extension, termed as CBM-Extend (illustrated in Figure 2), performs poorly in image retrieval when compared to the latent embeddings used from standard, CBM and HybridCBM [Mahinpei et al., 2021] embeddings with similar base architectures (ResNet-18). (see Figure 3). The significant drop in performance (Recall@k here, as defined in equation 1) is likely due to the lack of generalization with an increasing number of hidden layers (increased capacity), rendering the naive extensions unusable. While section 5 details what Recall@k signifies here, it is sufficient to say that there exists a clear gap in abilities and performance, thus giving a clear picture of how representations generated by CBMs compare to corresponding traditional models ((RQ1)).

4 Our Proposed Architecture: CHAIR

We have established the importance of enabling human-AI collaboration for tasks beyond classification. Keeping in mind (RQ2), which asks "How can we augment CBMs to enable human interventions on the representations" and (RQ3) that asks "How can we enable different levels of expertise for intervention", we now present CHAIR, a two-stage CBM-like architecture that helps address these questions. This section is organized into two parts: (a) Architecture and (b) Training.

Architecture: Figure 4 illustrates our proposed CHAIR architecture. Let ζ indicate the encoder that generates embeddings z from the input image x. ϕ(x) denotes the concept head that generates the concept vector c from the input image x and ψ(c) denotes the classification head that outputs the final class label y. The Concept Bottleneck comprises of the encoder and the concept head, that is ϕ(ζ(x)). The final class prediction in a CBM is as follows:

We propose a simple two-stage modification to the vanilla CBM architecture that enables integrating human interventions to improve retrieval. Specifically, as illustrated in Figure 4, our architecture introduces a Fusion Head, whose output is an additional edited embedding alongside the predicted class. The Fusion Head comprises a concept-toembedding projection layer - a linear layer (ω) that projects the predicted concepts into the same dimensional space as the embedding (z), and a classifier. To incorporate the concepts into the embedding, we add the projection and the previous embedding to give us an edited embedding, similar in spirit to a residual connection [He et al., 2016], thus the name Fusion Head. The objective of this fusion is to be able to learn a meaningful and better representation when these concepts are corrected, which helps us address (RQ2), that is, enabling concept encoding onto the generated embeddings. Lastly, the classifier (ψ) utilizes this edited embedding for the final classification, thus preserving the same ability as CBMs to incorporate interventions for improving classification.

Training: Given we have an architecture capable of incorporating human intervention in the latent space for image retrieval and classification, we now introduce a two-stage training that learns better embeddings with intervention while maintaining classification accuracy. Note that we adopt the standard classification loss function for image retrieval as used in [Sharif Razavian et al., 2014], which is shown to extend to tasks beyond classification. Furthermore, this helps us understand how our architecture performs against similar models trained on the same loss functions.

Stage 1: From [Sharif Razavian et al., 2014], we know that the embedding z obtained from training on cross-entropy like loss functions provides a good baseline embedding to enable image retrieval. Leveraging the fusion head architecture, the learned edited embedding encodes information from the already performant embeddings (z) and concepts. Mathematically, the fusion head consists of the ϕ(z) to generate concepts c. c is then projected to the same space as (z) using a simple linear layer ω(c). Consequently, this projection is added to the original embedding (z) to generate the final edited embedding z′ . Our proposed fusion head then utilizes a different classification head ψ′ (z′ ) to predict the final class label y′ as follows:

The goal of Stage 1 is to train the concept-to-embedding projection layer ω to enable learning better edited embeddings z′ , which improves classification and image retrieval.

Stage 2: Having now trained the concept-to-embedding projection layer, our model is capable of incorporating corrections in the latent space. Referring back to (RQ3), which aims to include different levels of expertise, we posit that Stage 1 training alone is insufficient to achieve this ability. Hence, we introduce this stage that performs random interventions on a select number of concepts for each mini-batch during training. This random correction of concepts simulates varying levels of expertise while training the edited embedding to learn better representations under partial interventions. Note that the concept head here is frozen to preserve the activation values that signify presence or absence of a concept, similar to [Koh et al., 2020].

Intuition: Revisiting the requirements of a good humanAI collaborative model, the reasoning behind our contribution is three-fold: (a) our architecture helps us retain the same abilities of CBMs to perform classification, (b) our Stage 1 training allows training ω, the concept-to-embedding projection layer, to enable incorporating concepts into the edited embedding (addressing (RQ2) and (c) Stage 2 training allows for learning quality embeddings under variable intervention, thus addressing (RQ3).

We adopt the two training modes outlined in [Koh et al., 2020]: Sequential (shortened as Seq hereafter), where the concept and classification heads are trained separately, and Joint where the heads are trained all at once. Note that these training modes mainly differ in Stage 1 training in the loss function. Algorithms 1 and 2 outline training our proposed CHAIR model for both modes. More specifically, class loss and concept loss correspond to the

Algorithm 1 Training CHAIR model

Input: Training data: D = [(x, c, y)], mode Initialize: ϕ, ω, ψ, ζ =pre-trained,pint Stage-1 for (xi, ci, yi) in D do zi ← ζ(xi) ▷ Get Embedding c′i ← ϕ(zi) ▷ Get Concepts loss = concept loss(ci, c′i ) if mode == "sequential" then z′i = zi + ω(c′i ) ▷ Get edited embedding y′i = ψ(z′i ) loss += class loss(yi, y′i ) loss.backward() Stage-2 Freeze: ζ, ϕ if mode == "sequential" then reset weights(ω) cint ← intervention values(ζ, ϕ, D) for (xi, ci, yi) in D do pi ← torch.rand(1) ▷ Partial intervention ˆ ci ← concept intervention(zi,ϕ,ci,pi,cint) z"i ← ω(ˆ ci) + zi y"i ← ψ(z"i) loss = class loss(yi, y"i) loss.backward() return ζ, ϕ, ω, ψ

same loss functions utilized in the original CBM training. While we observe the same performance pattern for the sigmoid activation function as noted by [Koh et al., 2020], all of the results we report here are from using the ReLU activation and cross-entropy loss for class loss and individually for each concept loss. Furthermore, the intervention values function enables calculating the activation values for concepts that indicate their presence or absence. We utilize the training data to calculate these values, similar to [Koh et al., 2020].

Intervention: Intervention in Stage 2 is performed by sampling from a uniform distribution, since we do not assume any level or category of expertise from the humans in the human-AI team. Therefore, all evaluations with interventions also assume expertise on a percentage of the available concepts sampled uniformly.

5 Results

5.1 Datasets and Evaluation

We conduct experiments on two real-world datasets (similar to [Espinosa Zarlenga et al., 2022]), the Caltech-UCSDBirds-200-2011 (CUB) dataset [Wah et al., 2011] and Largescale CelebFaces Attributes dataset (CelebA) [Liu et al., 2015] to demonstrate the effectiveness of our proposed architecture. Specifically, we utilize the CUB dataset to measure performance in both classification and image retrieval tasks, while the CelebA dataset is used for classification only. The CUB dataset comprises n = 11, 788 bird images belonging to a total of 200 possible species. There are 112 binary concepts

Algorithm 2 Stage 2 Intervention functions

Function intervention values(ζ, ϕ, D): Initialize: cmax int = [], cmin int = [] x, c, y = D ▷ Training data z ← ζ(x) c′ ← ϕ(zi) for c′i in c do ▷ activations when concept is present and absent cmax inti ←top-95-percentile(c'i) cmin inti ←bottom-5-percentile(c'i) cint = [cmax int , cmin int ] return cint Function concept intervention(z, ϕ, c, p, cint): c′ ← ϕ(z) idx = torch.randperm(c)[: p ∗ len(c)] cmax = c int[0][idx] ▷ Simulates partial correction cmin = c int[1][idx] c′ [idx] = torch.where(c[idx] == 1, cmax, c′ ) c′ [idx] = torch.where(c[idx] == 0, cmin, c′ ) ˆ c ← c′ return ˆ c

associated with each image. We follow the same experimental setup for the classification task as detailed in [Koh et al., 2020], where the dataset is divided into training, validation, and testing. Furthermore, we follow the data split established in the literature for image retrieval. The training data here consists of bird images from the first 100 classes. The images from the next 100 unseen classes are then used to create the gallery, which is used to retrieve images and measure performance. In contrast, each image in the CelebA dataset consists of 40 concept labels, and we utilize 1000 classes for this classification task. We note that the CelebA dataset was not used in evaluating retrieval, as previous works measure retrieval based on binary concepts, which deviates from the main target application of our work. We use ResNet-18 [He et al., 2016] pre-trained on ImageNet [Deng et al., 2009] for all experiments and train on a single NVIDIA Tesla V100 16GB GPU. 1

5.2 Retrieval Performance

Figure 5 shows how CHAIR performs against a standard ResNet-18 model and CBM. A typical CBM with the same underlying ResNet-18 architecture performs slightly worse than the standard ResNet model, which aligns with the behavior noted by [Koh et al., 2020] in classification. Our model provides 15-20 % Recall@k improvement for k = [1, 5, 10], indicating that our model has learned better-unedited embeddings as opposed to the embeddings from the standard and CBM models. Furthermore, we randomly select p% of concepts for every image in the gallery and each image during query to investigate how intervention improves retrieval and evaluate the Recall@k performance, as shown in Figure

1 Code (written in Python using PyTorch [Paszke et al., 2019]) and instructions to reproduce these results can be found here: https://github.com/realize-lab/CHAIR

6. We observe that as the number of interventions (corrections) increases, the retrieval performance improves, eventually reaching near-perfect Recall@k. A key observation that emerges from this plot is that when all the predicted concepts are set to their correct values, the method achieves nearly perfect recall when k = 10.

In the real world, it is possible that the labelers constructing the gallery embeddings could be less experienced than the users during query time (especially when the gallery is crowdsourced). Hence, we demonstrate how different levels of intervention in the query and gallery images impact the accuracy of our proposed retrieval architecture. Figure 8 shows a heatmap of % interventions on query images and gallery images on the x and y axes, respectively, and the RecallAccuracy@10 value for each intervention pair. We observe that even when the gallery is constructed with no intervention, any amount of human intervention during query time helps improve retrieval performance. Notably, this trend emerges across all interventions during query time when the interventions of the gallery embeddings are kept fixed. Note that we do not plot a similar heatmap for Recall@k, since our method already performs well without any intervention.

5.3 Importance of Stage 2

Recall that Stage 2 was introduced to help improve performance under partial intervention. To measure the impact of Stage 2, we calculate the recall values of both the training modes (Seq and Joint, as defined in Section 4) perform with and without Stage-2 training for each subset of intervention in Figure 9. This figure clearly shows the benefit acquired when Stage 2 training is employed, especially in the 0-80% range, where we obtain 5%-25% improvement in recall depending on the intervention. Note that the Seq. performance here without Stage 2 implies the classification head is trained on predicted concept activation values without any intervention.

5.4 Quality of Edited Representations

In order to visualize the representations at each intervention level, we use t-Distributed Stochastic Neighbor Embedding (t-SNE), a dimensionality reduction technique that helps visualize high-dimensional data, which in our case are the embeddings that are used for retrieval. We extract embeddings for each image present in our test data (all the classes in the

-50

test data are unseen during training) and apply t-SNE to reduce them to 2 dimensions. We then plot these reduced embeddings, as shown in Figure 7 for different levels of intervention. With accurate intervention on increasing random subsets of the concepts, the clusters for each class become more distinct. These visual representations also help provide evidence to support how the recall performance under no intervention is much lower when compared to performance under intervention.

5.5 Classification Performance

As referred to in (RQ2), our model must be able to match the performance of CBMs in classification and allow intervenable retrieval so that practitioners do not have a tradeoff when selecting the architectures. Figure 10 compares the classification performance of both sequential (Seq) and joint training modes under perfect concept intervention against the best-performing classifying CBM on both the CUB and the CelebA datasets. It is evident that our proposed CHAIR model outperforms the standard best-performing CBM in these real-world datasets. Furthermore, this plot shows how Stage 2 helps improves classification performance, thus showing its importance beyond improving retrieval under partial intervention.

6 Conclusion and FutureWork

In this work, we first establish that CBMs, a model that allows for human-AI collaboration, underperform when compared to standard neural networks on retrieval tasks. We propose CHAIR, a modification of both CBM architecture and training strategy, which can a) incorporate human input in the form of concept correction for image retrieval, b) allow varying levels of human input and expertise, and c) significantly improve retrieval performance while maintaining similar classification performance to vanilla CBMs. Furthermore, we show that the quality of the embeddings generated with some corrections performs better than the alternatives through t-SNE plots and improved retrieval performance. Finally, it is evident that while both the Seq and Joint training modes perform equally well with a high level of intervention, we recommend choosing the joint training mode for image retrieval-related tasks. Our work enables the expansion of CBMs to domains that move beyond classification, further improving human-AI collaboration in impactful applications such as wildlife population monitoring.

We envision future work in this area along multiple frontiers: a) incorporating probabilistic concepts, label prediction, and embeddings to capture the uncertainty of the prediction process better [Kim et al., 2023; Li et al., 2021], b) learning when and how to defer to humans that are interfacing with CBMs to achieve the best possible complementary performance [Bondi et al., 2022; Mozannar and Sontag, 2020], and c) conducting human studies to understand effective strategies of presenting CBM information to achieve complementarity.

Ethical Statement

Our contributions and future work involve including and complementing human efforts. Hence, we strongly advocate rigorously testing with all stakeholders before deploying CHAIR-like models.

Acknowledgements

We want to thank the reviewers for their insightful comments and suggestions. We also thank the e-Health and Artificial Intelligence (e-HAIL) initiative and the Advanced Research Computing (ARC) at the University of Michigan, Ann Arbor, for their support and for providing computing resources and services, respectively.