Set Valued Predictions For Robust Domain Generalization

Ron Tsibulsky, Daniel Nevo, Uri Shalit

TL;DR: We develop a method for robust out-of-domain set-valued prediction, give theoretical performance guarantees, and show our method works well on several WILDS datasets

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.

Lay Summary:

Modern machine learning systems often struggle when faced with new or different data than what they were trained on. This challenge is known as domain generalization. In this work, we explore how predicting sets of possible answers, rather than just a single answer, can help models remain accurate and reliable when tested on unfamiliar data.

We introduce a new way to define success for these set-based predictions and provide both theoretical foundations and practical tools to achieve it. A key challenge we tackle is the tradeoff between making prediction sets small, so they remain specific and useful, and ensuring they are large enough to consistently include the correct answer, even in environments the model hasn't seen before.

We evaluate our approach on real-world datasets and demonstrate that it achieves strong performance across diverse test conditions, while keeping the prediction sets as compact as possible. Our findings suggest a promising path for building machine learning systems that are both robust and adaptable in practice.

Application-Driven Machine Learning: This submission is on Application-Driven Machine Learning.

Link To Code: https://github.com/ront65/set-valued-ood

Primary Area: General Machine Learning->Transfer, Multitask and Meta-learning

Keywords: Set Valued Prediction, Domain Generalization, PAC Learnability

Submission Number: 12812