Distribution Shift-Aware Prediction Refinement for Test-Time Adaptation

Thank you for the positive feedback and summary of our work. We hope the responses to the weaknesses below resolve your concerns.

(W1) the results are limited to severe label shifts, while the effectiveness of the proposed method under only covariate shift is not well studied

We focus on the problem setup where both covariate and label distribution shifts occur during the test time. In Section 2, we first examine the impact of the label distribution shift on BNAdapt, which is known to effectively address covariate shifts, thus serving as the foundation for entropy minimization-based TTA methods including TENT. We observe that BNAdapt suffers from a significant performance degradation under the label distribution shift, since the updated BN statistics follow the test label distribution, which causes bias in the classifier. In Section 3, we propose a novel refinement scheme for the predictions generated by the BN-adapted classifiers under both the covariate and label distribution shifts. This refinement mainly corrects the misclassification generated from the label distribution shifts. Thus, without the label distribution shift, our module does not have a big impact in adjusting the classifier predictions. We actually observe this from the oracle analysis in Table 1 for CIFAR-10C-LT with $\rho=1$ (balanced case), where a limited gain of 0.3% is achieved even with the oracle method (using true test labels to find the oracle square matrix $T$). Therefore, DART, which uses the same prediction modification scheme, can only achieve limited gains when there is only covariate shift. We report the experimental results of DART-applied TTA methods on balanced CIFAR-10C in Appendix E.5. We’d like to point that similar trends of limited gain for balanced datasets were also reported for ODS (Zhou et al., ICML’23), another attempt to alleviate test-time class distribution shifts.

(W2) could the proposed method be extended to source-free domain adaptation like SHOT (Liang et al., ICML 2020) (more epochs)? More results are welcome to verify the versatility of the proposed method

The major drawback of the entropy minimization-based methods is the error accumulation in training due to the utilization of inaccurate pseudo-labels. Therefore, improving pseudo-label accuracy before test-time training, as in DART, can significantly improve performance. We expect that our method can also be utilized as a plug-in method for source-free domain adaptation (SFDA) methods that use pseudo-labels for training like SHOT (Liang et al., ICML’20) and NRC (Yang et al., NeurIPS’21). For a fair comparison, we fine-tune BN layers in pre-trained models by Adam optimizer with learning rate 1e-3 for those SFDA methods. To test the scalability of DART to SFDA methods under test-time label distribution shift, we conduct experiments on CIFAR-10C-LT of $\rho=100$ in the offline manner (multiple epochs). SHOT and NRC utilize the prediction diversity loss ($\mathcal{L}_{\text{div}} = KL(\bar{p}||u)$, where $\bar{p}$ is the averaged pseudo label distribution and $u$ is a uniform label distribution) to prevent the adapted classifier from predicting all data into some classes. Since the test label distributions are imbalanced, the loss may cause prediction degradation. Thus, we report the experimental results in Table R1 for both cases when this loss is used and not used. While TENT/SHOT/NRC showed almost similar test accuracy of 85.83/85.67/85.97% on balanced CIFAR-10C (i.e., $\rho=1$), all the methods showed degraded performance on CIFAR-10C-LT of $\rho=100$. Especially, SHOT showed poor performance on CIFAR-10C-LT due to inaccurate self-supervised pseudo labels caused by the severe class imbalance of test data. However, we can observe that DART consistently improves the performance of adapted classifiers for all the SFDA methods.

R1. Comparison under offline setup, which adapts the pre-trained models by multiple epochs. We report average accuracy (%) on CIFAR-10C-LT of $\rho=100$.

(W3) concerning the mapping module $g_\phi$ is the network design (like two-layer MLP or the hidden dimensional) sensitive?

We conducted experiments to check the sensitivity of DART over the hidden dimension $d_h$ and number of layers of $g_\phi$, and the results are summarized on CIFAR-10C-LT of $\rho=100$ in Table R2. We can observe that DART is robust against the change in the $g_\phi$ structure.

R2. Sensitivity analysis about the network design of $g_\phi$.

(W4) how about the sensitivity of the batch size $B$ in the proposed method?

We conducted experiments to check the sensitivity of DART over $B$, the test batch size, and the results are summarized in Table R3. We can observe that DART is robust against the change in $B$.

R3. Sensitivity analysis about the test batch size B. We report average accuracy (%) on CIFAR-10C-LT of $\rho=100$.

