Learning Variational Neighbor Labels for Test-Time Domain Generalization

We thank Reviewer 4tVC for their feedback.

Graphical model

We agree with the reviewer that the graphical model might be misleading. We meant to express the conditional dependence structure between random variables (data and latent variables). We changed it to “Probabilistic modeling graph” in the updated paper.

Estimation conditional probabilities

(1) For the prior distribution $p_\phi (w_t | X_t, \theta_s)$, we first generate the features $z_t$ of the batch of images $X_t$ and their pseudo labels $\hat{Y}_{t}$ through the source trained model $\theta_s$.

Then we generate the averaged features $z_t^c$ for each category according to the pseudo labels. The categorical averaged features $z_t^c$ are then fed to the MLP network $\phi$ to generate the mean and variance of $p_\phi(w_t)$. $X_t$ here is not a random variable. It denotes the batch of images, similar to the image $x_t$.

(2) For the posterior distribution $q_\phi (w_t | X_t, Y_t, \theta_s)$, the reviewer is right that we never observe $Y_t$, so the distribution is intractable at test time. Therefore we use only the prior distribution for the test and introduce meta-generalization to estimate the posterior distribution for the meta-target data during training. The meta-target domain is simulated by the multiple source domains, so we can get $Y_{t’}$ during training. Similar to the prior distribution, we first generate the features $z_{t’}$ of the batch of images $X_{t’}$. Differently, the categorical averaged features are generated according to the ground truth meta-target labels $Y_{t’}$. Then we use the same MLP $\phi$ to generate q(w_t) by the categorical averaged features.

Rationale of meta-generalization and graph

Meta-generalization is utilized to simulate distribution shifts during training by splitting the multiple source domains into meta-source and meta-target domains. By doing so, we learn to generalize over shifts between meta-source and meta-target distributions. At test time, we leverage this learned ability to generalize over the unseen target data. The meta-generalization mechanism alleviates the overfitting of directly learning a model on source data.

Under meta-generalization, the graph can be directly obtained by replacing the source and target variables in Figure 1 (c) with meta-source and meta-target variables.

Gradient descent makes it hard to construct a probabilistic term

We agree with the reviewer that gradient descent makes it hard to construct a probabilistic term. Since the distribution of the model parameters is unknown, we have to rely on an empirical Bayesian method with gradient descent to approximate $\theta_t$ (Chelsea Finn et al., NeurIPS 2018).

Note we only use this approximation to estimate the final model parameter $\theta_t$, the variational inference is used to generate the pseudo labels. During the generation of the variational neighbor labels, we never use gradient descent to approximate the probabilistic term.

We have clarified this in the methodology section.

Reason to use variational inference

Our key idea is to capture the uncertainty of the pseudo labels in the latent space for more robust generalization across distribution shifts. The probabilistic framework encodes the uncertainty and more target information for model generalization (e.g., the sample in Figure 6 in the appendix).

To optimize the probabilistic framework, we use variational inference to approximate the true posterior of the probabilistic pseudo labels, which inherently introduces more neighboring target information and categorical information during training. Our experimental results in Table 1 (as well as the additional ablations by Reviewer Hsmr) confirm the variational pseudo labels outperform the deterministic ones, especially when combined with our proposed meta-generalization.

We have added these discussions to the main paper.

Why does the method improve the pseudo label quality and when

By meta-generalization we simulate distribution shifts during training, so the model learns to generate more effective pseudo labels for fine-tuning the model across distribution shifts. The variational neighbor labels are further improved by considering more neighboring target information.

It works because we have multiple source domains available during training to mimic distribution shifts and the availability of a batch of target samples for neighboring information at test time.

We have added this discussion to the methodology section of the paper.

We thank Reviewer C6Ah for their feedback.

Ambiguity in Contribution

We agree that meta-generalization is important and helps us achieve a good improvement. By simulation of distribution shifts during training the model learns the ability to better encode uncertainty in the variational neighbor labels and improve generalization across distribution shifts, see also the requested ablation by Reviewer Hsmr. We will better stress its importance in the updated paper.

Limited validation

We perform experiments on seven datasets that cater for (test-time) domain generalization, as we require multiple domains for training (similar to Gulrajani & Lopez-Paz, 2020; Iwasawa & Matsuo., 2021; Xiao et al., 2022). The suggested datasets are all single-source and used for (test-time) robustness to image corruptions. Nevertheless, we conducted the experiment on Cifar-10-C and obtained 21.60 as the error rate, where the source model without adaptation achieves 29.14 and Tent (Wang et al., 2021) obtains 13.95. We have added this comparison in the appendix. The dependence on multiple source domains during training is both a strength and a weakness of our proposal. We consider single-source domain generalization at test time a worthwhile avenue for future work, as highlighted in the conclusion.

Regarding utilizing information from the target domain training samples
We would like to stress that we do not utilize any target samples during training. Instead, we mimic domain shifts during source training by obtaining a “meta target” set from the source domains. Specifically, during source training, we use one source domain as the meta-source and the other source domains as the meta-target in each iteration (this is the reason why we depend on datasets providing multiple domains).

Benefits of meta-generalization

