Introduction
Concept-based eXplainable AI (C-XAI) aims to make machine learning models more transparent by introducing an intermediate layer of human-understandable concepts. In a typical Concept Bottleneck Model (CBM), the system first predicts a set of discrete concepts (e.g., "has wings," "red breast") before using those predictions to determine a final class (e.g., "bird"). This structure allows for human intervention: if a model misidentifies a concept, a human can correct it, and the model should ideally update its final prediction accordingly.
However, traditional CBMs often suffer from a performance penalty compared to black-box models. To address this, Concept Embedding Models (CEMs) were introduced, replacing simple scalars with high-dimensional embeddings. While CEMs recover black-box accuracy, they introduce a new problem: "concept leakage." The embeddings often become so entangled with the raw input data that human interventions lose their effectiveness, particularly when the model encounters out-of-distribution (OOD) data. This paper introduces the Variational Concept Embedding Model (V-CEM), a framework designed to achieve high accuracy while maintaining robust intervenability through variational inference.
The Problem: The Intervenability-Accuracy Trade-off
The core issue in current concept-based modeling is the entanglement between high-dimensional embeddings and raw input features. In CEMs, the embeddings contain both the intended concept information and residual information from the input. When a user intervenes to change a concept label, the underlying embedding may still carry noisy or incorrect features from the original input, causing the downstream predictor to ignore the human correction. This failure is most prominent in OOD scenarios where the input data is noisy or shifted, rendering the model's "interpretability" practically useless for error correction.
Methodology: Variational Concept Embedding Models
To decouple concept embeddings from raw input data, the authors propose V-CEM. Unlike deterministic CEMs, V-CEM treats the generation of concept embeddings as a probabilistic process. The model assumes a generative process where embeddings depend strictly on the interpretable concepts rather than the raw input.
Variational Inference and Prior Matching
V-CEM is trained using an Evidence Lower Bound (ELBO) optimization objective. A critical component of this objective is the Prior Matching term, which utilizes Kullback-Leibler (KL) divergence. This term forces the approximate posterior distribution—which has access to both the input and the concept—to align with a prior distribution that only knows whether a concept is active or inactive. By penalizing the model when the posterior deviates from this concept-only prior, V-CEM effectively "scrubs" the raw input features from the embedding space.
Concept Representation Cohesiveness (CRC)
To quantify the quality of the learned representations, the authors introduce a new metric: Concept Representation Cohesiveness (CRC). Based on the silhouette score, CRC measures how well the embeddings cluster into distinct groups for active versus inactive concepts. High CRC indicates that the embeddings are dense and well-separated, which prevents the "leakage" of spurious input features into the concept bottleneck.
Experimental Evaluation
The authors evaluated V-CEM against CBMs and standard CEMs across several domains:
- Vision: MNIST Even/Odd, MNIST Addition (summing two digits), and CelebA (facial attribute recognition).
- NLP: CEBaB (aspect-level restaurant reviews) and IMDB (movie review sentiment based on cinematography or acting).
- OOD Robustness: To simulate real-world noise, the authors added progressive Gaussian noise to the inputs during inference to test how well human interventions could recover model performance.
Key Results
- High Accuracy: V-CEM maintains the performance benefits of embedding-based models. On the MNIST Addition task, V-CEM achieved 73.12% accuracy, significantly outperforming Probabilistic CBMs (27.31%) and linear CBMs (44.19%).
- Superior Intervenability: Under OOD conditions (Gaussian noise), standard CEM accuracy often flatlines regardless of human intervention. V-CEM, however, shows a linear increase in accuracy as more concepts are corrected, matching the reliable behavior of traditional CBMs.
- Representation Quality: V-CEM achieved near-perfect CRC scores (e.g., 0.98 on MNIST Even/Odd), compared to 0.65 for standard CEMs. This confirms that the variational prior successfully creates distinct, cohesive concept clusters that are resistant to input noise.
Discussion and Limitations
V-CEM successfully bridges the gap between the high performance of CEMs and the reliable intervenability of CBMs. By using variational inference to enforce a structured latent space, it ensures that human corrections are meaningful even when the input data is corrupted.
However, some limitations remain. V-CEM does not currently include a mechanism to automatically detect when data is OOD; it relies on a human to initiate the intervention. Additionally, the current model assumes that concepts are independent, which may not hold true in complex real-world datasets where concepts (e.g., "has feathers" and "can fly") are highly correlated. Future work may involve modeling these dependencies and extending the framework to multimodal settings where concept embeddings must be aligned across text and image data.