Concept-based Decision Tree Distillation
- The paper introduces a technique that distills black-box deep models into compact decision trees, maintaining above 92% predictive fidelity with significantly reduced complexity.
- It leverages interpretable concept representations and additive ensembles to enable adaptive test-time intervention and precise attribution of model decisions.
- Empirical evaluations across computer vision, NLP, and graph domains demonstrate the method’s effectiveness in balancing transparency with robust generalization.
Concept-based Decision Tree Distillation is a methodology for converting black-box deep learning models—specifically those relying on interpretable intermediate “concept” representations—into compact, human-interpretable decision-tree or logical classifiers. This approach bridges state-of-the-art performance with rigorous model transparency by taking advantage of the semantic structure provided by concept bottlenecks and recent advances in interpretable additive tree ensembles. Modern instantiations, such as FIGS-BD for Concept Bottleneck Models and logical decision tree distillation for GNNs, demonstrate that distilled tree-based surrogates can retain high predictive fidelity while offering precise attribution and facilitating test-time human intervention (Shen et al., 9 Mar 2025, Pluska et al., 2024).
1. Foundations: Concept Bottleneck Models and Motivation
Concept Bottleneck Models (CBMs) decompose prediction tasks through a two-stage process: an encoder extracts a -dimensional vector of human-interpretable “concepts” , which is then mapped to output space by a function . While CBMs promise transparency, in practice, their concept-to-label module often defaults to a black-box predictor (e.g., MLPs or Transformers), limiting actual interpretive utility. Concept-based decision tree distillation addresses this by learning an explicit, transparent surrogate such that , where is typically a binary transformation of concepts—enabling Boolean or logical rule sets over interpretable concepts (Shen et al., 9 Mar 2025).
For graph learning, logical distillation leverages the equivalence between message-passing neural architectures and the two-variable fragment of first-order logic with counting quantifiers (C2), directly translating node or graph predictions into interpretable logical structure (Pluska et al., 2024).
2. Formal Task Definition and Distillation Objective
The general distillation setup defines the following problem: given a pre-trained teacher model with access to concept representations , learn a student decision tree (or additive ensemble) over binarized (or otherwise discretized/logically structured) features to approximate the teacher’s outputs:
0
The distillation dataset is
1
where 2 for specified thresholds 3 (e.g., 4 for zero-centered image concepts).
The training objective incorporates both a fidelity term and a complexity penalty:
5
where 6 is squared or cross-entropy loss, and 7 counts decision rules or leaves (Shen et al., 9 Mar 2025).
3. Methodologies: Algorithmic Frameworks
3.1 Binary Distillation with FIGS
Fast Interpretable Greedy Sum-Trees (FIGS) is an ensemble of 8 shallow binary trees, each of depth at most 9:
0
Each tree is fit in a stage-wise, gradient-boosting fashion (fitting current residuals), but remains shallow for interpretability. Local split selection at each node optimizes a regularized gain:
1
The result is a compact, additive model whose per-path contributions are readily traceable (Shen et al., 9 Mar 2025).
Pseudocode Excerpt (see (Shen et al., 9 Mar 2025)):
0
3.2 Logical Decision Tree Distillation for GNNs
Given a GNN of depth 2, iterated decision trees (IDTs) are constructed layerwise, encoding modal count features over unary atomic concepts (graph node features or annotations). Each tree layer operates over features like “self,” “neighbors,” and “self + neighbors” modal counts (e.g., number of neighbors satisfying a sub-concept 3). Decision tree splits are scored by variance reduction (for regression) or information gain (for classification) (Pluska et al., 2024).
Pseudocode Excerpt (see (Pluska et al., 2024)):
1
4. Interpretability and Attribution Mechanisms
Decision paths in the distilled trees correspond to explicit conjunctions over binary concepts or modal logical formulas. In FIGS-BD, each path represents a conjunction of “concept 4 present/absent” conditions, with additive leaf constants summing to the final prediction. The decomposition into a small ensemble of shallow trees (30–50 trees of depth 5 3–4) enables direct inspection of which concept combinations drive the model’s output (Shen et al., 9 Mar 2025).
In the logical setting, each decision tree split over a modal-count feature (e.g., 6 for “at least 7 neighbors satisfy 8”) corresponds to an interpretable logical clause in C2, often recoverable as known ground-truth logical formulas (Pluska et al., 2024).
5. Empirical Evaluation and Model Compactness
Extensive evaluations have shown that concept-based decision tree distillation achieves high predictive fidelity with dramatic reductions in model size and improved transparency:
| Dataset | Teacher (CBM/TBM) Accuracy | Student (FIGS-BD) Accuracy |
|---|---|---|
| CUB (200-way) | 79.8% | 75.9% (95.1%) |
| TravelingBirds | 51.8% | 47.9% (92.4%) |
| AGNews | 89.6% | 88.8% (99.1%) |
| CEBaB (R²) | 0.868 | 0.871 (100.3%) |
The FIGS-BD models used ≤200 rules versus thousands for dense XGBoost baselines, maintaining ≥92.5% of teacher performance, sometimes exceeding it in generalization tasks (Shen et al., 9 Mar 2025). Logical IDTs distilled from GNNs similarly matched or outperformed the original model fidelity, often providing succinct decision trees of ≤10 internal nodes (Pluska et al., 2024).
6. Adaptive Test-Time Intervention
A salient property of concept-based decision tree surrogates is their suitability for adaptive test-time intervention (ATTI). FIGS-BD enables per-sample ranking of the most critical concept interactions by analyzing additive contributions across the tree ensemble. At test time, end-users can be prompted to validate or correct only the top-ranked binary concepts, providing targeted intervention that yields sharp performance improvements, even when only a handful of concept corrections are allowed:
- Empirically, intervention on just 2–3 groups yields substantial gains, outperforming random or baseline ATTI methods.
- Pseudocode for ATTI (see (Shen et al., 9 Mar 2025)) is provided to facilitate practical deployment.
7. Theoretical Guarantees and Limitations
The logical distillation framework is provably lossless in expressive power for properties definable in C2—that is, for any message-passing GNN, there exists an equivalent decision-tree over modal-count features reconstructing its predictions (Pluska et al., 2024). Complexity in both frameworks is determined by the number of concepts, binary features, and tree depth:
- FIGS-BD: 9 total complexity
- Logical IDT: polynomial in data and GNN size due to limited tree depth and pruning
Compactness and interpretability are ensured by restricting maximum tree depth and number; however, if the underlying concepts are not truly disentangled or if the teacher relies on highly intricate interactions, some fidelity loss or increased tree complexity may occur.
8. Applications and Extensions
Concept-based decision tree distillation is deployed in:
- Computer vision (CUB, TravelingBirds), where concepts correspond to semantic parts or attributes
- NLP (AGNews, CEBaB), with one-hot or binary concepts denoting semantic themes or entities
- Graph-based domains (AIDS, PROTEINS), mapping GNNs to logical classifiers via iterated decision trees
A plausible implication is that such distillation frameworks can extend to any architecture for which an interpretable concept bottleneck exists or can be constructed, allowing unified post-hoc explainability while preserving predictive utility (Shen et al., 9 Mar 2025, Pluska et al., 2024).