Thank the authors for the responses. However, two major concerns remain, leading to a final score in the range of 5-6. (The main idea is interesting but the proposed method is not solid.)

• As acknowledged by the authors, the proposed method seems to fail in common covariate shifts with the existence of labels shift, heavily limiting its application in real-world TTA problems.

• As stated in the original review, the reviewer is curious about the performance under small batch sizes. However, only results under larger batch sizes are provided in the rebuttal.

Besides, the results of source-free domain adaptation are also not persuasive. Without the strong long-tail distribution shift, would the proposed method work? Since the label distribution in common domain adaptation (DA) datasets is also not balanced, I strongly suggest the authors validate the effectiveness of the proposed method on standard DA benchmarks like OfficeHome and DomainNet.

If the authors could provide evidence about the effectiveness of the proposed method beyond long-tail datasets, I would increase my score.

Thanks for your reply.

1. The Oracle is the best built-in variant of your method, which does not mean that the negligible gains are acceptable. It is disappointing to find even the Oracle variant of the proposed method cannot address covariate shifts.

2. Even for the ImageNet testbed, a high imbalance ratio is required in Table 3, $\alpha>1000$. OfficeHome is much not larger, each domain only has several thousand images, which makes it easy for a quick comparison under naturally imbalanced label shift. By the way, is a dataset consisting of only 7 classes hard?

3. For the sensitivity of the batch size, would you like to provide the results except for long-tail datasets, e.g., CIFAR10-C, CIFAR100-C, or Digits?

Generally speaking, the proposed method is limited to highly imbalanced datasets. It would be much more suitable if the authors studied the test time prior shift and compared the proposed method with counterparts [1-2] in that field.

[1]. Beyond invariance: Test-time label-shift adaptation for distributions with" spurious" correlations

[2]. Self-supervised aggregation of diverse experts for test-agnostic long-tailed recognition

Continued response to (W2)

In Appendix D, we also reported the experimental results showing that the traditional LLN method is not effective in TTA for refining the predictions. In particular, we considered HOC (Zhu et al., 2021), one of the recent LLN methods that is known to successfully estimate the noise transition matrix under LLN scenarios. We found that HOC failed to estimate the transition matrix for CIFAR-10C-LT with the label distribution shift of $\rho=100$. HOC estimates the transition matrix by using the empirical pseudo label distribution of nearest neighbors of each example. However, as observed in Figure 6, the nearest neighbors in the embedding space already have the same pseudo labels/predictions for the BN-adapted classifier, which makes it impossible to estimate a correct $T$ depending on the label distribution shift.

(Q1) It would be beneficial if the authors provided a more detailed explanation or key intuition behind the use of averaged pseudo labels as inputs in the distribution shift-aware module during intermediate time.

As explained in our response to (W2 of the reviewer), the distribution shift-aware module is trained to output a square matrix of the class dimension, only using the averaged pseudo labels, to fix the misclassification of the BN-adapted model originated from the class distribution shift. We designed the module to take the averaged pseudo label distribution as an input since this is the only available information at the test time. In the intermediate time, the module experiences several batches of diverse class distribution shifts and is optimized to generate a square matrix depending on each pseudo label distribution that can be multiplied to the outputs of BN-adapted classifier to minimize the cross-entropy loss between the true labels of training examples and the modified predictions. Through this process, the module learns the misclassification pattern of the BN-adapted classifier, originated from the label distribution shift, just from the observation of the average model predictions (pseudo labels). To verify this argument, in Appendix E.2, we reported the experimental results of the iterative version of DART, where the DART takes the iteratively refined pseudo labels as input. Specifically, for $i \in \mathbb{N}$, $T_i = g_\phi(\mathbb{E}\_{x \in D_{\text{test}}}[\text{softmax}(f_\theta(x) \Pi_{j=0}^{i-1}T_{j})])$, where $T_0$ is set to an identity matrix. Thus, $T_1$ is the output of our original DART. In Table 7 of the manuscript, we reported the pseudo label accuracy as the iteration increases. We can observe that (1) the accuracy of the refined prediction is maximized with $T_1$, and (2) the accuracy gradually decreases as the number of iterations increases. This indicates that our distribution-shift aware module indeed requires the average model predictions (pseudo labels) of BN-adapted classifier (not a more refined test label distribution) as it is trained to fix the misclassification pattern of the BN-adapted classifier, just using the noisy pseudo labels of the model itself.

(Q2) I wonder whether T_test is updated on every batch during test time, or if it is constructed just once across the entire test dataset.

