Joint Variational Inference Network for domain generalization

Abstract

Domain generalization is a machine learning task that involves training a model on a set of domains with the goal of achieving high performance on unseen domains. While most domain generalization methods focus on extracting domain-invariant features, they may ignore domain-specific information beyond object styles which is label-relevant for classification. In this paper, we first propose a two-step data generation process in which the domain label and observed data are sampled from two distributions sequentially, and accordingly develop an end-to-end Joint Variational Inference Network (JVINet) for domain generalization. JVINet is a framework that admits a variational-based discriminative structure, which the domain-specific latent variable is learned and integrated to enhance the feature for classification. When testing on unseen target domains, the potential beneficial domain information is hence utilized to improve generalization ability. The overall objective is to optimize the variational lower bound of the conditional joint likelihood functions for both class and domain labels. We provide theoretical proof that JVINet can achieve an optimal lower bound using a variational-based discriminative approach. To evaluate the effectiveness of our method, we compare it with state-of-the-art methods on both simulated and real-world datasets.