Set Valued Predictions For Robust Domain Generalization

TL;DR

This paper proposes a novel set-valued prediction framework for domain generalization, aiming to improve robustness across unseen domains by balancing prediction set size and performance, supported by theoretical insights and empirical validation.

Contribution

It introduces a theoretical framework and practical optimization method for set-valued predictions in domain generalization, enhancing robustness and providing new insights into achievable generalization conditions.

Findings

Set-valued predictors improve robustness on unseen domains.

The proposed method balances small prediction sets with high performance.

Empirical results on WILDS datasets demonstrate effectiveness.

Abstract

Despite the impressive advancements in modern machine learning, achieving robustness in Domain Generalization (DG) tasks remains a significant challenge. In DG, models are expected to perform well on samples from unseen test distributions (also called domains), by learning from multiple related training distributions. Most existing approaches to this problem rely on single-valued predictions, which inherently limit their robustness. We argue that set-valued predictors could be leveraged to enhance robustness across unseen domains, while also taking into account that these sets should be as small as possible. We introduce a theoretical framework defining successful set prediction in the DG setting, focusing on meeting a predefined performance criterion across as many domains as possible, and provide theoretical insights into the conditions under which such domain generalization is achievable. We further propose a practical optimization method compatible with modern learning architectures, that balances robust performance on unseen domains with small prediction set sizes. We evaluate our approach on several real-world datasets from the WILDS benchmark, demonstrating its potential as a promising direction for robust domain generalization.