Good Teachers Explain: Explanation-Enhanced Knowledge Distillation

Overview

• Knowledge Distillation (KD) is a technique for compressing large teacher models into smaller student models while maintaining similar accuracy.
• While student models can achieve similar performance as teachers, they often do not learn the same underlying function.
• It is desirable for students to learn the "right features" from teachers by basing predictions on similar input features.

Plain English Explanation

Knowledge Distillation (KD) is a way to take a large, complex machine learning model (the "teacher") and distill its knowledge into a smaller, more efficient model (the "student"). The goal is for the student model to achieve similar accuracy to the teacher, but use fewer computational resources.

One challenge is that even though the student and teacher models may have similar performance, the student doesn't always learn the same underlying function as the teacher. In other words, the student may be making the right predictions, but for the wrong reasons.

The researchers in this paper explored whether they could address this by not just optimizing the classic KD loss, but also encouraging the student to learn explanations similar to the teacher's. Their "explanation-enhanced" KD (e²KD) approach:

1. Consistently improves the student's accuracy and agreement with the teacher's predictions.
2. Ensures the student learns to be "right for the right reasons" by matching the teacher's explanations.
3. Works robustly across different model architectures, amounts of training data, and even with approximate pre-computed explanations.

By aligning the student's explanations with the teacher's, the researchers found a way for the student to truly learn the same underlying function as the teacher, not just mimic its outputs.

Technical Explanation

The key innovation in this paper is the "explanation-enhanced" Knowledge Distillation (e²KD) approach. Typical KD optimizes the student model to match the teacher's output predictions, but the researchers hypothesized that this alone does not guarantee the student learns the same underlying function.

To address this, they proposed also optimizing the student to match the teacher's

for its predictions, in addition to the predictions themselves. These explanations could come from techniques like feature importance or activation maps.

Through extensive experiments, the researchers showed that e²KD:

1. Improves Accuracy and Agreement: e²KD consistently led to higher student model accuracy compared to standard KD. It also increased the agreement between the student and teacher's predictions.

2. Learns "Right Reasons": e²KD ensured the student learned to make predictions based on the same underlying factors as the teacher, not just mimicking outputs.

3. Robustness: e²KD worked well across different model architectures, amounts of training data, and even when using approximate pre-computed explanations.

Overall, the key insight is that optimizing not just prediction outputs but also explanation similarity can help the student truly learn the teacher's function, not just replicate its behavior.

Critical Analysis

The researchers provide a comprehensive evaluation of their e²KD approach, including comparisons to standard KD and ablation studies. The results are compelling and the method seems to reliably improve student performance.

One potential limitation is the reliance on having access to the teacher's explanations. In some cases, these may not be available or easy to compute. The researchers did show their method can work with approximate explanations, but this could still be a practical challenge.

Additionally, the paper focuses on classification tasks. It would be interesting to see how e²KD performs on other types of machine learning problems, such as regression or structured prediction.

Finally, the researchers mention that e²KD could be combined with other KD techniques, like label revision or feature-based distillation. Exploring these hybrid approaches could lead to even stronger student models.

Overall, this is a well-executed study that offers a promising new direction for improving knowledge distillation by aligning student and teacher explanations.

Conclusion

This paper presents an "explanation-enhanced" Knowledge Distillation (e²KD) approach that goes beyond just matching prediction outputs between teacher and student models. By also optimizing the similarity of their explanations, the researchers found a way for students to truly learn the same underlying function as their teachers, not just mimic their behavior.

The e²KD method consistently improves student accuracy and agreement with teachers, while ensuring students learn the "right reasons" for their predictions. This robustly works across different model architectures and data regimes, making it a valuable tool for compressing large, complex models into more efficient versions without sacrificing performance or interpretability.

As machine learning models grow ever larger and more powerful, techniques like e²KD will be increasingly important for deploying these capabilities widely while managing computational and energy constraints. This work represents an important step forward in that direction.