Learn from A Rationalist: Distilling Intermediate Interpretable Rationales

AI summary

Overview

• Research area: Interpretable machine learning — rationale extraction (RE) combined with knowledge distillation (KD), applied to language and vision classification.
• Technical level: Advanced. The paper assumes familiarity with select-predict architectures, Gumbel-Softmax gradient estimators, and temperature-scaled knowledge distillation.
• Scope: The paper introduces REKD, a framework in which small "student" rationale-extraction models learn from a larger teacher's rationales and predictions in addition to their own RE objective, and validates it on IMDB, CIFAR 10, and CIFAR 100 with BERT and ViT variants.

What This Paper Is About

Rationale extraction builds neural networks that first select a small subset of input features (a rationale) and then predict using only those features, guaranteeing that the explanation is faithful. Training this select-predict architecture is hard because the generator and predictor must bootstrap each other from only the final task loss — a "chicken and egg" problem that gets worse when the underlying neural networks are small. The paper's goal is to improve the predictive performance of these lightweight student RE models by having them learn from a stronger teacher RE model's rationales and predictions.

Key Contributions

1. The authors identify the "chicken and egg" dilemma in training the select-predict architecture of RE, explaining why lightweight students struggle to learn effective rationales and achieve good predictive performance without external guidance.
2. They propose REKD (Rationale Extraction with Knowledge Distillation), a neural-model agnostic distillation framework for rationale extraction that exploits the intrinsic curriculum of Gumbel-Softmax temperature annealing.
3. They validate the approach on language (IMDB movie reviews) and vision (CIFAR 10 / CIFAR 100) classification using multiple variants of BERT and ViT as RE backbones, showing significant improvement in student predictive performance.
4. They argue for an objective evaluation metric — predictive performance under a rationale-ratio constraint — instead of plausibility measured as overlap with human annotations.

Main Findings

• Consistent gains over RE baselines: REKD improves student accuracy across all three datasets. On CIFAR 10, ViT Small goes from 0.889@RE to 0.968 and ViT Tiny from 0.797@RE to 0.936. On CIFAR 100, ViT Small goes from 0.779@RE to 0.845 and ViT Tiny from 0.645@RE to 0.777. On IMDB, BERT Small goes from 0.881@RE to 0.906 and BERT Mini from 0.863@RE to 0.892.
• Capacity matters for the "chicken and egg" problem: The drop from classification (CLS) to RE is larger for smaller backbones. On CIFAR 10, ViT Small falls 0.092 (0.981@CLS to 0.889@RE) while ViT Base falls only 0.020 (0.984@CLS to 0.964@RE). On CIFAR 100 the drops are 0.117 for ViT Base, 0.165 for ViT Small, and 0.258 for ViT Tiny.
• Students can match or exceed the teacher's average: On CIFAR 10, using the ViT Base teacher of seed=2029 (accuracy 0.969), the average ViT Small REKD accuracy of 0.968 slightly outperforms the average of the ViT Base models, 0.964. On CIFAR 100, with a ViT Base teacher from seed=2031 (accuracy 0.848@RE), the average ViT Small REKD result of 0.845 exceeds the average ViT Base result of 0.830.
• REKD students can beat their own CLS versions on IMDB: REKD yields 0.906 for BERT Small versus 0.889@CLS, and 0.892 for BERT Mini versus 0.877@CLS, suggesting the teacher delivers predictive patterns not acquired by the black-box student CLS models.
• IMDB needs few features: Under a 10% rationale constraint, degradation from CLS to RE is small — 0.002 for BERT Base, 0.008 for BERT Small, 0.014 for BERT Mini — indicating the IMDB task does not require many features for accurate prediction.
• CIFAR 100 is harder under the same constraint: With the same 15% rationale ratio as CIFAR 10, ViT RE models degrade more from CLS to RE, which the authors attribute to the more fine-grained 20-class task.
• Rationale length correlates with accuracy: Varying p_target from 5% to 75% for ViT Base, Small, and Tiny on CIFAR 10 (Figure 3, mean of 10 runs) shows accuracy drops faster for less capable ViT models as p_target decreases.
• Reduced variance: The paper reports that REKD improves predictive performance "with reduced variance," as shown in the standard deviations in Tables 1, 2, and 3.
• Ablations are deferred: Three ablation questions are raised (pure supervised KD without RE; whether gains come from more stable optimization versus teacher knowledge; and whether rationale or prediction distillation matters more), with details placed in Appendix C, which is not included in the available content.

Methodology in Plain English

