🤖 machine learning

Learn from A Rationalist: Distilling Intermediate Interpretable Rationales

This paper proposes REKD, a knowledge distillation framework that enhances the predictive performance and interpretability of rationale extraction models by training smaller student networks to learn from the rationales and predictions of a more capable teacher model, thereby overcoming the computational challenges of learning feature subsets under remote supervision.

Original authors: Jiayi Dai, Randy Goebel

Published 2026-05-29

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jiayi Dai, Randy Goebel

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Chicken and Egg" of AI

Imagine you are trying to teach a student how to solve a complex puzzle. The student has two jobs:

Pick the pieces: They must decide which puzzle pieces are actually important.
Solve the puzzle: They must use only those picked pieces to solve the puzzle.

The problem is that the student doesn't know which pieces are important until they try to solve the puzzle, but they can't solve the puzzle well until they know which pieces are important. This is the "Chicken and Egg" dilemma.

In the world of Artificial Intelligence (specifically Deep Neural Networks), this happens when we try to make AI "explainable." We want the AI to highlight the specific words in a movie review or the specific pixels in a photo that led to its decision. This is called Rationale Extraction (RE).

Usually, this works great if the AI is very smart (a "large" model). But if the AI is smaller or less powerful (a "student" model), it gets stuck. It picks the wrong pieces, fails to solve the puzzle, gets confused, and never learns the right way to pick pieces.

The Solution: Learning from a "Rationalist" Teacher

The authors of this paper propose a clever fix called REKD (Rationale Extraction with Knowledge Distillation).

Think of it like a master chef (the Teacher) teaching an apprentice (the Student).

The Old Way: The apprentice tries to guess which ingredients are important and cook the dish alone. They often fail because they don't have enough experience.
The REKD Way: The master chef doesn't just say, "Here is the recipe." Instead, the chef says: "Look, I picked these specific herbs and spices (the Rationale), and because I used them, the dish tastes perfect (the Prediction)."

The apprentice learns two things at once:

What to pick: They watch the teacher's hand to see exactly which ingredients were selected.
Why it works: They see that those specific ingredients lead to a delicious result.

By copying the teacher's "intermediate" choices (the rationales), the student learns much faster and makes fewer mistakes, even if the student is smaller and less powerful than the teacher.

How It Works: The "Soft" to "Hard" Training Camp

The paper uses a special training technique involving a "temperature" setting, which acts like a training curriculum:

Phase 1: The Soft Start (High Temperature): At the beginning, the teacher's choices are a bit "fuzzy" or "soft." The teacher might say, "I'm 70% sure this word is important, but maybe that one too." This helps the student explore and get a general feel for the task without getting overwhelmed by strict rules.
Phase 2: The Hard Finish (Low Temperature): As training goes on, the "temperature" drops. The teacher's choices become sharp and definite. "This word is 100% important; ignore the rest." The student then learns to make precise, confident decisions, just like the teacher.

This process ensures the student learns the logic of selection, not just the final answer.

What They Found (The Results)

The researchers tested this on two types of tasks:

Reading: Analyzing movie reviews (IMDB dataset) to see if they are positive or negative.
Seeing: Identifying objects in photos (CIFAR datasets).

They used powerful models (like BERT and Vision Transformers) as teachers and smaller, weaker versions as students.

The Results were clear:

Without help: The small students struggled. When forced to pick only a few "important" features (like 10% of the words or pixels), their accuracy dropped significantly. They couldn't figure out what mattered.
With REKD: The small students learned from the teacher's "rationales." Their accuracy jumped up, often performing almost as well as the big, powerful teachers, even though they were much smaller.
Bonus: In some cases, the students learned from the teacher's explanations so well that they actually performed better than they did when they were allowed to use all the data without any explanation constraints.

Why This Matters

The paper argues that we shouldn't just judge AI explanations by whether they match what a human thinks is important. Instead, we should judge them by whether the AI can actually make a good prediction using only the parts it highlighted.

By using a "Rationalist" teacher to guide the student, we can build smaller, faster AI models that are not only accurate but also capable of showing us why they made a decision, without needing massive computing power to do it.

In short: If you want a small AI to be smart and explainable, don't let it guess alone. Let it shadow a smart teacher who shows it exactly which clues to follow.

Technical Summary: Learn from A Rationalist: Distilling Intermediate Interpretable Rationales

