Bayesian Concept Bottleneck Models with LLM Priors

We thank the reviewer for their thoughtful comments and are heartened to hear that you find the method quite neat! This work was inspired by our real-world collaboration with the local hospital to build risk prediction models, where there is a major need for methods that learn fully interpretable clinical risk prediction models for unstructured data. We believe BC-LLM, by bringing together the world knowledge of LLMs and the rigor of Bayesian inference, has the potential to address this critical gap. Thus we sincerely appreciate the reviewer's feedback on how the manuscript can be further improved and address the reviewer's points below:

No thorough analysis is conducted on the LLM extracted concept. It is unclear how interpretable or trustworthy these concepts actually are.

Thank you for this question! We can use the semi-synthetic MIMIC experiment to check how trustworthy the extracted concepts are, because we have the human annotations for the true concepts. In particular, we have human annotations for “Does the note mention the patient consuming alcohol in the present or the past?”, “Does the note mention the patient smoking in the present or the past?”, “Does the note mention the patient using recreational drugs in the present or the past?”, “Does the note imply the patient is unemployed or retired?”. BC-LLM converged to similar concepts with the following agreement rates with the human annotations:

• For the alcohol concept, BC-LLM learned concept questions like “Does the note mention this patient having a history of alcohol use disorder?” and “Does the note indicate that the patient consumes alcohol?”. The agreement rate of LLM extractions with human annotations is 91.5%.
• For the tobacco concept, BC-LLM learned concept questions like “Does the note mention the patient’s smoking history?” and “Does the note mention the patient as currently smoking?”. The agreement rate of LLM extractions with human annotations is 85.9%.
• For the recreational drugs concept, BC-LLM learned concept questions like “Does the note mention the patient having a history of cocaine use?”. The agreement rate is 95.1%.
• For the unemployment/retired concept, BC-LLM learned concept questions like "Does the note mention the patient being retired?" and "Does the note mention the patient as unemployed?". The agreement rate is 73%.

So in this experiment, we found that the LLM extracted concepts are quite trustworthy. In manual reviews for other experiments, we also found the extracted concepts to be quite accurate. That said, in settings where the LLM cannot extract certain concepts accurately, we expect BC-LLM to converge towards concepts that it can accurately extract and are predictive of the label.

In the iterative procedure of generating new concepts, in each iteration, all training data are used. This can be highly inefficient once the data set gets larger and/or the iteration becomes large.

Because the primary goal of this manuscript is to introduce a Bayesian formalization for learning CBMs, there is still room for improvement when it comes to the scalability of BC-LLM. Currently, the runtime BC-LLM is linear in the number of observations since it uses all the training data in each iteration. A straightforward way to speed this up is by using mini-batches instead. Our preliminary experiments show that BC-LLM achieves similar results when doing this for the greedy warm-up procedure. We are now speeding up the Gibbs sampling procedure with mini-batching while preserving statistical guarantees, using methods that estimate the MH acceptance ratio (see e.g. Seita 2018 and Bardenet 2014). We believe this is beyond the scope of this paper and thus leave it to future work.

Calibration is measured by Brier, which seems to me is an odd choice. ECE is a more standard measurement of calibration in the literature?

Both Brier and ECE are standard measurements of calibration. We had selected Brier because it is also a proper scoring rule, so the best possible model will also attain the lowest Brier score. ECE tends to be very noisy for smaller dataset sizes since the ECE is highly sensitive to binning strategy; in addition, inaccurate models can also attain low ECE values.

In the case of the semi-synthetic MIMIC experiment, the true probabilities are known, so the Brier score–defined here as the mean squared error between the true and predicted probabilities–is exactly what we want. As shown in the results, BC-LLM achieves a lower Brier score than other methods.

