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

TL;DR

This paper introduces a Laplace approximation-based method to improve group fairness in knowledge distillation by better identifying challenging instances, especially in early-exit neural networks, demonstrated on a Bert model with MultiNLI.

Contribution

It proposes using Laplace approximation for uncertainty estimation to reweight instances, enhancing group fairness in knowledge distillation with early-exit neural networks.

Findings

Laplace approximation improves identification of difficult instances.

Enhanced group fairness observed in Bert-based models on MultiNLI.

Outperforms margin-based reweighting methods in fairness metrics.

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.