Problem Statement
Deep neural networks (DNNs) have achieved superior performance in high-stakes domains but remain opaque "black boxes," raising trust concerns. While post-hoc explainability tools (e.g., LIME, SHAP) exist, they often fail to provide faithful explanations. Rationale Extraction (RE) offers a "select-predict" architecture where a generator selects features (rationales) and a predictor makes decisions based solely on those features, ensuring faithfulness. However, training RE models faces a "chicken and egg" dilemma: the generator needs the predictor to identify important features, while the predictor needs the generator to provide a clean input. This problem is exacerbated when using lightweight or less capable neural networks (students), which struggle to learn effective rationales and achieve high predictive performance without external guidance.

Methodology: REKD
The authors propose REKD (Rationale Extraction with Knowledge Distillation), a framework where a student RE model learns from a teacher RE model (a "rationalist") in addition to its own optimization. The method is neural-model agnostic, allowing any black-box architecture to serve as a backbone.

Rationale Extraction (RE) Component:
- Utilizes a Straight-Through Gumbel-Softmax estimator to make discrete feature selection differentiable.
- A generator produces unnormalized logits for feature selection, sampled via Gumbel-Softmax with a temperature parameter $\tau$ .
- A binary mask is derived via argmax (approximated as identity during backpropagation) to zero out unselected features.
- The predictor classifies based only on the masked rationale.
- The objective combines task loss (cross-entropy) and a selection loss (penalizing deviation from a target rationale ratio).
Knowledge Distillation (KD) Component:
- The student learns from the teacher's rationales (feature selection distributions) and predictions (class probabilities).
- Generator Distillation: Minimizes the Kullback-Leibler (KL) divergence between the student's and teacher's Gumbel-Softmax distributions across all features.
- Predictor Distillation: Minimizes the KL divergence between the student's and teacher's temperature-scaled output distributions.
Synchronized Temperature Annealing:
- A key innovation is the synchronization of the KD temperature with the Gumbel-Softmax annealing scheduler.
- The temperature $\tau$ decays exponentially from a high initial value ( $\tau_0$ ) to a low final value ( $\tau_K$ ).
- This creates an implicit curriculum learning strategy:
  - Early Phase (High $\tau$ ): The student absorbs "soft" knowledge from the teacher's broad distributions, guiding exploration.
  - Late Phase (Low $\tau$ ): The student transitions to mimicking the teacher's "hard," high-confidence feature selections as the distributions sharpen, enforcing precision during discretization.
- The total loss is a weighted combination of the RE objective and the KD objective.

Key Contributions

Identification of the Dilemma: The paper formally identifies and analyzes the "chicken and egg" training difficulty in RE, specifically highlighting how it hinders lightweight student models.
REKD Framework: Proposes a distillation framework that leverages the intrinsic curriculum of Gumbel-Softmax annealing to transfer interpretable feature selection knowledge from a teacher to a student.
Empirical Validation: Demonstrates the approach's viability across language (IMDB) and vision (CIFAR 10/100) tasks using BERT and Vision Transformer (ViT) variants, showing significant performance gains for student models.

Experimental Results
Experiments were conducted on IMDB (sentiment analysis), CIFAR 10, and CIFAR 100 (image classification) using BERT and ViT architectures.

Performance Improvement: REKD significantly improves the predictive accuracy of student RE models compared to standard RE training.
- CIFAR 10: ViT Small improved from 0.889 (RE) to 0.968 (REKD); ViT Tiny improved from 0.797 to 0.936.
- CIFAR 100: ViT Small improved from 0.779 to 0.845; ViT Tiny improved from 0.645 to 0.777.
- IMDB: BERT Small improved from 0.881 to 0.906; BERT Mini improved from 0.863 to 0.892.
Surpassing Teachers: In some cases (e.g., ViT Small on CIFAR 10 and CIFAR 100), the student REKD model achieved higher average accuracy than the teacher model itself, suggesting the distillation process effectively refines the student's ability to utilize the teacher's rationales.
Robustness: The method reduces the variance in performance across different random seeds.
Ablation: The authors note that the performance gain stems from the combination of RE and KD, with the teacher's rationales providing a strong regularization signal that prevents the student from learning spurious correlations.

Significance and Claims
The paper claims that REKD offers a practical solution to the training instability of interpretable models, particularly for resource-constrained scenarios (e.g., mobile devices). By framing the problem as learning from a "rationalist," the method aligns with human learning processes where novices learn from verified, interpretable knowledge provided by experts.

The authors emphasize that their evaluation metric focuses on predictive performance under rationale constraints rather than alignment with human annotations. They argue that faithfulness to the data (even if the data contains biases) is a more objective measure of a rationale's quality than human plausibility, as it reveals the model's actual reasoning. The method is presented as a step toward making DNNs more trustworthy in high-stakes domains without sacrificing the predictive power of large models, by allowing smaller models to inherit the "reasoning" of larger ones.