For the experiments where the true probability is not known, all the methods have similar ECE values (using 5 bins). For CUB, we find that the BC-LLM attains an ECE of 0.25 with a 95% CI of (0.17, 0.32). This overlaps with the comparators LLM+CBM (ECE=0.23), Boosting LLM+CBM (ECE=0.24), Human+CBM (ECE=0.26), LLM+CBM (no keyphrases) (ECE=0.31), ResNet (ECE=0.26). For the readmission risk prediction experiment in Section 4.3, we find that BC-LLM attains an ECE of 0.184 with a 95% CI of (0.176, 0.193), which again overlaps with comparators LLM+CBM (ECE=0.182), Boosting LLM+CBM (ECE=0.180), and Bag-of-Words (ECE=0.177).

I still don't fully understand how to use the proposed method to make prediction on a new data point. I think we need to first let the LLM annotate the new data point based on learnt concepts, but then how do you make prediction? Also, how do you account for the error that might occur during the LLM annotation?

Correct – To make a prediction on a new data point, the LLM annotates it with the learnt concepts. Each CBM in the ensemble (i.e., the posterior samples) then outputs a prediction per its fitted parameters. The final prediction is an average across the ensemble.

BC-LLM does not additionally adjust for annotation error. Instead, it relies on the Bayesian inference procedure to implicitly do this adjustment. That is, concepts that the LLM cannot annotate accurately will have a low MH acceptance rate because these concepts will not be predictive of the target in held-out data. Instead, BC-LLM will converge towards concepts that both (i) the LLM is able to correctly annotate and (ii) are predictive of the target.

The writing could be improved.

Thanks for this feedback. We will revise the notation to use only bold symbols for vectors, add a notation table for clarity, and clean up writing wherever possible to make the paper easier to follow.

We thank the reviewer for their detailed comments and for recognizing that our method is both novel and well-motivated. Before addressing each of the reviewer’s points, we believe it would be fruitful to highlight that the goals of this work significantly differ from prior work. The goal is to learn a fully-interpretable non-leaky CBM in the “low-knowledge regime” where it is difficult to list potentially relevant concepts upfront. This is a major unaddressed need in healthcare, as clinical providers often want fully interpretable, faithful models that can predict (often noisy) outcomes using the unstructured data and have limited insight into which candidate concepts are good to include. In fact, BC-LLM is inspired by our real-world collaboration with Anonymous hospital to build a risk prediction model, which is the experiment in Section 4.3 of the paper. As illustrated in this experiment, existing CBM methods do poorly because they are not designed for the low-knowledge regime. In contrast, BC-LLM does much better.

Choice of experiments: The reason we do not copy existing experiments from other CBM works is that they are not designed to test the ability of a CBM to learn concepts in low-knowledge regimes. Existing works test the setting where the LLM already has good knowledge of the prediction task and is able to list a comprehensive list of candidate concepts. In contrast, our experiments are designed to simulate the low-knowledge regime as follows:

• In the MIMIC example, we have real-world clinical notes and simulate the label Y based on a hand-annotated set of concepts. Because the label Y is simulated and we only tell the LLM the outcome is “some label Y”, the LLM is unlikely to know what set of concepts we actually used in the simulation. Furthermore, using this semi-synthetic setup, we can evaluate if BC-LLM converges to the known true concepts.
• In the readmission risk prediction example based on real-world clinical notes, we use a dataset that is private to the Anonymous hospital, thereby eliminating the possibility of an LLM having seen the data. Furthermore, readmission is known to be very difficult to predict, and factors influencing readmission are known to be highly variable across different hospitals. The factors picked up by BC-LLM are indeed uniquely relevant to Anonymous hospital, where a substantial subgroup of patients exhibit substance use.
• For the CUB example, we designed the task to be slightly different from prior setups to minimize the chance of the LLM already knowing the answer. To do this, we do not reveal to the LLM that the task is to predict bird species and change the prediction task. Rather than predicting one of 200 bird species, the task is to predict a bird species within its bird family (a set of closely related bird species, e.g. a Bunting bird family is composed of three different Bunting bird species). Furthermore, predicting within bird families is better aligned with the fact that (i) our primary goal is to learn highly interpretable CBMs with fewer concepts and (ii) there likely exists a smaller CBM for distinguishing bird species within a family. In contrast, a CBM with a large number of concepts is no longer easily interpretable.

Additional experiments, including OOD performance: We have added the following