$T_\text{test}$ is constructed just once across the entire test dataset using the averaged pseudo-label distribution of the whole test data as described in the last paragraph of Section 3. As an additional remark, we also considered DART for online TTA, assuming that each test data sample is encountered only once during test time, in Section 4.3. Since obtaining $T_\text{test}$ for the whole test data is unavailable for online TTA, we modified DART to take the averaged pseudo-label distribution of the first test batch to output $T_\text{test}$, which is then used throughout the entire test time. We summarized the experimental results of this variant of DART for online TTA in Table 4 (last row) and Appendix F. The results indicate that this online variant of DART performs similarly to the original DART but with a slight decrease in performance.

However, for the online label distribution shift setup on ImageNet-C, the class distributions within test batches are significantly different. Thus we computed the square matrix $T$ for each test batch as described in Appendix A.5.

Thank you for the constructive feedback. We hope the responses to the questions and weaknesses below resolve your concerns.

(W1) Using labeled data at intermediate time for training is a recently introduced protocol. However, compared to traditional TTA methods that cannot access labeled data, this approach may not be entirely fair. Test time adaptation, where the model learns directly during the testing phase, might be a more desirable direction.

We agree with the reviewer’s point that the intermediate-time training requires additional cost compared to the traditional TTA, but we’d like to emphasize that the additional cost in training our module for DART is very low and such a low additional cost can bring a significant performance gain. In particular, the cost in training DART is low since (1) it utilizes the training dataset, not a novel auxiliary dataset and (2) it requires a short runtime (about 2 hours for CIFAR-10C). With this low additional cost, our method exhibits consistent and significant performance gains when combined with the traditional TTA methods as shown in Table 2 of the manuscript. For example, it shows performance gains of 5.52% and 16.44% on CIFAR-10C-LT of $\rho=10$ and $100$, respectively.

(W2) The distribution-shift aware matrix for refinement has been frequently employed in handling label-noise datasets (Natarajan et al., 2013; Patrini et al., 2017; Zhu et al., 2021). Although the authors argue that TTA and these tasks differ, TTA essentially involves adding noise to the original data. Therefore, the nature of the problem between TTA and handling label noise is fundamentally similar. The authors need to further elucidate the methodological distinctions between their proposed approach and existing methods (Natarajan et al., 2013; Patrini et al., 2017; Zhu et al., 2021). Additional considerations should also be clearly addressed.

We’d like to first clarify that the nature of the two problems, Learning with Label Noise (LLN) and Test-Time Adaptation (TTA), are different since the causes of the class-wise confusion are different. In LLN, it is assumed that there exists a noise transition matrix $T$, which determines the label-flipping probability of a sample from one class to other classes. For LLN, two main strategies have been widely used in estimating $T$: 1) using anchor points, which are defined as the training examples that belong to a particular class almost surely, and 2) using the clusterability of nearest-neighbors of a training example belonging to the same true label class. LLN uses the empirical pseudo label distribution of the anchor points or nearest-neighbors to estimate $T$. For TTA, on the other hand, the misclassification occurs not based on a fixed label-flipping pattern, but from the combination of covariate shift and label distribution shift. To adjust the pre-trained model against the covariate shifts, most TTA methods apply the BN adaptation, which updates the Batch Norm statistics using the test batches. However, when there exists a label distribution shift in addition to the covariate shift, since the updated BN statistics follow the test label distribution, it induces bias in the classier (by pulling the decision boundary closer to the head classes and pushing the boundary farther from the tail classes as described in Appendix C “motivating toy example”). Thus, the resulting class-wise confusion pattern depends not only on the class-wise relationship in the embedding space but also on the classifier bias originated from the label distribution shift and the updated BN statistics. Such a classifier bias has not been a problem for LLN, where we don’t need to modify the BN statistics of the classifier at the test time.

Our proposed method, DART, focuses on this new class-wise confusion pattern and is built upon the idea that if the module experiences various batches with diverse class distributions before the test time, it can develop the ability to refine inaccurate predictions resulting from label distribution shifts. Based on this intuition, we train a distribution shift-aware module during the intermediate time, by exposing several batches with diverse class distributions using the training datasets. As described in Equation (1) of the manuscript, the module is trained using the labeled training dataset to output a square matrix of the class dimension for prediction refinement. In this process, the module takes the averaged pseudo-label distribution as an input to learn the class-wise confusion pattern of the BN-adapted classifier depending on the label distribution shift.

