[Paper Review] NBDT: Neural-Backed Decision Trees

Alvin Wan, Lisa Dunlap+7 more|arXiv (Cornell University)|Apr 1, 2020

Explainable Artificial Intelligence (XAI)56 references66 citations

TL;DR

NBDTs replace a neural net’s final layer with a differentiable oblique decision tree, enabling high accuracy while providing interpretable, path-based explanations.

ABSTRACT

Machine learning applications such as finance and medicine demand accurate and justifiable predictions, barring most deep learning methods from use. In response, previous work combines decision trees with deep learning, yielding models that (1) sacrifice interpretability for accuracy or (2) sacrifice accuracy for interpretability. We forgo this dilemma by jointly improving accuracy and interpretability using Neural-Backed Decision Trees (NBDTs). NBDTs replace a neural network's final linear layer with a differentiable sequence of decisions and a surrogate loss. This forces the model to learn high-level concepts and lessens reliance on highly-uncertain decisions, yielding (1) accuracy: NBDTs match or outperform modern neural networks on CIFAR, ImageNet and better generalize to unseen classes by up to 16%. Furthermore, our surrogate loss improves the original model's accuracy by up to 2%. NBDTs also afford (2) interpretability: improving human trustby clearly identifying model mistakes and assisting in dataset debugging. Code and pretrained NBDTs are at https://github.com/alvinwan/neural-backed-decision-trees.

Motivation & Objective

• Aim to jointly improve predictive accuracy and interpretability for image classification tasks.
• Replace the final linear layer of neural networks with a differentiable oblique decision tree.
• Introduce a tree supervision loss and induced hierarchies to learn high-level concepts.
• Enable path-probability based inference that tolerates uncertain intermediate decisions.
• Demonstrate improved generalization to unseen classes and more trustworthy explanations.

Proposed method

• Replace final linear layer with a differentiable oblique decision tree using leaf weights tied to class predictions.
• Use soft (probabilistic) path traversal to allow recovery from uncertain early decisions (soft inference).
• Seed node weights with corresponding neural network weights and compute child probabilities via softmax inner products.
• Construct induced hierarchies by performing hierarchical clustering on pre-trained class weight vectors and averaging leaf weights to form inner-node weights.
• Label internal nodes using WordNet concepts to provide semantic meaning where possible.
• Train with a tree supervision loss that combines standard cross-entropy with a path-probability based cross-entropy over hierarchy path distributions, with time-varying weights.

Experimental results

Research questions

• RQ1 Can neural-backed decision trees achieve comparable or superior accuracy to modern neural networks on CIFAR, TinyImageNet, and ImageNet?
• RQ2 Do induced hierarchies based on model weights outperform data-driven or WordNet-based hierarchies for NBDTs?
• RQ3 Does the tree supervision loss improve original model accuracy and help learning high-level decisions?
• RQ4 Do NBDTs provide more helpful and trustworthy explanations than saliency maps in identifying misclassifications and ambiguous labels?
• RQ5 Can NBDTs generalize better to unseen classes compared to standard neural networks?

Key findings

• NBDTs match or outperform modern networks on CIFAR, TinyImageNet, and ImageNet in accuracy.
• NBDTs generalize to unseen classes by up to 16% and can improve the original model’s accuracy by up to 2%.
• Induced hierarchies built from pre-trained weights outperform WordNet and data-driven hierarchies in accuracy.
• Tree supervision with path probabilities improves learning and yields better performance than hierarchical softmax.
• NBDT explanations help users identify model mistakes more accurately than saliency maps and increase trust in challenging tasks.
• Zero-shot superclass generalization shows NBDT outperforms the backbone in several superclass distinctions.