On the Connection between Invariant Learning and Adversarial Training for Out-of-Distribution Generalization

TL;DR

This paper explores the connection between invariant risk minimization and adversarial training, proposing a new method called Domainwise Adversarial Training (DAT) that improves out-of-distribution generalization by removing domain-specific features.

Contribution

It reveals the relationship between IRM and adversarial training and introduces DAT, a novel adversarial training-based approach for better domain generalization.

Findings

DAT effectively removes domain-varying features.

DAT improves OOD generalization under various shifts.

Experimental results show superior performance of DAT.

Abstract

Despite impressive success in many tasks, deep learning models are shown to rely on spurious features, which will catastrophically fail when generalized to out-of-distribution (OOD) data. Invariant Risk Minimization (IRM) is proposed to alleviate this issue by extracting domain-invariant features for OOD generalization. Nevertheless, recent work shows that IRM is only effective for a certain type of distribution shift (e.g., correlation shift) while it fails for other cases (e.g., diversity shift). Meanwhile, another thread of method, Adversarial Training (AT), has shown better domain transfer performance, suggesting that it has the potential to be an effective candidate for extracting domain-invariant features. This paper investigates this possibility by exploring the similarity between the IRM and AT objectives. Inspired by this connection, we propose Domainwise Adversarial Training (DAT), an AT-inspired method for alleviating distribution shift by domain-specific perturbations. Extensive experiments show that our proposed DAT can effectively remove domain-varying features and improve OOD generalization under both correlation shift and diversity shift.