Technical Summary: Mixture of Concept Bottleneck Experts (M-CBEs)

Problem Statement

Concept Bottleneck Models (CBMs) have emerged as a promising approach to align Deep Learning models with human reasoning by decomposing predictions into two stages: a concept encoder that maps raw inputs to human-interpretable variables (concepts), and a task predictor that maps these concepts to the final output. However, existing CBMs typically constrain the task predictor to a single, global function with a functional form set a priori (e.g., a single linear layer or a fixed Boolean expression).

This rigidity presents two critical limitations:

1. Limited Expressiveness: A single global function often fails to capture the true, potentially complex, concept-to-task generating process, leading to suboptimal predictive accuracy.
2. Inflexible Interpretability: Interpretability is inherently user-dependent. Fixing the functional form (e.g., forcing a linear model when the relationship is non-linear, or a Boolean model when the relationship is continuous) may unnecessarily sacrifice predictive performance or fail to align with the specific mental models of different users.

Current approaches often route inputs to mixtures of experts (MoE) to improve accuracy, but these typically route to black-box predictors or fix the functional form of the experts, failing to address the joint optimization of accuracy and user-aligned interpretability.

Methodology

The authors propose Mixture of Concept Bottleneck Experts (M-CBEs), a unified framework that generalizes CBMs along two dimensions:

1. Number of Experts: Instead of a single global function, the task predictor is modeled as a mixture of $M$ specialized functions (experts).
2. Functional Form: The framework allows explicit control over the functional form of each expert, ranging from simple linear expressions to complex symbolic expressions defined by user-specified operator vocabularies.

Core Framework

M-CBEs represent the task predictor as a random variable $T$ (an expression tree) conditioned on the input $x$. The joint distribution factorizes as:
$$p(y, c | x) = p(c | x) \sum_{m=1}^{M} p(m | x) p(y | c, t_m)$$
Where:

• $p(c | x)$ is the concept encoder.
• $p(m | x)$ is a learnable selector (router) that assigns input $x$ to a specific expert index $m$.
• $t_m$ is a unique, fixed expression tree associated with expert $m$.
• $p(y | c, t_m)$ is the task predictor executing the selected expression $t_m$ on the concepts $c$.

By enforcing a finite set of expression trees, M-CBEs ensure global interpretability, allowing users to inspect all reasoning paths in finite time, unlike models that generate unique expressions for every input (which result in infinite complexity).

Two Novel Instantiations

The paper introduces two specific models to explore this design space:

1. Linear M-CBE (Lin-M-CBE):

   • Structure: Restricts the expression trees to a linear functional form ($y = w^T c + b$).
   • Mechanism: Learns a finite set of $M$ distinct linear expressions. The selector routes inputs to the most appropriate linear expert.
   • Goal: Provides a robust balance for high-dimensional classification tasks where linear combinations of concepts are sufficient but a single global linear model is insufficient.

2. Symbolic M-CBE (Sym-M-CBE):

   • Structure: Leverages Symbolic Regression (SR) to discover the structure of the expression trees.
   • Mechanism: Users define a vocabulary of interpretable operators $W$ (e.g., $\{+, -, \times, \sin, \exp\}$). The model first trains a "placeholder" black-box (e.g., MLP) for each expert. Then, symbolic regression is applied to the data subsets assigned to each expert to distill the optimal symbolic expression tree $t_m$ using the user-defined vocabulary.
   • Adaptability: Supports post-hoc adaptation. Once the encoder and selector are trained, users can change the operator vocabulary $W$ to generate new symbolic expressions aligned with their domain expertise without retraining the entire model.

Key Contributions

1. Unified Framework: M-CBEs subsume existing concept-based methods (like standard CBM, CMR, DCR) as special cases while exposing a largely unexplored two-dimensional design space (number of experts $\times$ functional form).
2. Theoretical Guarantees: The authors provide propositions demonstrating that:
   • A single non-linear expert can approximate any continuous function (Universal Approximation).
   • A mixture of linear experts can approximate any continuous function if $M$ is sufficiently large, with approximation error bounds decreasing as $M$ increases.
3. Global Interpretability: By learning from a finite set of expressions, M-CBEs avoid the "infinite memory" problem of instance-specific predictors (like LICEM or DCR), enabling full model inspection and verification.
4. Post-hoc Adaptability: Sym-M-CBE allows users to modify the complexity and type of expressions (e.g., switching from trigonometric to polynomial forms) after initial training, catering to diverse user expertise levels.

Experimental Results

The framework was evaluated on four synthetic and five real-world datasets (including MNIST-Arithm, dSprites-Exp, Pendulum, MAWPS, AWA2, CUB-200, and CIFAR-10) covering both classification and regression tasks.

• Accuracy-Interpretability Trade-off: Exploring the design space allows finding the Pareto frontier.
  • Symbolic M-CBE excels in regression tasks (e.g., Pendulum), discovering compact, high-accuracy expressions that dominate the Pareto frontier.
  • Linear M-CBE provides a robust balance for high-dimensional classification tasks, outperforming rigid Boolean logic (CMR, DCR) which degrades significantly as the number of concepts increases (e.g., on CUB-200 and CIFAR-10).
• Scalability: Boolean functional forms fail to scale with high-dimensional concept spaces due to the brittleness of conjunctive rules. Algebraic forms (Linear and Symbolic) maintain robust performance.
• Intervenability:
  • Sym-M-CBE achieves near-ideal intervenability on regression tasks, recovering ground-truth mechanisms (e.g., $8.0 \sin(\theta) + 10.0$) and showing sharp drops in error when concepts are corrected.
  • Lin-M-CBE is highly responsive to interventions in classification tasks, outperforming other globally interpretable models.
• Handling Incomplete Concepts: Increasing the number of experts ($M > 1$) compensates for incomplete concept bottlenecks or varying task logic across the input space, recovering predictive performance where single-expert models fail.
• Complexity: M-CBEs frequently match or exceed the accuracy of black-box MLPs while being substantially less complex (fewer nodes in expression trees).

Significance and Claims

The paper claims that M-CBEs establish a principled approach for developing interpretable models that adapt to both task requirements and diverse user needs.

• Democratizing AI: By enabling users to align model predictive mechanisms with computations they can interpret (from complex mathematical expressions for domain experts to simplified polynomials for lay users), the framework democratizes access to AI.
• Responsible AI: It advances responsible AI deployment by ensuring predictions remain grounded in finite, inspectable sets of interpretable expressions, addressing the opacity of black-box models in high-stakes domains.
• Optimal Trade-off: The authors assert that no single architectural choice is universally superior; rather, the ability to navigate the design space of expert cardinality and functional form is essential for finding the optimal accuracy-interpretability trade-off for a specific problem.

Limitations Acknowledged:

• The selector network $p(m | x)$ operates as a black-box, limiting transparency regarding why a specific expert is chosen.
• Symbolic regression is computationally intensive, particularly for large operator vocabularies.
• The number of experts $M$ must be specified as a hyperparameter, requiring manual balancing of expressiveness and interpretability.