• Original CUB-200 dataset
• Functional Map of the World (fMoW) (Christie et al., 2018): RGB satellite images classified into 62 building/land use categories. Models were trained on satellite images from the US, with 100 images sampled from each category. CBMs were trained to have 60 concepts. Models were evaluated in terms of in-distribution performance (i.e. images from the USA) as well as out-of-distribution performance when evaluated on satellite images from China. As another reviewer mentioned, these satellite images may require more domain expertise than the CUB dataset included in the original submission.
• Imagenette: 10 classes from ImageNet. CBMs were trained to have 10 concepts.

The table below reports the accuracy of BC-LLM as well as comparators. As we see, BC-LLM outperforms the CBM methods and even outperforms ResNet on the fMoW dataset.

Adequate comparison to SoTA CBM approaches: To our knowledge, the paper includes all major non-leaky CBM comparators: (A) Fitting a CBM on human-selected and annotated concepts (referred to as “Human+CBM”), (B) Using an LLM to suggest and annotate concepts and then fitting a variable selection method on these concepts (referred to as “LLM+CBM (no keyphrases)”), (C) Using an LLM to brainstorm keyphrases that describe each observation, fitting a sparse model on these keyphrases, and having the LLM suggest concepts based on the top keyphrases (“LLM+CBM”), and (D) Using an LLM to suggest concepts in a boosting procedure (“Boosting LLM + CBM”). Note that we do not include CBM methods that rely on “soft” (i.e. continuous-valued) concept extractors, as they can suffer from information leakage and thus are not fully interpretable.

Interpretability evaluation: We appreciate these concerns and will ask for input from more clinicians for the revised manuscript. That said, we would like to highlight that the average relevance score for BC-LLM is >1 standard error higher than that for the other methods. Getting input from clinicians is rare for AI/ML conference papers, and having input from four clinicians is substantially more than most. We will include error bars to Fig 4 in the revised manuscript.

Rigorous statistical inference and uncertainty quantification: BC-LLM is supported by rigorous theoretical analyses, faster convergence to the correct concepts, lower Brier score, and higher entropy in OOD settings. We have also demonstrated better performance in OOD settings in the additional experimental results shown above.

Computational cost: The average wall times for BC-LLM are: 1 minute per epoch for Section 4.1 (short clinical notes), 3 minutes per epoch Sec 4.2 (CUB), and 1.2 minutes per epoch for Section 4.3 (long clinical notes). Each BC-LLM experiment was run for 5 epochs. The cost for BC-LLM is around \$0.0015 per request for CUB. On average, this translates to around \\$3 per experiment. The cost for MIMIC is around \$0.0005 per request, which translates to around \\$1 per experiment. We believe these costs are highly reasonable. We note that this is all without requiring human annotations for the concepts, which is significantly more costly. Furthermore, the concepts in BC-LLM are not extracted through embedding similarity-based measures, which are known to be leaky and more error-prone.

Ablation studies to study sensitivity to LLM dependence: We have investigated sensitivity to the choice of the LLM in Section D of our appendix, where we show results from using BC-LLM with Cohere’s Command-R and Claude 3.5 Haiku instead. We consistently find that BC-LLM outperforms comparator methods. We used a temperature of zero to minimize stochasticity and default parameters otherwise. We used very basic prompts (Section C of our appendix) to reflect the “low-knowledge” regime. The revised manuscript will explore how performance changes if we conduct further prompt engineering.

Related works: We have additionally added some discussion on variational methods such as [2]

We hope our responses clarify any points of confusion and demonstrate the novelty and importance of this work. If possible, we kindly ask the reviewer to reassess their score.

$1$ Ragkousis et al. "Tree-Based Leakage Inspection and Control in Concept Bottleneck Models" arXiv:2410.06352

$2$ Panousis et al. "Coarse-to-Fine Concept Bottleneck Models" NeurIPS 2024

We thank the reviewer for their thoughtful comments and for recognizing the novelty of this Bayesian framework for learning CBMs. This work was inspired by our real-world collaboration with the local hospital to build risk prediction models, where there is a major need for methods that learn fully interpretable clinical risk prediction models for unstructured data. We believe BC-LLM, by bringing together the world knowledge of LLMs and the rigor of Bayesian inference, has the potential to address this critical gap. Thus we sincerely appreciate the reviewer's feedback on how the manuscript can be further improved and address the reviewer's points below:

