Explainability-Aligned Distillation

• Explainability-Aligned Distillation is a paradigm that ensures a student model replicates both the output predictions and the reasoning of its teacher using auxiliary alignment constraints.
• Methodologies such as independent imitation and joint loss functions balance prediction accuracy with explanation fidelity, maintaining consistency across different architectures.
• Empirical studies on benchmarks like MNIST and Zillow demonstrate improved explanation alignment and predictive performance, reinforcing its applicability in high-stakes AI.

Explainability-Aligned Distillation is a paradigm in knowledge distillation that requires a student (distilled) model to not only approximate the output predictions of its teacher model, but also to align its underlying reasoning or explanations with those of the teacher. In contrast to classical distillation—which focuses purely on matching predictive performance—explainability-aligned distillation introduces explicit constraints or auxiliary tasks to ensure that the student’s rationale for each decision closely mirrors that of its teacher. This alignment serves to enhance model trust, interpretability, and potentially generalization, especially when the teacher and student architectures differ substantially.

1. Core Principles of Explainability Alignment

The central tenet is that a distilled student model must deliver “similar predictions for similar reasons.” For any sample $x$, the teacher generates an output $\hat{y}_T(x)$ and an accompanying explanation vector $\Phi_T(x)$, typically comprising feature attributions or importances (e.g., via structure-based methods like Sabbas or SHAP). Explainability-aligned distillation enforces that the student produces both a comparable prediction $\hat{y}_S(x)$ and an explanation $\Phi_S(x)$, where the latter is optimized to approximate $\Phi_T(x)$.

This paradigm unifies fidelity in prediction (the “what”) and fidelity in rationale (the “why”), so that even if model architectures diverge, the decision justifications are preserved post-distillation.

2. Methodological Frameworks

Several methodological variants have been proposed to operationalize explainability alignment:

• Independent Imitation: The student first learns embeddings or intermediate representations optimized for output prediction; subsequently, a separate mapping is trained to map from the student’s internal representations to the teacher’s explanation vector, typically using an $l_2$ or MSE loss:

$$\min_{w, w_0} \mathbb{E}_x\, [l_2(w^\top NN(x; \theta_{ped}) + w_0,\, \Phi_T(x))]$$

• Joint Loss Functions: A unified objective explicitly balances prediction matching and explanation alignment:

$$\min_w \mathbb{E}_x\,[\lambda\, l_1(\text{prediction loss}) + (1-\lambda) l_2(\text{explanation loss})]$$

where $\lambda \in [0, 1]$ trades off accuracy versus interpretability.

• Case Study—GBDT2NN: In the examined work, distillation from tree ensembles (GBDT) to neural networks is performed by embedding the leaf indices for each tree and then combining these with loss terms for both prediction output ($l_1$) and explanation vectors ($l_2$), derived from trusted structure-based methods.

These methods ensure alignment both at output and interpretability levels, and can be extended with more complex multi-task or adaptive loss balancing strategies.

3. Practical Implementation and Empirical Findings

Explainability-aligned distillation has been empirically validated on benchmarks such as AutoML-3, OnlineNewsPop, MNIST (as digit-0 vs digit-4), and Zillow. Performance is assessed via both standard predictive metrics (e.g., AUC for classification, MSE for regression) and explanation fidelity measures:

• NDCG (Normalized Discounted Cumulative Gain): Measures the ranking concordance of important features between teacher and student explanations.
• Top-$k$ Coverage: Evaluates the intersection of top-ranked features in teacher and student, with higher overlap indicating better explanation alignment.

Empirically:

• Both independent and joint alignment methods improve the agreement of student explanations with those of the teacher, as demonstrated by higher NDCG and top-$k$ coverage.
• Joint training additionally improves prediction performance (e.g., slight gains in AUC/MSE and faster convergence), indicating synergy between explanation and output alignment.
• Visualization (e.g., for MNIST) shows the student capturing critical image regions identified by the teacher.

4. Theoretical Foundations and Objective Formulations

Several key formulations underpin the approach:

Element	Mathematical Formulation	Description
Leaf Embedding Loss	$$\min_{\theta_{ped}} \mathbb{E}_x\,[l_1(NN(x; \theta_{ped}), L_t(x))]$$	Matches NN hidden state to tree leaf embeddings
Explanation Mapping	$$\min_{w, w_0} \mathbb{E}_x\,[l_2(w^\top NN(x; \theta_{ped}) + w_0,\, \Phi_T(x))]$$	Aligns student inferences with teacher attributions
Joint Loss (2-view)	$$\min_w \mathbb{E}_x\,[\lambda\, l_1(...) + (1-\lambda) l_2(...)]$$	Simultaneous optimization for both losses

These losses provide explicit and tunable mechanisms to incorporate interpretability into model compression.

5. Implications, Generalization, and Future Directions

Imposing interpretability constraints during distillation is significant when:

• The architectures are heterogeneous (e.g., tree-to-NN, GBDT2NN), as student “black-boxes” otherwise lose access to structured teacher rationale.
• Stakeholder trust and model auditability are critical, since the student model can now expose its reasoning traceable to established, transparent teacher decisions.
• Model debugging and feature importance analysis are needed post-compression.

Potential research extensions include:

• Extending explainability-aligned distillation to other teacher–student pairings (e.g., transformer-to-CNN).
• Adaptive or confidence-weighted loss balancing.
• Evaluating explanation stability versus model complexity/noise or integrating different explainer methods.
• Applying this paradigm in sequence modeling, time-series forecasting, or natural language understanding.

A plausible implication is that as the field advances, explainability-aligned distillation may become the standard for high-stakes AI deployment, where trust, accountability, and the clear tracing of predictions to interpretable rationales are paramount.

6. Summary

Explainability-Aligned Distillation, as initiated in the context of GBDT2NN and formalized through empirical and joint loss frameworks, reformulates the knowledge distillation objective to include the transfer of reasoning or attribution vectors, not just label outputs. Empirical and theoretical analysis demonstrate improved interpretability and, often, improved predictive performance. By constraining student models to “think like” their teachers even across architectural boundaries, this paradigm enhances transparency and trust in compressed and deployed AI systems (Huang et al., 2020).