Adaptive Group Robust Ensemble Knowledge Distillation

TL;DR

This paper introduces AGRE-KD, a novel ensemble knowledge distillation method that improves performance on underrepresented subgroups by selectively combining teacher models, outperforming traditional ensemble approaches.

Contribution

The paper proposes AGRE-KD, a new ensemble distillation strategy that enhances subgroup robustness by selectively integrating debiased teacher models based on gradient analysis.

Findings

AGRE-KD outperforms traditional ensemble distillation in subgroup performance.

The method effectively leverages an additional biased model to improve worst-case subgroup accuracy.

Experiments show AGRE-KD surpasses classic ensemble methods like majority voting.

Abstract

Neural networks can learn spurious correlations in the data, often leading to performance degradation for underrepresented subgroups. Studies have demonstrated that the disparity is amplified when knowledge is distilled from a complex teacher model to a relatively ``simple'' student model. Prior work has shown that ensemble deep learning methods can improve the performance of the worst-case subgroups; however, it is unclear if this advantage carries over when distilling knowledge from an ensemble of teachers, especially when the teacher models are debiased. This study demonstrates that traditional ensemble knowledge distillation can significantly drop the performance of the worst-case subgroups in the distilled student model even when the teacher models are debiased. To overcome this, we propose Adaptive Group Robust Ensemble Knowledge Distillation (AGRE-KD), a simple ensembling strategy to ensure that the student model receives knowledge beneficial for unknown underrepresented subgroups. Leveraging an additional biased model, our method selectively chooses teachers whose knowledge would better improve the worst-performing subgroups by upweighting the teachers with gradient directions deviating from the biased model. Our experiments on several datasets demonstrate the superiority of the proposed ensemble distillation technique and show that it can even outperform classic model ensembles based on majority voting. Our source code is available at https://github.com/patrikken/AGRE-KD