Distribution Shift Inversion for Out-of-Distribution Prediction

TL;DR

This paper introduces a novel Distribution Shift Inversion method that enhances out-of-distribution prediction by transforming test samples towards the training distribution using a diffusion model, without needing the test distribution during training.

Contribution

The paper proposes a portable distribution shift inversion algorithm that improves OoD prediction by transforming test samples without access to the test distribution during training.

Findings

Improves performance across multiple OoD algorithms.

Effective on both multiple-domain and single-domain datasets.

Theoretically justified approach.

Abstract

Machine learning society has witnessed the emergence of a myriad of Out-of-Distribution (OoD) algorithms, which address the distribution shift between the training and the testing distribution by searching for a unified predictor or invariant feature representation. However, the task of directly mitigating the distribution shift in the unseen testing set is rarely investigated, due to the unavailability of the testing distribution during the training phase and thus the impossibility of training a distribution translator mapping between the training and testing distribution. In this paper, we explore how to bypass the requirement of testing distribution for distribution translator training and make the distribution translation useful for OoD prediction. We propose a portable Distribution Shift Inversion algorithm, in which, before being fed into the prediction model, the OoD testing samples are first linearly combined with additional Gaussian noise and then transferred back towards the training distribution using a diffusion model trained only on the source distribution. Theoretical analysis reveals the feasibility of our method. Experimental results, on both multiple-domain generalization datasets and single-domain generalization datasets, show that our method provides a general performance gain when plugged into a wide range of commonly used OoD algorithms.