Test-Time Adaptation by Causal Trimming

Thank you for the extensive feedback and for recognizing our empirical analysis. We discuss the weakness and questions raised in the review below.

Weakness

W1. "no TTA'' baseline and training time augmentation

The "no TTA'' baseline of BirdCalls, Camelyon17, and CivilComments were not trained with the augmentations used by TACT to identify and reduce non-causal features. Following the instructions of the data publisher, no data augmentation was used when training the baseline with empirical risk minimization. To understand the effectiveness of training time augmentation, we trained the models with the data augmentations used by TACT. The result is shown in the table below. Training time augmentation is not always sufficient, as it achieves worse performance than TACT on BirdCalls and Camelyon17. We further show that, when training time augmentation is performed, TACT can still improve the performance. These results suggest TACT's usefulness when tackling distribution shift.

For ImageNet-R and ImageNet-V2, the "no TTA'' baseline provided by torchvision was trained with AutoAugment using the ImageNet policy. The table below shows the performance of TACT on ImageNet-R and ImageNet-V2 with the same augmentation. TACT improves the performance.

W2. Performance selected by validation data

For Birdcalls, Camelyon, and CivilComments, we use the validation domain for model selection. For ImageNet-R and ImageNet-V2, we use the validation set of ImageNet for model selection. The performance is shown in the table below. TACT performs the best on all datasets except for Birdcalls. For Birdcalls, TACT performs the second best.

W3. Computation cost

We compared the computation cost of TACT with other backpropagation-free methods on Birdcalls. The table below shows that TACT requires more time and memory than other backpropagation-free methods. However, this increase in computational cost leads to a significant improvement in performance shown in Table 1 of the main paper, which we believe justifies the trade-off. In many practical applications, the improved performance is critical and outweighs the added cost. Further optimization techniques on efficient eigendecomposition for PCA could be applied to reduce the overhead if needed, which we leave as future work.

Questions

Q1. ''no TTA" baseline

Please refer to W1.

Q2. Performance selected by validation data

Please refer to W2.

Q3. Removing components based on a threshold

We consider using the variance that the top PC accounts for as a threshold ($\tau$) to decide whether causal trimming is conducted or not. When the augmentation only changes non-causal features and causal features remain unchanged, datapoints that are invariant to augmentations should have smaller variance of the top PCs. Thus, if the variance is smaller than a threshold, causal trimmings will not be conducted on the data. As the range of variance is not known and it could change significantly, setting a numerical threshold might not be feasible. We consider normalized variance, where we divide the variance of top PCs by the sum of variances of all PCs.

The table below shows the performance. Current experiments on Birdcall with threhsold does not improve the performance. Future work can study this issue in more details.

Dear Reviewer DW91,

Thank you so much for the feedback you provided. We hope that our reply addresses your concerns. Please let us know if you have any further concerns or suggestions—we're more than happy to discuss them.

Thank you for the comprehensive feedback and for praising the proposed method as well as our empirical and theoretical analysis. A detailed discussion of the concerns and questions is presented below.

Weakness

Quality-W1. Reliance on augmentation and generalizability

We agree that augmentations need to be carefully selected to achieve good performance. In practice, the augmentation can be treated as a hyperparameter to search over. Our basic assumption is that knowledge about the prediction target is known, based on which augmentations that are semantically invariant with respect to the prediction target can be designed. For example, for the commonly studied image classification task, we recommend searching over general image augmentations, such as AutoAugment and RandomAugment. These augmentations preserve the key causal features, particularly the shape information of objects [1]. Our experiments in the main paper of using them with ImageNet-R and ImageNet-V2 show state-of-the-art performance. Our experiment of using AutoAugment and RandomAugment with Birdcalls and Camelyon shows that they can also offer performance gain compared to no adaptation, suggesting the potential generalization ability of TACT.

Quality-W2. Non-causal features dominate variance directions and entangled representation space

We select data augmentation strategies to heuristically maximize variations along non-causal directions and minimize variations along causal directions, which leads to improved disentanglement after applying PCA. Therefore, it is critical to select the augmentations carefully.

Theoretically, TACT does not rely on perfect disentanglement. We show that when causal features are partially trimmed, if sufficient causal information is preserved and the learned decision boundary is directionally aligned with the causal decision boundary, TACT would perform well. Empirically, in Section 6, we show that even approximate separation of causal and non-causal features achieved by PCA was sufficient to yield gains in performance and robustness across diverse datasets.

Clarity-W1. Tuning guidance of number of trimmed PCs

When the oracle is not accessible, m in general can be set to 1.

Clarity-W2. Failure cases

When causal features and spurious features are not orthogonal, removing spurious features will affect the remaining causal features, which likely also degrade performance. In the extreme, feature removal will not work if spurious features are perfectly correlated with causal features in the training data. This also corresponds to our theoretical analysis where sufficient causal features need to be preserved to make correct predictions.

Questions

Q1. Robustness when the assumption is only approximately satisfied

Theoretically, Propositions 2 and 3 suggest that when causal features vary, resulting in the trimming of causal features, sufficient causal information needs to be preserved for TACT to perform well.

