DomED: Redesigning Ensemble Distillation for Domain Generalization

Domain generalization aims to improve model performance on unseen, out-of-distribution (OOD) domains, yet existing methods often overlook the crucial aspect of uncertainty quantification in their predictions. While ensemble learning combined with knowledge distillation offers a promising avenue, naively combining these techniques is non-trivial and remains largely unexplored in the context of domain generalization. In this work, we systematically investigate different ensemble and distillation strategies for domain generalization tasks and design a tailored data allocation scheme. This approach trains teacher models on distinct subsets of domains and performs distillation on complementary (unseen) subsets, thereby fostering model diversity and training efficiency. Moreover, our theoretical analysis demonstrates that distilling from teachers on unseen domains effectively filters out domain-specific spurious correlations. To address the accuracy degradation often observed with standard uncertainty distillation, we further develop a novel technique that decouples uncertainty distillation from the standard distillation process, enabling accurate uncertainty estimation without compromising model accuracy. Our proposed method, Domain-aware Ensemble Distillation (DomED), is extensively evaluated against state-of-the-art domain generalization and ensemble distillation techniques across multiple benchmarks, achieving competitive accuracies and substantially improved uncertainty estimates.

We study how to improve artificial intelligence models when they encounter new environments that differ from their training data. In these situations, models often make incorrect predictions while remaining overly confident. We wanted to address this by transferring the knowledge of multiple models into a single efficient one, but doing this for new environments is challenging and frequently reduces accuracy. Our approach solves this by training several "teacher" models on different, specific environments. We then gather their predictions on environments they have never seen. We show that combining these diverse predictions effectively filters out irrelevant, environment-specific details. Next, we teach a single "student" model to reproduce both the teachers' predictions and their uncertainty. To prevent a drop in accuracy, we introduce a technique that separates learning the correct answer from learning the uncertainty level. Our experiments demonstrate that this method helps the final model achieve competitive accuracy on unseen data while significantly improving its ability to estimate its own uncertainty. This provides a practical way to build efficient models that can better indicate when their predictions might be incorrect.