Debugging Concept Bottleneck Models through Removal and Retraining

1. Problem Statement

Concept Bottleneck Models (CBMs) are interpretable architectures that predict final labels via an intermediate layer of human-understandable concepts (e.g., "has a beak," "blue wings"). While this allows domain experts to validate predictions and intervene at test time (e.g., correcting a misidentified concept), current CBMs suffer from systemic misalignment.

• The Core Issue: Models often learn spurious correlations (shortcuts) from biased data. For example, a bird classifier might rely on the background (water vs. land) rather than the bird's features.
• Limitation of Existing Solutions:
  • Test-time intervention: Correcting a single prediction does not fix the underlying reasoning flaw; the model will make the same mistake on new data.
  • Simple Removal: Simply deleting a spurious concept (e.g., removing "water" from the concept list) is insufficient. The model may still implicitly encode that information through other concepts, or performance may collapse if the model relied too heavily on the removed shortcut.
  • Current Retraining: Standard fine-tuning on the biased dataset often causes the model to re-learn the same spurious correlations.

The paper addresses the need for a framework that allows experts to globally edit a model's reasoning to remove reliance on undesired concepts and retrain the model to align with expert knowledge.

2. Methodology: The CBDebug Framework

The authors propose a two-step framework: Removal and Retraining, implemented via a novel method called CBDebug (Concept Bottleneck Debugger).

Step 1: Removal

• Process: A domain expert inspects the learned concept explanations (e.g., image patches for ProtoPNets or text descriptions for VLM-CBMs).
• Action: The expert identifies and selects a subset of concepts ($C_{spur}$) that are spuriously correlated with the label (e.g., "beach" background for bird classification) and marks them for removal.
• Outcome: These concepts are temporarily removed from the model's active set.

Step 2: Retraining (CBDebug)

Simply removing concepts is not enough. CBDebug treats the expert's feedback as a causal intervention to approximate a counterfactual distribution where the undesired concepts do not influence the label. It consists of three sub-stages:

1. Label (Auxiliary Label Generation):

   • The system uses the trained concept extractor $\phi$ to generate activation scores for the spurious concepts ($C_{spur}$) across the entire training dataset.
   • These scores form an approximated auxiliary label matrix ($\hat{V}$). Unlike traditional bias mitigation which requires discrete group labels, $\hat{V}$ is continuous and high-dimensional.

2. Reweight (Permutation Weighting):

   • The method employs Permutation Weighting (Arbour et al., 2021) to compute sample weights ($U$).
   • It constructs a "confounded" dataset (original $Y$ and $\hat{V}$) and an "unconfounded" dataset (where $Y$ is permuted relative to $\hat{V}$).
   • A binary predictor is trained to distinguish between these two distributions. The resulting weights ($u_i$) represent the likelihood that a sample belongs to the unconfounded distribution.
   • Goal: Down-weight samples where the spurious concept strongly correlates with the label (e.g., a waterbird on a beach) and up-weight samples where the correlation is weak.

3. Augment (Targeted Data Augmentation):

   • To address the instability of reweighting (where minority groups might get extremely high weights), CBDebug performs targeted augmentation.
   • Samples with low weights (highly biased) are selected for augmentation based on a probability derived from the inverted weights.
   • Implementation:
     • ProtoPNets (PIP-Net): Uses CutMix to replace spurious patches (e.g., background) with patches from the concept bank.
     • VLM-CBMs: Uses Mixup with images generated from a text-to-image concept bank corresponding to the spurious concept.
   • Result: This creates a richer training distribution that forces the model to learn task-relevant features rather than shortcuts.

Finally, the model ($\phi, h$) is fine-tuned on the augmented dataset ($X_{aug}$) using the computed weights ($U$).

3. Key Contributions

1. General Debugging Framework: A two-step "Removal and Retraining" process that extends beyond specific architectures (like ProtoPNets) to a general class of unsupervised CBMs.
2. CBDebug Algorithm: A novel retraining method that bridges interpretable debugging and bias mitigation. It uniquely converts high-level concept-level feedback into sample-level auxiliary labels without requiring manual group annotation.
3. Causal Intervention Approach: By using permutation weighting and targeted augmentation, CBDebug effectively approximates a counterfactual distribution, reducing reliance on spurious concepts more effectively than standard fine-tuning or simple removal.

4. Experimental Results

The framework was evaluated on Waterbirds, MetaShift, CelebA, and ISIC datasets using two model architectures: PIP-Net (patch-based) and Post-hoc CBM (text-based).

• Performance Gains:
  • CBDebug significantly outperformed baselines (Original, Removal only, Standard Retrain, ProtoPDebug).
  • Worst-Group Accuracy (WGA): On the Post-hoc CBM for Waterbirds, CBDebug improved WGA by 26.1% over the original model (from 25.8% to ~51.9% in some configurations, with up to 26% improvement noted in the abstract).
  • On MetaShift, it improved WGA by 4.8% (Post-hoc) and 4.9% (PIP-Net).
• Robustness to Feedback Sources:
  • Real Users: Tested with 6 human experts; CBDebug consistently improved performance across all sessions.
  • Automated Feedback: Tested using LLMs (GPT-3.5-turbo) to identify spurious concepts. CBDebug maintained robust performance, outperforming the original model by up to 42.6% on CelebA, demonstrating scalability.
• Qualitative Analysis:
  • Visualizations confirmed that CBDebug successfully removed spurious concepts (e.g., "beach," "bamboo") from the top-activated concepts and replaced them with task-relevant features (e.g., "beak," "wing shape").
  • Baseline retraining often replaced one spurious concept with another (e.g., swapping "beach" for "water"), whereas CBDebug eliminated the reliance on background entirely.

5. Significance

• Bridging Interpretability and Robustness: The paper demonstrates that interpretability is not just for post-hoc analysis but can be actively used as a mechanism to debug and improve model robustness.
• Human-in-the-Loop Efficiency: It provides a practical workflow for domain experts to correct systemic model biases without needing to re-label the entire dataset or understand complex hyperparameters.
• Automation Potential: The success with LLM-generated feedback suggests that this framework can be scaled to automate bias mitigation in large-scale systems where human experts are unavailable for every dataset.
• Superiority over Unsupervised Methods: Unlike unsupervised bias mitigation methods (e.g., JTT, LfF) that guess at spurious correlations based on training dynamics, CBDebug leverages explicit expert knowledge, leading to more reliable and targeted corrections.

In summary, CBDebug offers a rigorous, interpretable, and effective solution to the "shortcut learning" problem in Concept Bottleneck Models, ensuring models are not only accurate but also "right for the right reasons."