[Paper Review] Distilling a Neural Network Into a Soft Decision Tree

Nicholas Frosst, Geoffrey E. Hinton|arXiv (Cornell University)|Nov 27, 2017

Machine Learning and Data Classification5 references266 citations

TL;DR

The paper presents a method to distill knowledge from a trained neural network into a soft decision tree that makes hierarchical decisions, improving interpretability while retaining reasonable accuracy.

ABSTRACT

Deep neural networks have proved to be a very effective way to perform classification tasks. They excel when the input data is high dimensional, the relationship between the input and the output is complicated, and the number of labeled training examples is large. But it is hard to explain why a learned network makes a particular classification decision on a particular test case. This is due to their reliance on distributed hierarchical representations. If we could take the knowledge acquired by the neural net and express the same knowledge in a model that relies on hierarchical decisions instead, explaining a particular decision would be much easier. We describe a way of using a trained neural net to create a type of soft decision tree that generalizes better than one learned directly from the training data.

Motivation & Objective

• Motivate the tension between deep nets' generalization and interpretability.
• Propose a soft hierarchical decision tree distilled from a neural network.
• Show that the distilled tree generalizes better than a tree trained on raw data.
• Demonstrate the approach on MNIST and other datasets with qualitative interpretability benefits.

Proposed method

• Use a soft binary decision tree with learned filters at inner nodes and leaf distributions Q_ell over classes.
• Each inner node i computes p_i(x) = sigma(beta(x w_i + b_i)) as the probability to go right.
• Leaves hold class distributions Q^ell_k = exp(phi^ell_k) / sum_k' exp(phi^ell_k').
• Train tree by mini-batch gradient descent to minimize L(x) = -log( sum_ell P^ell(x) sum_k T_k log Q^ell_k ).
• Regularize to encourage balanced use of sub-trees via a depth-dependent cross-entropy penalty tied to alpha_i, the average path probability to node i.
• Optionally distill from neural network predictions by using soft targets T that mix true labels with NN outputs.
• At test time, use leaf with maximum path probability for the final predictive distribution.

Experimental results

Research questions

• RQ1 Can a soft decision tree mimic a neural network's input-output function while remaining interpretable?
• RQ2 Does distillation from a neural network improve a soft decision tree's accuracy over training directly on data?
• RQ3 How do regularizers and depth-related penalties affect learning and generalization?

Key findings

• On MNIST, a depth-8 soft decision tree trained on true targets reaches 94.45% test accuracy.
• A neural net with CNN layers achieves 99.21% on MNIST, higher than the soft tree.
• Soft targets from the neural net improve the tree to 96.76% test accuracy, midway between the NN and the tree trained on true targets.
• Soft trees generalize better than trees trained directly on data due to data distribution sparsity at lower nodes.
• Across datasets, distillation enabled reasonable accuracy even with interpretable models, e.g., Connect4: 80.60% vs 78.63% (NN-free baseline); Letter: 78.0% (depth 9, raw) and 81.0% (distilled from NN ensemble).
• The approach yields interpretable visualizations of decision paths and learned filters, aiding explanation of individual predictions.