Improving Group Fairness in Knowledge Distillation via Laplace Approximation of Early Exits

TL;DR

This work tackles group fairness degradation in knowledge distillation by using Laplace-approximation-based uncertainty estimates from early-exit layers to reweight training losses. Building on margin-based reweighting, it treats uncertain early predictions with a probabilistic Last-Layer Laplace approach, yielding a predictive distribution $p(z_i|x_i) = \mathcal{N}(\hat{\mathbf{W}}_{\text{MAP}}^T\hat{\boldsymbol{\phi}}_i, \hat{\boldsymbol{\phi}}_i^T\Sigma_{\phi}\hat{\boldsymbol{\phi}}_i \cdot \mathbf{I}_C)$ and using the entropy of the averaged softmax to adjust losses. In experiments on MultiNLI with a BERT teacher and DistilBERT student, the method achieves small accuracy gains, with best results when the auxiliary reweighting occurs at an earlier exit (layer 3), and shows altered confidence margins that suggest different calibration dynamics. The approach is computationally lightweight and extends the DeDiER framework toward more robust, fairness-aware distillation, with clear avenues for broader validation across datasets and more sophisticated auxiliary architectures.

Abstract

Knowledge distillation (KD) has become a powerful tool for training compact student models using larger, pretrained teacher models, often requiring less data and computational resources. Teacher models typically possess more layers and thus exhibit richer feature representations compared to their student counterparts. Furthermore, student models tend to learn simpler, surface-level features in their early layers. This discrepancy can increase errors in groups where labels spuriously correlate with specific input attributes, leading to a decline in group fairness even when overall accuracy remains comparable to the teacher. To mitigate these challenges, Early-Exit Neural Networks (EENNs), which enable predictions at multiple intermediate layers, have been employed. Confidence margins derived from these early exits have been utilized to reweight both cross-entropy and distillation losses on a per-instance basis. In this paper, we propose that leveraging Laplace approximation-based methods to obtain well-calibrated uncertainty estimates can also effectively reweight challenging instances and improve group fairness. We hypothesize that Laplace approximation offers a more robust identification of difficult or ambiguous instances compared to margin-based approaches. To validate our claims, we benchmark our approach using a Bert-based model on the MultiNLI dataset.

Figures (3)

• Figure 1: Confidence margin at different feature levels for the worst-performing group compared to other groups in the original DEDIER paper.
• Figure 2: Confidence margins per layer on worst group predictions and all predictions
• Figure 3: Confidence margins in student and teacher on the predictions that were wrong