Adaptive Risk Minimization: Learning to Adapt to Domain Shift

TL;DR

This paper introduces adaptive risk minimization (ARM), a framework that enables models to adapt to domain shifts at test time using unlabeled data, improving robustness in real-world applications.

Contribution

The paper proposes ARM, a novel training framework that explicitly optimizes models for test-time adaptation to domain shifts, outperforming prior robustness methods.

Findings

ARM improves test accuracy by 1-4% on domain-shifted image classification tasks.

Models trained with ARM adapt effectively to new, unseen domains.

ARM outperforms existing invariance and robustness approaches in experiments.

Abstract

A fundamental assumption of most machine learning algorithms is that the training and test data are drawn from the same underlying distribution. However, this assumption is violated in almost all practical applications: machine learning systems are regularly tested under distribution shift, due to changing temporal correlations, atypical end users, or other factors. In this work, we consider the problem setting of domain generalization, where the training data are structured into domains and there may be multiple test time shifts, corresponding to new domains or domain distributions. Most prior methods aim to learn a single robust model or invariant feature space that performs well on all domains. In contrast, we aim to learn models that adapt at test time to domain shift using unlabeled test points. Our primary contribution is to introduce the framework of adaptive risk minimization (ARM), in which models are directly optimized for effective adaptation to shift by learning to adapt on the training domains. Compared to prior methods for robustness, invariance, and adaptation, ARM methods provide performance gains of 1-4% test accuracy on a number of image classification problems exhibiting domain shift.