The approach keeps the standard RE structure of two networks — a generator that decides which input features to keep, and a predictor that classifies using only the kept features — and adds a teacher.

• Differentiable selection: For each feature, the generator produces a two-way (selected / not selected) Gumbel-Softmax distribution. Soft samples are discretized by argmax into a binary mask, and gradients are passed through using the Straight-Through estimator, so the mask is discrete during the forward pass but differentiable during the backward pass.
• Length control: A selection loss penalizes squared deviation between the number of selected features and a target of L × p_target; this is combined with the cross-entropy prediction loss into the RE objective.
• Distilling rationales: The student generator is trained to match the teacher's temperature-scaled Gumbel-Softmax distributions feature by feature, using a summed Kullback–Leibler divergence across all L features.
• Distilling predictions: The student predictor matches the teacher's temperature-scaled softmax output distribution, again with KL divergence, scaled by τ² following standard KD practice.
• Shared temperature schedule: The KD temperature is synchronized with the Gumbel-Softmax annealing schedule, which decays exponentially (from τ_0 to τ_K over K steps). This creates an implicit curriculum: the student first absorbs broad, softened teacher knowledge, then converges to sharp, high-confidence feature selections as the temperature anneals.
• Combined objective: The total loss is a weighted mix of the RE loss and the KD loss, controlled by α, with a separate weight λ_R on the rationale distillation term.
• Setup: The teacher is bert-base-uncased for IMDB and vit-base-patch16-224 for CIFAR; students are bert-small/bert-mini and vit-small-patch16-224/vit-tiny-patch16-224, all initialized from pretrained weights. Training used 35 epochs for RE and REKD and 20 epochs for classification, learning rate 1e-5, batch size 32, τ_0 = 5, τ_K = 0.1 with temperature adjusted every 100 training steps, λ_R = 0.5, p_target of 15% for ViTs on CIFAR 10/100 and 10% for BERTs on IMDB (BERT Small@RE used 9.5%). Experiments were repeated over 10 random seeds, 2026 through 2035.

Why This Matters

• Impact on research: REKD is described as the first approach applying knowledge distillation to rationale extraction. Because the feature-selection layer acts as a universal, architecture-agnostic interface, it sidesteps the dimension-mismatch problems that require projection functions or channel collapsing in conventional feature-based distillation. It also reframes evaluation around predictive performance under a rationale constraint rather than human-annotation overlap.

• Real-world applications (as suggested by the paper):

  • Resource-constrained deployment where interpretability matters, such as mobile devices for healthcare.
  • High-stakes domains generally, including healthcare and finance, where opaque model knowledge raises trustworthiness concerns.
  • Any setting that needs faithful, faithful-by-design explanations rather than post-hoc approximations such as LIME, SHAP, Integrated Gradients, or Grad-CAM.
  • Distillation of discrete latent structures learned with Gumbel-Softmax, such as relational graphs.

• Industry relevance: The method is model-agnostic, so off-the-shelf pretrained backbones (the exact HuggingFace labels are listed in the paper) can be plugged in as teacher or student. Public code is released at https://github.com/JiayiDai/REKD.

Future Directions

• Cross-architecture distillation: The paper explicitly notes that all experiments distill between models of the same general architecture (ViT-to-ViT or BERT-to-BERT), even with differing hidden dimensions and layer counts. Testing architectures with fundamentally different inductive biases, such as ViT to ResNet, remains open — and would require that both models share identical input features, e.g., identical tokenization.
• Better rationale-quality validation: The authors acknowledge that confirming rationale quality without ground-truth feature maps is inherently difficult, and that select-predict architectures can be vulnerable to learning covert communication channels instead of informative structures. Human-centric usability studies are suggested as a valuable complement.
• Distillation of other discrete latent structures: Beyond rationale extraction, the approach could extend to other Gumbel-Softmax-based discrete structures such as relational graphs.
• Hyperparameter guidance: The paper reports that λ_select and α were examined experimentally across model variants, with details placed in appendices not included in the available content.

Target Audience

Researchers and practitioners working on explainable AI, interpretable-by-design neural architectures, and model compression will benefit most. The paper is also relevant to engineers who need faithful explanations on limited compute, and to readers interested in knowledge distillation methods for structured or discrete intermediate representations. A background in Gumbel-Softmax estimators and distillation losses is needed to follow the method sections in full.

Authors’ 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 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 REKD (Rationale Extraction with Knowledge Distillation) where a student RE model learns from the rationales and predictions of a teacher (i.e., a 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.