V-CEM: Bridging Performance and Intervenability in Concept-based Models

TL;DR

V-CEM addresses the interpretability-performance tension in concept-based models by marrying the strengths of concept embeddings with variational inference. It introduces a probabilistic framework where concept embeddings are generated via an approximate posterior q(\mathbf{c}|x,c) conditioned on both inputs and concept predictions, while a prior p(\mathbf{c}|c) enforces concept-specific clustering; this yields an ELBO-based objective L = (1/k)L_c + \lambda_t L_t + \lambda_p L_p that balances concept accuracy, task accuracy, and embedding regularization. The model enables targeted concept-embedding interventions that can override input-dependent effects, improving OOD intervenability while preserving strong ID performance comparable to black-box models and CEMs. Experimental results across vision and NLP tasks demonstrate V-CEM’s cohesive embedding space (CRC), robust OOD intervention responsiveness, and competitive ID accuracy, suggesting practical benefits for reliable, human-guided AI systems. The work also introduces metrics for evaluating concept representations and discusses avenues for extending V-CEM to multimodal data and causal modeling to further enhance interpretability and robustness.

Abstract

Concept-based eXplainable AI (C-XAI) is a rapidly growing research field that enhances AI model interpretability by leveraging intermediate, human-understandable concepts. This approach not only enhances model transparency but also enables human intervention, allowing users to interact with these concepts to refine and improve the model's performance. Concept Bottleneck Models (CBMs) explicitly predict concepts before making final decisions, enabling interventions to correct misclassified concepts. While CBMs remain effective in Out-Of-Distribution (OOD) settings with intervention, they struggle to match the performance of black-box models. Concept Embedding Models (CEMs) address this by learning concept embeddings from both concept predictions and input data, enhancing In-Distribution (ID) accuracy but reducing the effectiveness of interventions, especially in OOD scenarios. In this work, we propose the Variational Concept Embedding Model (V-CEM), which leverages variational inference to improve intervention responsiveness in CEMs. We evaluated our model on various textual and visual datasets in terms of ID performance, intervention responsiveness in both ID and OOD settings, and Concept Representation Cohesiveness (CRC), a metric we propose to assess the quality of the concept embedding representations. The results demonstrate that V-CEM retains CEM-level ID performance while achieving intervention effectiveness similar to CBM in OOD settings, effectively reducing the gap between interpretability (intervention) and generalization (performance).

Figures (7)

• Figure 1: Probabilistic Graphical Models of a) CBMs, b) the CEMs, and c) the proposed V-CEM architecture. Solid lines represent the data generation process, while dotted lines represent inference.
• Figure 2: Illustration of the V-CEM architecture. Given an image of a parrot with a red breast, V-CEM concept encoder $p(c|x)$ assigns a high probability to the "Red Breast" concept and a low probability to "Blue Feathers", which is absent. The approximate posterior $q(\mathbf{c}|x,c)$ maps concept prediction to concept embeddings clustered around $\mu^+_{\text{Red Breast}}$ and $\mu^-_{\text{Blue Feathers}}$, respectively. These embeddings are then employed to condition $p(y \space|\space \mathbf{c})$ and enable a correct label prediction ("Red Breasted Parrot").
• Figure 3: The solid lines represent the mean task accuracy under random interventions at probability $p_{int}$, while the shaded areas indicate the standard deviation of each method. Results are reported across different models and datasets, under varying levels of input noise $\theta \in [0,1]$. The Black-box model is not shown since it does not allow human interventions.
• Figure 4: 2D t-SNE visualization of the concept embedding space $\mathbf{c}$ for the CEBaB dataset, comparing V-CEM, Prob-CBM and CEM. V-CEM concept representation is much denser than the ones of CEM and Prob-CEM.
• Figure 5: Variation in V-CEM's ID accuracy across different values of $\lambda_p$ on the CEBaB and CelebA datasets.