Which Invariance Should We Transfer? A Causal Minimax Learning Approach

TL;DR

This paper introduces a causal minimax learning framework to identify the optimal subset of stable information for transfer, improving model robustness under dataset shifts.

Contribution

It provides a causal analysis for selecting the best stable information subset to transfer, along with an efficient algorithm for subset selection based on worst-case risk minimization.

Findings

The whole stable set is not always optimal for transfer.

The proposed algorithm efficiently finds the subset with minimal worst-case risk.

Method improves robustness in synthetic and Alzheimer's disease diagnosis data.

Abstract

A major barrier to deploying current machine learning models lies in their non-reliability to dataset shifts. To resolve this problem, most existing studies attempted to transfer stable information to unseen environments. Particularly, independent causal mechanisms-based methods proposed to remove mutable causal mechanisms via the do-operator. Compared to previous methods, the obtained stable predictors are more effective in identifying stable information. However, a key question remains: which subset of this whole stable information should the model transfer, in order to achieve optimal generalization ability? To answer this question, we present a comprehensive minimax analysis from a causal perspective. Specifically, we first provide a graphical condition for the whole stable set to be optimal. When this condition fails, we surprisingly find with an example that this whole stable set, although can fully exploit stable information, is not the optimal one to transfer. To identify the optimal subset under this case, we propose to estimate the worst-case risk with a novel optimization scheme over the intervention functions on mutable causal mechanisms. We then propose an efficient algorithm to search for the subset with minimal worst-case risk, based on a newly defined equivalence relation between stable subsets. Compared to the exponential cost of exhaustively searching over all subsets, our searching strategy enjoys a polynomial complexity. The effectiveness and efficiency of our methods are demonstrated on synthetic data and the diagnosis of Alzheimer's disease.