Improving Domain Generalization with Domain Relations

TL;DR

The paper tackles distribution and domain shifts by introducing D3G, which builds domain-specific predictors for each training domain and uses domain-relations to reweight these predictors at test time. It combines a multi-head training architecture with a relation-aware consistency loss, and learns a domain-similarity matrix by blending fixed meta-data-based relations with learned relations derived from domain attributes $m_i$, via $a_{ij}=\beta a_{ij}^g+(1-\beta)a_{ij}^l$. The authors provide a theoretical excess-risk bound for the test domain and demonstrate that relation-informed weighting reduces generalization error, complemented by extensive experiments on DG-15 and real-world domain-shift datasets (TPT-48, FMoW, ChEMBL-STRING) showing an average improvement of $10.6\%$ over prior methods. Overall, D3G leverages domain meta-data and relational refinement to achieve robust out-of-domain generalization, with reproducibility and code release planned.

Abstract

Distribution shift presents a significant challenge in machine learning, where models often underperform during the test stage when faced with a different distribution than the one they were trained on. This paper focuses on domain shifts, which occur when the model is applied to new domains that are different from the ones it was trained on, and propose a new approach called D$^3$G. Unlike previous methods that aim to learn a single model that is domain invariant, D$^3$G leverages domain similarities based on domain metadata to learn domain-specific models. Concretely, D$^3$G learns a set of training-domain-specific functions during the training stage and reweights them based on domain relations during the test stage. These domain relations can be directly obtained and learned from domain metadata. Under mild assumptions, we theoretically prove that using domain relations to reweight training-domain-specific functions achieves stronger out-of-domain generalization compared to the conventional averaging approach. Empirically, we evaluate the effectiveness of D$^3$G using real-world datasets for tasks such as temperature regression, land use classification, and molecule-protein binding affinity prediction. Our results show that D$^3$G consistently outperforms state-of-the-art methods.

Figures (6)

• Figure 1: An illustration of D3G. (a) The multi-headed architecture of D3G, where each training domain is associated with a single head for prediction. (b) The relation extraction module, where fixed relations are extracted from domain meta-data and refined through learning from the same meta-data. (c) The training stage of D3G, where $x$ represents a single example from domain $d$, and the loss is composed of both a supervised loss and a consistency loss. (d) The test stage, where the weighting of all training domain-specific functions is used to perform inference for each test example.
• Figure 2: Results of domain shifts on toy task (DG-15). Figures (a) and (b) illustrate the training and test distributions, where datapoints in circles with the same color originate from the same domain. Figures (c) and (d) show the predicted distribution of the strongest single model method (GroupDRO) and D3G. Bottom Table reports averaged accuracy over all test domains (see full table with standard deviation in Appendix \ref{['app:exp_toy']}). We bold the best results and underline the second best results.
• Figure 2: Comparison between D3G with domain-specific fine-tuning. Full results: Appendix \ref{['app:domain_ft']}.
• Figure 3: Performance comparison w.r.t. consistency regularization. Only fixed relations are used.
• Figure 4: Analysis of relation learning. Top figures show the multi-unit residential areas of Turkey, Syria, and Saudi Arabia. Bottom figures illustrate both fixed relations and learned relations.

Theorems & Definitions (2)

• Theorem 4.1
• Proposition 4.2