Generalizability to Other Image Domains and Applicability to Domain-Specific Tasks: Can this approach be applied to different image classification datasets beyond those used in the paper? How would this method perform on tasks that require highly domain-specific knowledge, where LLMs may lack sufficient prior understanding, such as satellite imagery or industrial/factory imaging?

Thank you for these two questions about the robustness of BC-LLM to other real-world image datasets! To illustrate the generalizability of BC-LLM to other image domains, we have run experiments on two additional datasets:

• Functional Map of the World (fMoW) (Christie et al., 2018): RGB satellite images classified into 62 building/land use categories. Models were trained on satellite images from the US, with 100 images sampled from each category. CBMs were trained to have 60 concepts. Models were evaluated in terms of in-distribution performance (i.e. images from the USA) as well as out-of-distribution performance when evaluated on satellite images from China. As the reviewer mentioned, these satellite images may require more domain expertise than the CUB dataset included in the original submission.
• Imagenette (https://github.com/fastai/imagenette): 10 classes from ImageNet. CBMs were trained to have 10 concepts. Note that we do not run on all of ImageNet, since BC-LLM is not intended for learning CBMs with a very large number of classes.

The table below reports the accuracy of BC-LLM as well as comparators (we exclude human+CBM because there are no available human-generated concepts for these datasets). As we see, BC-LLM outperforms the CBM methods and even outperforms ResNet on the fMoW dataset.

While these results are promising, we acknowledge that for highly specific problem domains, BC-LLM may require some tailoring. For instance, rather than using a general LLM as we do in the experiments, one may instead need to use a domain-specific LLM (e.g. a medical LLM). Furthermore, if there is no LLM that can accurately extract values for domain-specific concepts in a zero-shot manner, one may need to train/tune a model specifically to extract domain-specific concepts.

Prompt Design for Concept Annotation: Could you provide the exact prompt used in Step 3 for concept annotation in the bird image classification task (Appendix)? How are concept annotations generated across different modalities?

Here is the prompt used for concept annotation (Step 3) in the bird image classification task. The image is uploaded to the LLM through the API alongside the prompt

"You will be given an image. I will give you a series of questions. Your task is to answer each question with 0 or 1. You may output a probability if you are unsure, but avoid this in general. Summarize the response with a JSON that includes your answer to all of the questions. Questions:

1. <CONCEPT QUESTION 1>

2. <CONCEPT QUESTION 2>

3. …

Output your answer in JSON"

We fill in the prompt with the concept questions for which we want annotations.

Entropy in OOD Settings: In Table 1, the ResNet model shows the highest entropy in the out-of-distribution (OOD) setting. Could you clarify why this might be the case?

ResNet can achieve higher entropy than CBM models because CBMs are highly constrained: CBMs have only a small number of input features by design, and the values of these input features are further constrained to be interpretable (e.g., a probability). As a consequence, a CBM may not be able to fully represent the additional uncertainty in OOD settings. Nevertheless, because a CBM is easy to audit by humans, it is also easier to debug than a ResNet model. By simply observing the extracted concept values, one is likely able to detect whether an observation is OOD. Furthermore, as shown in the new experiment with fMoW, we find that BC-LLM can outperform ResNet in terms of OOD performance in certain cases.

Related Work Suggestion: You may want to cite the paper "Constructing Concept-based Models to Mitigate Spurious Correlations with Minimal Human Effort" (ECCV 2024), which also uses LLMs for concept candidate generation, filtering, and annotation. It seems closely related and could help contextualize your contribution.

Thank you for this suggestion – we have added a discussion on it to our related works section. The paper highlights studying correlated concepts and is very similar to the “LLM+CBM” baseline we compare against. The iterative Bayesian framing we propose provides an extra layer of uncertainty quantification and enables finding more concise models than the method proposed in the suggested paper.

We thank the reviewer for the thoughtful comments. We are heartened to hear that you found the Bayesian framework to be interesting and nontrivial and the experiments to be convincing. This work was inspired by our real-world collaboration with the local hospital to build risk prediction models, where there is a major need for methods that learn fully interpretable clinical risk prediction models for unstructured data. We believe BC-LLM, by bringing together the world knowledge of LLMs and the rigor of Bayesian inference, has the potential to address this critical gap. Thus we sincerely appreciate the reviewer's feedback on how the manuscript can be further improved and address the reviewer's points below:

Unclear implication of assumption H.1. It would be great to include a high-level summary of the justification for this.

Assumption G.1. is used to prove Theorem 3.1, which states that MH sampling with a potentially misspecified LLM prior and proposal distributions can still be robust and converge to a desired target distribution. This result is highly unique, because most works on Bayesian analysis analyze an explicitly defined posterior and do not consider the effect of misspecification in the sampling process on the target distribution. Thus additional assumptions are required. Most components in Assumption G.1 are either standard regularity conditions ((i), (ii), (iv), (v), (vi), (viii), (ix)) that are often made in statistical analysis or are hyperparameter choices ((iii) and (x)). The key new assumption is (vii), which asserts that at each step, there is at least a fixed nonzero probability of the LLM proposing one of the K “best” concepts conditional on the data. This helps to ensure at every outer loop of the sampler (i.e. for every cycle through all the positions of the concept vector), enough mass flows to the optimal set of concepts. We believe this is a reasonable assumption even in regimes where the LLM has limited knowledge, as long as the degree of knowledge is non-zero.

What if all the initial concepts suggested by the LLM are not related to the task? As with all metropolis hastings like schemes, the sampling complexity here might be high.

Thank you for this question! First, we would like to highlight that this work was motivated by this very concern you have raised. Most existing methods ask LLMs for a set of candidate concepts or rely on a pre-existing concept list, but this is not effective when LLMs or humans have limited prior knowledge of which concepts are truly relevant. BC-LLM addresses this by iteratively refining the set of concepts during posterior inference.

Second, we agree that if an MH sampling procedure is initialized poorly, sampling complexity will be high. Because of this very concern, the initialization and warm-up steps of BC-LLM are specifically designed to accelerate convergence. For initialization, the LLM explores a large variety of concepts by brainstorming keyphrases that describe each observation and then selects concepts based on the keyphrases that are most predictive. For warm-up, the LLM greedily chooses concepts that improve prediction performance, again based on keyphrases that are most associated with the residual. Together, these two steps ensure the MH sampling procedure starts closer to the posterior.

Scalability to a large number of concepts (K) is a challenge… How do you expect the computational cost to change with increasing concept size?

Because the primary goal of this manuscript is to introduce a Bayesian formalization for learning CBMs, there is still room for improvement when it comes to the scalability of BC-LLM. Currently, the runtime of BC-LLM is linear in the number of concepts K, because each iteration of Gibbs refines/replaces a single concept. A straightforward way to speed this up is by refining J>1 concepts in each iteration, which reduces the runtime by a factor of J. Our preliminary experiments show that BC-LLM achieves similar results when doing this for the greedy warm-up procedure. However, extending the Gibbs sampling procedure to refine/replace multiple concepts in a single iteration while preserving statistical guarantees is beyond the scope of this paper, and we leave it to future work.

Can the LLM hallucinate concepts valid for training, but not relevant for test? How should the uncertainty adapt in this case?

During the posterior sampling procedure, it is certainly possible for the LLM to propose irrelevant concepts (e.g., because it fits the training data well). Nevertheless, for a sufficiently large dataset, irrelevant concepts are likely to be rejected since they don’t fit the held-out data well. For small datasets where there is insufficient sample size to distinguish between relevant and irrelevant concepts, irrelevant concepts are more likely to be accepted, but the posterior sample will also reflect this increased uncertainty by including a larger variety of concepts. Indeed, we observe such behavior in Figure 3 (right): the posterior samples for the smaller dataset reflect high uncertainty (many unique concepts) while the posterior samples for the larger dataset reflect low uncertainty (few unique concepts).