Fishr: Invariant Gradient Variances for Out-of-Distribution Generalization

TL;DR

Fishr introduces a novel regularization technique that aligns domain-specific gradient variances to improve out-of-distribution generalization, outperforming existing methods on benchmark datasets.

Contribution

The paper proposes Fishr, a new regularization method that enforces domain invariance in gradient space, linking gradient covariance with Fisher Information and Hessian to enhance OOD generalization.

Findings

Fishr outperforms state-of-the-art on DomainBed benchmark.

Fishr consistently beats Empirical Risk Minimization.

Gradient variance alignment improves robustness across domains.

Abstract

Learning robust models that generalize well under changes in the data distribution is critical for real-world applications. To this end, there has been a growing surge of interest to learn simultaneously from multiple training domains - while enforcing different types of invariance across those domains. Yet, all existing approaches fail to show systematic benefits under controlled evaluation protocols. In this paper, we introduce a new regularization - named Fishr - that enforces domain invariance in the space of the gradients of the loss: specifically, the domain-level variances of gradients are matched across training domains. Our approach is based on the close relations between the gradient covariance, the Fisher Information and the Hessian of the loss: in particular, we show that Fishr eventually aligns the domain-level loss landscapes locally around the final weights. Extensive experiments demonstrate the effectiveness of Fishr for out-of-distribution generalization. Notably, Fishr improves the state of the art on the DomainBed benchmark and performs consistently better than Empirical Risk Minimization. Our code is available at https://github.com/alexrame/fishr.