Thank you for your interest in our method. We hope that the responses to the questions resolve your concerns.

(W1) In Notation part, why do the test data have labels? In Section 2, the authors calculate the cross-entropy using the test labels $CE[\text{softmax}(f_\theta(x)T, y)]$. Is it a pseudo label?

In Section 2, we first presented an oracle attempt using the test labels to verify the effectiveness of multiplying a square matrix $T$ to the logits in refining the predictions generated by the BN-adapted classifier. Through this oracle attempt, we showed in Table 1 of the manuscript that the prediction refinement effectively enhances the performance of the BN-adapted classifier. However, during the test time, we don’t have access to test data labels. Thus, in this work, we proposed a novel test-time adaptation (TTA) method DART, which estimates the square matrix $T$ just with the pseudo labels of the test data using the distribution shift-aware module, as presented in Section 3.

(W2) The recent work SAR also considers the case with class distribution shifts. What is the advantage of the proposed method over SAR? I found that the authors prepared the imbalanced ImageNet-C dataset following SAR. However, I failed to find the empirical comparisons between the proposed method and SAR.

DART can be integrated as a plug-in method with any TTA methods that rely on pseudo-labels obtained from the classifiers, including SAR, since DART focuses on effectively modifying the inaccurate predictions/pseudo-labels caused by test-time label distribution shifts. Thus, DART can also be used with SAR, and we summarize the experimental results on CIFAR-10C-LT in Table R4 and ImageNet-C-imbalance in Table R5. We can observe that the performances of SAR are worse/better than those of TENT on CIFAR-10C-LT/ImageNet-C-imbalance, respectively. However, DART consistently improves the performance of the SAR in a similar way as it improved the performances of other TTA methods, since DART improves the accuracy of the initial pseudo-labels used for SAR.

R4. Average accuracy (%) on CIFAR-10C-LT

R5. Average accuracy (%) on ImageNet-C-imbalance

(W3) In my understanding, the authors seek to train a distribution shift-aware module to generate a matrix for prediction refining. In this sense, the data to train such a module should be class-imbalanced. However, as mentioned in the paper, the dataset D_int seems to be class-balanced.

Our method exposes various batches of diverse class distributions to the distribution shift-aware module by sampling the balanced $D_\text{int}$ using the Dirichlet sampling as illustrated in Figure 3 of the manuscript. We make $D_\text{int}$ class-balanced to avoid unintentional bias when sampling through the Dirichlet distribution. For example, if the intermediate dataset of two classes (1,2) has class distribution [0.9, 0.1], the sampled batches from the dataset are likely to have many more samples from class 1 even when Dirichlet sampling is applied.

(W4) What is the computational cost to train a distribution shift-aware module in the “intermediate time”? Does it take a long time?

As reported in Appendix A.6, it takes only about 2 hours to train the distribution shift-aware module for CIFAR-10C. Since the pre-trained classifier weights are frozen and the distribution shift-aware module has a simple structure (2-layer MLP), training the module takes a much shorter time compared to the pre-training step. Additionally, the computational overhead is not significant since we train the module only once for each pre-trained classifier regardless of the test corruption types or label-distribution shifts as described in Section 3 “training of $g_\phi$”.

(W5) In the case without class distribution shifts, what is the performance of the proposed method compared with existing methods? Better or worse?

Our method refines the predictions of the classifier by multiplying a square matrix $T$ to the logits, and this process adjusts the incorrect predictions mainly originated from the class distribution shift. Thus, without the class distribution shift, multiplying $T$ does not have a big impact in adjusting the classifier predictions. For the balanced CIFAR-10C ($\rho=1$), where there is no label distribution shift in test time, we report the experimental results of DART in Appendix E.5. We observe that DART-applied TTA methods show comparable or slightly worse performance than the original TTA methods. This is also attributed to the limited gain even with the oracle method (using true test labels to find the oracle square matrix $T$ that will be multiplied to the logits) as shown in Table 1 ($\rho=1$). Specifically, the oracle achieved only a marginal performance gain of 0.3% on average even with the labels of test data on balanced CIFAR-10C. Therefore, DART, which uses the same prediction modification scheme, can only achieve limited gains even when generating square matrices similar to the one with the oracle. We’d like to point out that similar trends of limited gain for balanced datasets were also reported for ODS (Zhou et al., ICML’23), another attempt to alleviate test-time class distribution shifts.