The benefit of meta-generalization is to simulate distribution shifts during training. By doing so, the model learns the ability of generalization across domain shifts. At test time, the method can use the learned generalization ability for the (unseen) distribution shift between source and target data, leading to more robust generalization compared with methods without meta-learning. To further show the benefits of mimicking distribution shifts by meta-generalization. We conduct more ablations based on meta-generalization. As in the answer of Reviewer Hsmr, meta-generalization can also improve the performance of common pseudo-labeling and probabilistic pseudo-labeling, showing the effectiveness of mimicking distribution shifts. Meta-generalization of variational neighboring labels gains more improvement and the best overall results.

We also considered meta-generalization with fewer distribution shifts during training. When we train on two instead of three source domains from PACS we obtain 80.32±1.8 on the “art-painting” target set, where using all three domains gives 83.81. Showing the importance of mimicking distribution shifts for meta-generalization.

Difference between \L_{ce} and \L_{meta}

The negative sign is a typo, it should be a positive sign. We have updated it accordingly. $L_{ce}$ calculates the cross-entropy between neighbor labels and the ground truth, while $L_{meta}$ calculates the cross-entropy between the updated model's prediction and ground truth.

Performance degradation on Office-Home

The reviewer is correct. We attribute the performance degradation for Office-Home to the larger number of 65 categories, demanding a larger model capacity $\phi$. We have varied $\phi$ with different numbers of layers in the MLP for Office-Home. We obtain a mean accuracy of 57.1 with 2 layers and 64.3 with 3 layers. We have discussed this in the main paper and kept the details of the experiment in the appendix. Thank you.

Clarification about n samples in Algorithm 1

We just use the numbers “n” provided by the public datasets. Our method is not sensitive to “n” since we randomly select the meta-source and meta-target samples in each iteration.

Additional hyperparameters

Similar to recent methods, we have ablated standard test-time hyperparameters such as batch size. Additionally, since we use a small $\phi$ network we have also ablated its model capacity and added it in the main paper. There are no additional hyperparameters involved. We have clarified this in the main paper.

Code
We added more implementation details and updated the code, and will include a link to the public github repository in the final version.

We thank Reviewer JCc1 for their feedback.

Scale of datasets

We agree with the reviewer that the commonly used datasets PACS (9,991 images with 7 classes) and VLCS (10,729 images with 5 classes) are relatively small-scale. That is why we also conducted experiments on the larger-scale datasets Office-Home (15,588 images with 65 classes) and mini-DomainNet (140,000 images with 126 classes) in the appendix. On all datasets our method achieves competitive performance.

Since we utilize meta-learning to mimic domain shifts and neighboring target information, we indeed require multiple source domains (Iwasawa & Matsuo., 2021) during training as supported by the domain generalization datasets PACS, VLCS, Terra Incognita, Office-Home and mini-DomainNet. Consequently, it is difficult to conduct our method on single-source datasets intended for image corruption robustness, like CIFAR-10C (70,000 images with 10 classes) and ImageNet-C (14,197,122 images with 1000 classes). Nevertheless, we conducted the experiment on Cifar-10-C and obtained 21.60 as the error rate, where the source model without adaptation achieves 29.14 and Tent (Wang et al., 2021) achieves 14.30. We have added this comparison in the appendix. The dependence on multiple source domains during training is both a strength and a weakness of our proposal. We consider single-source domain generalization at test time a worthwhile avenue for future work, as highlighted in the conclusion.

Improvement over prior work

We achieve state-of-the-art on 6 of 7 domain generalization datasets, and the improvement per dataset varies from 0.4% (VLCS) to 2.1% (Terra Incognita). We observe our peers report similar improvement gains (e.g., Iwasawa & Matsuo, 2022, Chen et al., 2023b, Jang et al., 2023, Xiao et al., 2022), while achieving state-of-the-art on fewer datasets (e.g., Iwasawa & Matsuo, 2022 on 2 out of 4 datasets reported; Chen et al., 2023 on 3 out of 5 datasets reported). Alongside its empirical performance, our method also achieves better calibration (Figure 2) and addresses complex scenarios (Figure 5).

Presentation

Thank you for the suggestions. We have relocated related work and used the more meaningful codename ‘VNL’ for our method in the updated paper.

Other datasets hyperparameters

We use similar settings and hyperparameters for all domain generalization benchmarks that are reported in the main paper. We have moved the description from the appendix to the main paper.

Inference speed and computational cost

Although we are not as fast at inference as the online classifier adjustment method by Iwasawa & Matsuo., (2021) and BN fine-tuning methods (Wang et al., 2021 (BN)), we are comparable to other test-time domain generalization methods (Jang et al., 2023, Liang et al., 2020) as shown in the following table.

We introduce about 1% more parameters than the backbone model, while keeping computational costs low. We have mentioned the result of this experiment and the computational cost in the main paper and kept the Table in the appendix.

We thank Reviewer Hsmr for their feedback.

delta symbol

The delta symbol indicates the Dirac delta distribution, also known as unit impulse. The approximation of $p(\theta_t|\theta_s, x_t)$ by $\delta(\theta_t=\theta_t^*)$ indicates we use the deterministic MAP value $\theta_t^*$ to approximate the distribution $p(\theta_t)$. We have clarified this in the updated paper.

More ablations for meta-generalization

Following your suggestion, we report additional ablations on PACS: meta-generalization with pseudo-labeling, and meta-generalization with probabilistic pseudo-labeling. Meta-generalization with the proposed variational pseudo-labeling improves the most. We have added these results in the experiments section.