Learn from A Rationalist: Distilling Intermediate Interpretable Rationales

TL;DR

This work tackles the interpretability challenge of deep models by addressing the training chicken-and-egg problem in rationale extraction (RE). It introduces REKD, a neural-model-agnostic distillation framework that lets a student RE model learn from a teacher’s rationales and predictions while exploring its own explanations, unified by a temperature-annealing schedule. Through extensive experiments on IMDB and CIFAR-10/100 with BERT and ViT backbones, REKD consistently improves the predictive performance of lightweight RE models and, in some cases, surpasses teacher baselines while maintaining interpretability. The proposed approach offers a practical path to deploy interpretable, robust models in data- and compute-constrained settings, with potential extensions to cross-architecture distillation and discrete latent structure learning.

Abstract

Because of the pervasive use of deep neural networks (DNNs), especially in high-stakes domains, the interpretability of DNNs has received increased attention. The general idea of rationale extraction (RE) is to provide an interpretable-by-design framework for DNNs via a select-predict architecture where two neural networks learn jointly to perform feature selection and prediction, respectively. Given only the remote supervision from the final task prediction, the process of learning to select subsets of features (or \emph{rationales}) requires searching in the space of all possible feature combinations, which is computationally challenging and even harder when the base neural networks are not sufficiently capable. To improve the predictive performance of RE models that are based on less capable or smaller neural networks (i.e., the students), we propose \textbf{REKD} (\textbf{R}ationale \textbf{E}xtraction with \textbf{K}nowledge \textbf{D}istillation) where a student RE model learns from the rationales and predictions of a teacher (i.e., a \emph{rationalist}) in addition to the student's own RE optimization. This structural adjustment to RE aligns well with how humans could learn effectively from interpretable and verifiable knowledge. Because of the neural-model agnostic nature of the method, any black-box neural network could be integrated as a backbone model. To demonstrate the viability of REKD, we conduct experiments with multiple variants of BERT and vision transformer (ViT) models. Our experiments across language and vision classification datasets (i.e., IMDB movie reviews, CIFAR 10 and CIFAR 100) show that REKD significantly improves the predictive performance of the student RE models.

Figures (3)

• Figure 1: The schematic of REKD. The student RE model learns from the teacher's supervision on feature selections and predictions in addition to its own RE exploration via Straight-Through Gumbel-Softmax. The red arrows indicate the gradient flows and the red 1 represents the Straight-Through estimation, i.e., $\frac{\partial \textbf{M}}{\partial \textbf{S}} \approx 1$. The definitions for the computation and the variables are in Section \ref{['sec:rekd']}.
• Figure 2: Rationale ratio vs. accuracy for ViT Base, Small and Tiny on CIFAR 10 by varying $p_\text{target}$ from 5% to 75%. The x-axis is the actual rationale ratio and the y-axis is the accuracy. The target rationale ratio $p_\text{target}$ is used as the label for each data point. The label "cls" represents the classification result. Each data point is an average of 10 runs. See Table \ref{['tab:target_acc']} for the detailed numbers.
• Figure 3: Examples of rationale extraction for ViT Base on CIFAR 10, CIFAR 100 and BERT Base on IMDB movie reviews (arranged from top to bottom).