Learning Variational Neighbor Labels for Test-Time Domain Generalization

This paper strives for domain generalization, where models are trained exclusively on source domains before being deployed on unseen target domains. We follow the strict separation of source training and target testing, but exploit the value of the unlabeled target data itself during inference. First, we propose probabilistic pseudo-labeling of target samples to generalize the source-trained model to the target domain at test time. We formulate the generalization at test time as a variational inference problem, by modeling pseudo labels as distributions, to consider the uncertainty during generalization and alleviate the misleading signal of inaccurate pseudo labels. Second, we learn variational neighbor labels that incorporate the information of neighboring target samples to generate more robust pseudo labels. Third, to learn the ability to incorporate more representative target information and generate more precise and robust variational neighbor labels, we introduce a meta-generalization stage during training to simulate the generalization procedure. Experiments on seven widely-used datasets demonstrate the benefits, abilities, and effectiveness of our proposal. As soon as test data distributions differ from the ones experienced during training, deep neural networks start to exhibit generalizability problems and accompanying performance degradation (Geirhos et al., 2018; Recht et al., 2019). To deal with distribution shifts, domain generalization (Li et al., 2017; 2020; Motiian et al., 2017b; Muandet et al., 2013) has emerged as a promising tactic for generalizability to unseen target domains. However, as methods are only trained on source domains, this may still lead to overfitting and limited performance guarantees on unseen target domains. To better adapt models to target domains - without relying on target data during training - test-time adaptation (Liang et al., 2023; Sun et al., 2020; Varsavsky et al., 2020; Wang et al., 2021) was introduced. It provides an alternative learning paradigm by training a model on source data and further adjusting the model according to the unlabeled target data at test time.