Towards Unified and Effective Domain Generalization

TL;DR

UniDG is a unified, inference-time finetuning framework that improves out-of-distribution generalization of foundation models across various architectures by unsupervised learning and a penalty to prevent catastrophic forgetting.

Contribution

It introduces a novel inference-stage finetuning method with a penalty to enhance domain generalization without additional training.

Findings

Average accuracy improvement of +5.4% on DomainBed

Effective across 12 diverse visual backbones

Reduces catastrophic forgetting during finetuning

Abstract

We propose $\textbf{UniDG}$, a novel and $\textbf{Uni}$fied framework for $\textbf{D}$omain $\textbf{G}$eneralization that is capable of significantly enhancing the out-of-distribution generalization performance of foundation models regardless of their architectures. The core idea of UniDG is to finetune models during the inference stage, which saves the cost of iterative training. Specifically, we encourage models to learn the distribution of test data in an unsupervised manner and impose a penalty regarding the updating step of model parameters. The penalty term can effectively reduce the catastrophic forgetting issue as we would like to maximally preserve the valuable knowledge in the original model. Empirically, across 12 visual backbones, including CNN-, MLP-, and Transformer-based models, ranging from 1.89M to 303M parameters, UniDG shows an average accuracy improvement of +5.4% on DomainBed. These performance results demonstrate the superiority and versatility of UniDG. The code is publicly available at https://github.com/invictus717/UniDG