Discovering Fine-Grained Visual-Concept Relations by Disentangled Optimal Transport Concept Bottleneck Models

TL;DR

The paper addresses the lack of fine-grained visual-concept alignment in vanilla Concept Bottleneck Models, which leads to spurious correlations and weak localization. It introduces DOT-CBM, a two-stage framework that uses a disentangled optimal transport between image patches and concepts, with transport plan $\mathbf{Q}$ and costs $d_{ij}=1-\langle p_i,t_j\rangle$, augmented by orthogonal projection losses and saliency/co-occurrence priors to learn robust, interpretable mappings. DOT-CBM yields concept activation predictions and explicit concept inversion masks, achieving state-of-the-art results on image classification, part detection, and out-of-distribution generalization, while providing transparent heatmaps of visual-concept relations. This approach enhances interpretability and reliability in concept-based vision systems, offering practical benefits for transparent AI deployments.

Abstract

Concept Bottleneck Models (CBMs) try to make the decision-making process transparent by exploring an intermediate concept space between the input image and the output prediction. Existing CBMs just learn coarse-grained relations between the whole image and the concepts, less considering local image information, leading to two main drawbacks: i) they often produce spurious visual-concept relations, hence decreasing model reliability; and ii) though CBMs could explain the importance of every concept to the final prediction, it is still challenging to tell which visual region produces the prediction. To solve these problems, this paper proposes a Disentangled Optimal Transport CBM (DOT-CBM) framework to explore fine-grained visual-concept relations between local image patches and concepts. Specifically, we model the concept prediction process as a transportation problem between the patches and concepts, thereby achieving explicit fine-grained feature alignment. We also incorporate orthogonal projection losses within the modality to enhance local feature disentanglement. To further address the shortcut issues caused by statistical biases in the data, we utilize the visual saliency map and concept label statistics as transportation priors. Thus, DOT-CBM can visualize inversion heatmaps, provide more reliable concept predictions, and produce more accurate class predictions. Comprehensive experiments demonstrate that our proposed DOT-CBM achieves SOTA performance on several tasks, including image classification, local part detection and out-of-distribution generalization.

Figures (4)

• Figure 1: Comparison between vanilla CBM and our proposed DOT-CBM, including architecture design and inversion heatmap for concept predictions. (a) Due to the black-box mapping from images to concepts, Vanilla CBM needs Grad-CAM selvaraju2017grad techniques to locate concept predictions back to image space while DOT-CBM can provide an explicit inversion heatmap visualization. (b) Due to lack of fine-grained alignment, vanilla CBM produces spurious correlations in three levels of granularity (part-background, part-object and part-part level) that mislocalizes the local concept to the background, the whole object, and incorrect local region.
• Figure 2: Overview of our proposed DOT-CBM. The overall CBM framework proceeds from left to right. In the first stage, the model transforms input images into concept activation vectors. In forward pass, local embeddings from a pre-trained Vision Transformer (ViT) and a text encoder are processed through learnable adapters to generate two feature sets (Sec. \ref{['local']}). Two orthogonal projection losses are applied to both feature sets to constrain Adapter training. An Optimal Transport (OT) algorithm optimizes the Assignment Matrix, representing the explicit correlation between concepts and image patches. Concept activation values are derived by combining this matrix with a cost matrix, and the loss is supervised by concept labels for training (Sec. \ref{['ot']}).To address data bias, we use the Saliency Map from the pre-trained ViT for rough foreground-background separation as the visual prior in OT. A Conditioned Co-occurrence Matrix, based on frequency statistics, serves as the prior for concepts, enhancing the model's ability to distinguish co-occurring concepts (Sec. \ref{['prior']}). In the second stage, consistent with the vanilla CBM framework, concept activation values are used to predict final class labels via a linear classification network, which is trained using class labels (Sec. \ref{['overall']}).
• Figure 3: Conditioned co-occurrence matrix of concept labels.
• Figure 4: Qualitative results for the concept inversion. VanillaCBM and CoopCBM adopt Grad-CAM to provide a concept inversion heatmap, while our proposed DOT-CBM can generate the inversion mask on its own.