Empirically, in our experiment with ImageNet-R and ImageNet-V2, we only tried to maintain the key causal features, which are object structure and shape. Other causal features that could be helpful in inferring objects, such as color when inferring strawberries, are altered. This experiment shows that TACT is still able to improve performance when the assumption is only approximately satisfied.

Q2.1. Performance on larger models when representation entanglement is common

The table below shows the performance with larger backbones, including ViT-B/16 for images and BERT for texts. TACT achieves the best performance compared to other backpropagation-free methods on all datasets except ImageNet-R. TACT achieves the second best on ImageNet-R. The result suggests that TACT can be robust to changes in the representation space.

The models for Birdcalls and Camelyon are trained under the same setting as that for ViT-B/32 used in the paper. We follow the guidance of CivilComments' publisher to train BERT. ViT-B/16 backbone for ImageNet-R and ImageNet-V2 is published by torchvision.

Q2.2. Alternative direction finding methods

Although kernelPCA and ICA overcome the limitations of PCA, they incur additional requirements.

• kernelPCA transforms the representation space via a kernel, which needs to be selected such that causal and non-causal information are retained and better separated. Due to this complexity, we encourage further exploration on this in future works.
• ICA looks for statistically independent components and assumes each component follows a non-Gaussian distribution. Causal and non-causal features might not follow the non-Gaussian distribution assumption under augmentations that vary non-causal features.

We experiment with using ICA with TACT on Birdcalls. We rank the independent components by the variance of the scalars of features on the components. We remove the top independent components that have maximum variance. The experiment with ICA is searched over the same number of augmentations and number of components. The table below shows the result. ICA performs inferior to PCA, but better than no adaptation.

Q3. Selection of the number of removed PCs

Please refer to Clarity-W1.

[1] Geirhos et al., ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness, ICLR 2019.

Dear Reviewer QQc2,

We thank you sincerely for the feedback. We hope that our reply resolves your questions. If there are any further questions or suggestions, please do not hesitate to let us know. We would be glad to engage in any additional discussion or clarification you require.

Thank you for the detailed feedback and for acknowledging that TACT is a method that addresses the root cause of distribution shift backed by theoretical analysis. The weaknesses and questions raised in the review are discussed below.

Weakness

W1.1. Augmentation sensitivity analysis

The table below shows the result. The choice of augmentation would influence performance, but we see that general augmentations such as AutoAugment and RandomAugment could provide improvements over no adaptation across datasets. To our surprise, Birdcalls even performs better with RandomAugment than the existing color jitter.

W1.2. Automated strategy to select augmentation

In practice, the augmentation can be treated as a hyperparameter to search over. The data collection process that raises variation and features that affect the prediction target should be analyzed to propose a set of augmentations that are semantically invariant with respect to the prediction target, yet introduce variability in other, non-causal aspects. For example, for the commonly studied image classification task, we recommend searching over general image augmentations, such as AutoAugment and RandomAugment, that preserve the key causal features, particularly the shape information of objects [1]. The most effective way to select the augmentation is to test on a small subset of labeled test data.

W2.1. Orthogonality and co-varying of causal and spurious features

TACT does not rely on perfect orthogonality or disentanglement. The theoretical analysis in section 5 does not rely on the orthogonality assumption. We show that even when causal and non-causal features are not orthogonal, and causal features are partially trimmed, TACT performs well as long as top PCs alone cause the prediction error, or sufficient causal information is preserved after trimming, and the learned decision boundary is directionally aligned with the causal decision boundary. Empirically, in Section 6, we show that even approximate separation achieved by PCA was sufficient to yield gains in performance and robustness across diverse datasets.

While it is true that causal and spurious features can co-vary in the training data, it is important to note that when we apply augmentations at test time, we are trying to minimize the variation of causal features, so that the influence of non-causal features can be isolated.

W2.2. Alternative direction finding methods

We experiment with using ICA with TACT on Birdcalls. We rank the independent components by the variance of the scalars of features on the components. We remove the top independent components that have maximum variance. The experiment with ICA is searched over the same number of augmentations and number of components. The table below shows the result. ICA performs inferior to PCA, but better than no adaptation. We agree that ICA can overcome the orthogonality assumption. However, ICA only looks for statistically independent components and assumes each component follows a non-Gaussian distribution. Thus, ICA might not be better at identifying causal and non-causal features, which might not follow the non-Gaussian distribution assumption under augmentations that vary non-causal features.

Questions

Q1. Number of PCs to project out

We explain the projection in section 4.2 (Causal Trimming to Reduce Non-Causal Feature). The top-m PCs are removed. We maintain the same m for a dataset across batches. $m$ is a hyperparameter to be selected. We empirically observe that setting $m=1$, i.e., removing the top PC is frequently adequate. For datasets like ImageNet-R with more complex or layered distribution shifts, removing more PCs could further boost the performance.

Q2. Statistical significance test

To clarify, all experiments are repeated three times, and we report the mean and standard deviation in Table 1. We use the t-test to compute the statistical significance between the results of TACT-adapt and the strongest baselines. The p-value shows how likely the results of the two methods come from the same distribution. The table below shows the result. Camelyon, CivilComments, and ImageNet-V2 have a p-value smaller than 0.15.

We acknowledge that the improvement on Birdcalls and ImageNet-R does not seem statistically significant, which is likely due to the small number of runs. Therefore, we performed two additional runs to compute the result. The table below shows the updated result. P-values for the two datasets are improved. We plan to conduct more runs and update the results in the revision.

Q3. Hyperparameter and performance

Figure 4 shows the performance over various $(n, m)$ combinations. The colored lines denote the number of augmentations, highlighting the role of $n$, while the x-axis represents $m$, the number of removed PCs. Performance values (y-axis) are determined by the intersection of each line with the respective m value. Empirically, we find that values of $n \in$ {128, 256, 512} provide sufficient performance. For the number of removed principal components $m$, we observe that removing the top principal component, which typically captures the dominant non-causal variation, is frequently adequate.

[1] Geirhos et al., ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness, ICLR 2019.

Dear Reviewer 5gmS,

We are truly thankful for the feedback you provided. We hope that our reply clarifies your questions. If there are any additional concerns or questions you may have, please do not hesitate to let us know. We are more than willing to discuss any further suggestions or issues you might have.

Dear Reviewer 5gmS,

Thank you for raising the question regarding the scalability and efficiency of our method.

We use torch's torch.linalg.svd() to perform PCA as discussed in our appendix. The function is able to find PCs of a batch of high-dimensional data. The internal implementation of torch.linalg.svd() relies on highly optimized linear algebra backends. As a result, the practical efficiency tends to scale well on modern hardware, particularly for high-dimensional data in batches commonly considered in nowadays machine learning applications. In our experiments, we did not encounter scalability issues. We experimented with the representation of dimension 768 and batch size up to 128. The largest dataset we tested is of size 133,782.

As for latency, we compared the latency of TACT with other backpropagation-free methods on Birdcalls. The table below shows that TACT requires more time than other backpropagation-free methods. The increase in latency leads to a significant improvement in performance shown in Table 1 of the main paper, which we believe justifies the trade-off.

Further optimization techniques on efficient eigendecomposition for PCA could be applied to reduce the overhead if needed, which we leave as future work.

Thank you for the response. I think method’s substantially higher latency compared to other approaches may limit its applicability in low-latency settings. While further PCA optimizations are suggested, they are left as future work. I will update my score. But given these practical constraints, I update the score from 3 to 4.

Thank you for the thoughtful feedback and for recognizing our novelty and theoretical analysis. The limitations raised in the review were acknowledged in section 7 of the main paper. We discuss the two limitations in more detail below.

1. Domain knowledge to design augmentation

We agree that requiring prior knowledge of the data is a limitation of the proposed method. Under the setting of TTA, where only a trained model and test data are given, in the absence of any task-specific or domain knowledge, it is inherently difficult—if not impossible—to distinguish causal from non-causal features [1, 2].

Data augmentation is one approach to identifying non-causal features. In practice, when insufficient domain knowledge is given and the optimal augmentation strategies are unknown, the augmentations can be treated as a hyperparameter to search over. Specifically, our basic assumption is that, given a dataset and a task, we at least have the knowledge about the prediction target. Based on this, we can design augmentations that are semantically invariant with respect to the prediction target. For example, for the commonly studied image classification task, we recommend searching over general image augmentations, such as AutoAugment and RandomAugment. These augmentations preserve the key causal features, particularly the shape information of objects [3]. Our experiments in the main paper of using them for ImageNet-R and ImageNet-V2 show state-of-the-art performance. Birdcall's audio clip is converted to Mel spectrogram, which is treated as image as well. Our experiment of using AutoAugment and RandomAugment with Birdcalls and Camelyon show that they can also offer performance gain compared to no adaptation.

The paper wants to highlight that reducing non-causal features by adjusting the learned representations is helpful for TTA. Alternative to data augmentation, if other information, such as a small set of labeled data from multiple distributions, is given, non-causal features can also be identified by comparing the representations of data with the same labels.

2. Linearity and orthogonality of the principal directions

We use a simplified representation space that satisfies the linearity and orthogonality assumptions of causal and non-causal features to approximate the learned representation space, so that a practical method could be developed. Nonetheless, the theoretical analysis in section 5 does not rely on the assumption. We show that even when causal and non-causal features are not orthogonal, and causal features are partially trimmed, TACT performs well as long as top PCs alone cause the prediction error, or sufficient causal information is preserved after trimming, and the learned decision boundary is directionally aligned with the causal decision boundary. Empirically, the models used for TTA in section 6 do not have their representation space explicitly enforced to satisfy the assumption.

[1] Wang et al., Sound and Complete Causal Identification with Latent Variables Given Local Background Knowledge, NeurIPS 2022.

[2] Schölkopf et al., Toward causal representation learning, Proceedings of the IEEE 2021.

[3] Geirhos et al., ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness, ICLR 2019.

Dear Reviewer ptsx,

We sincerely appreciate the feedback you provided. We hope that our discussion on the limitations of TACT addresses your concerns. We're happy to clarify anything further, so please don't hesitate to share any remaining questions or suggestions you may have.