Paper: Un-Mixing Test-Time Normalization Statistics: Combatting Label Temporal Correlation (ICLR-2024)

Reviewer comment (Experiments):
Why the hyper-parameter \alpha is set to be 0.5 in Equation (5) for implementation while initializing UNMIX-TNS components?

Author response:
Thank you for your query regarding the choice of setting the hyperparameter $\alpha$. As mentioned in Appendix A.2 of our manuscript, we analyze the source variance within the context of UnMix-TNS. Our analysis indicates that the source variance is contributed by (1) the average variance of individual components, and (2) the variance of the means of individual components. By setting $\alpha=0.5$, we ensure that these two contributions to the source variance are balanced and have equal weight. This balance is crucial for the effective functioning of UnMix-TNS.

Reviewer comment (Reproducibility):
In Section 2.2.3, the hyper-parameter \lambda is not explained clearly.

Author response:
Thank you for highlighting the need for a clearer explanation regarding the hyperparameter $\lambda$ in Section 2.2.3 of our paper. Equations (12) and (13) are derived based on the principles of momentum batch normalization (MBN), and the hyperparameter $\lambda$ is set according to the MBN guidelines proposed by Yong et al. (2020). The primary objective of MBN is to ensure a consistent noise level introduced by BN across varying batch sizes. This consistency is crucial to maintain equivalent noise levels in both small and large batch scenarios. The tuning of $\lambda$ is integral to achieving this standardization based on the specific batch size employed. For a more comprehensive understanding of $\lambda$ and its significance in our method, we have included an in-depth explanation in the newly added Appendix A.4 of the manuscript.

Reviewer comment (Experiments):
While Tables 1 and 2 include 11 methods for comparison, the experiments on the video dataset (Table 3) only include seven methods.

Author response:
Thank you for your comment regarding the number of methods included in the comparisons for the video dataset (Table 3). We aimed for a fair and relevant evaluation by aligning our experimental setup with that of LAME (Boudiaf et al., 2022), selecting top-performing baselines for comparison. The methods listed in Tables 1 and 2, which are not included in Table 3, either necessitate extensive reimplementation for application to video datasets or are not directly applicable to such datasets. This decision was made to ensure the most meaningful and practical comparison within the constraints of our experimental setup.

Reviewer comment (Reproducibility):
Does the UNMIX-TNS update the same parameters with other TTA methods? And have the affine parameters in BN layers been changed within the optimization? Does the method require only single forward inference or multiple forward passes?

Author response:
In response to the reviewer's queries about the UnMix-TNS method, as a standalone method, UnMix-TNS is designed to update only its internal statistics and requires only a single forward pass. However, when integrated with other TTA methods, UnMix-TNS solely replaces the Batch Normalization (BN) layers and behaves like the replaced BN in the same manner implemented in those TTA methods, including parameter updates. For instance,  in a TTA method like TENT, the affine parameters of the BN layers are updated, and UnMix-TNS will follow this same update protocol. Conversely, in methods like LAME, the focus is on updating only the internal statistics of normalization layers. In summary, UnMix-TNS can either operate independently with a single forward pass or align with the update mechanisms of other TTA methods it is paired with.

Reviewer comment (Theory):
The authors mentioned that “only components that closely align with the instance-wise statistics undergo updates”, but equations (12) and (13) update all K BN statistics components with different assignment probabilities.

Author response:
We acknowledge that our initial phrasing may have been misleading. In practice, we update all K BN statistics components, each with varying assignment probabilities. We have updated the manuscript to clarify this point more accurately.

Reviewer comment (Theory):
**Lack of Theoretical Insight for Temporal Correlation Handling:** The abstract and introduction do not provide clear theoretical insights or explanations on how UnMix-TNS effectively deals with test data streams having temporal correlation. A deeper understanding of the underlying principles is crucial for assessing the robustness and reliability of the method.

Author response:
We appreciate your thoughtful review and valuable suggestion to provide a deeper theoretical understanding of how our model accounts for temporal correlation. In response, we have expanded our theoretical analysis. In our revised manuscript, we have added a dedicated section (A.3 in the Appendix) to address this topic comprehensively. We also compare UnMix-TNS with the standard test-time normalization method (TBN) in section A.3. Furthermore, we have highlighted why UnMix-TNS can help mitigate the bias introduced in the estimation of feature normalization statistics, which arises due to the time-correlated feature distribution.

Reviewer comment (Experiments):
**Potential for Bias in Instance Selection:** The effectiveness of UnMix-TNS might be influenced by the selection of instances from incoming test batches for updating statistics. If the selection is not well-designed, it could introduce bias, affecting the generalization of the model.

Author response:
We would like to clarify any confusion in the submitted paper regarding the potential for bias in instance selection. For a given batch of test feature-map instances, represented as $\mathbf{z}\in \mathbb{R}^{B\times C\times H \times W}$ (where $B$ is the number of instances, $C$ the number of channels, $H, W$ the width and height of an instance feature-map), **we do not perform any selection of the test features for updating the components statistics**. Instead, every individual feature map $z[b]$ $\in \mathbb{R}^{C\times H \times W}$ is soft-assigned to the $K$ components of UnMix-TNS. This assignment is based on the cosine similarity between the mean of each instance-wise feature map and the centers of the components (Equation 7). Based on this similarity given by $p_{b,k}$, we update the parameters $\mu_k^t$, $\sigma_k^t$ of all $K$ components (Equations 12 \& 13 in the paper). We acknowledge the importance of this clarification.

Reviewer comment (Theory):
Could the authors provide more detailed theoretical insights or foundational principles on how UnMix-TNS adapts to test data streams with temporal correlation? Specifically, how does the method theoretically account for and mitigate issues that arise due to temporal dependencies in the data?

Author response:
We appreciate your thoughtful review and valuable suggestion to provide a deeper theoretical understanding of how our model accounts for temporal correlation. In response, we have expanded our theoretical analysis. In our revised manuscript, we have added a dedicated section (A.3 in the Appendix) to address this topic comprehensively. We also compare UnMix-TNS with the standard test-time normalization method (TBN) in section A.3. Furthermore, we have highlighted why UnMix-TNS can help mitigate the bias introduced in the estimation of feature normalization statistics, which arises due to the time-correlated feature distribution.

Reviewer comment (Theory):
What is the theoretical rationale behind the 'unmixing' approach in handling temporally correlated data? How does this strategy compare with traditional methods in terms of theoretical robustness against temporal variations in data streams?

Author response:
We appreciate your thoughtful review and valuable suggestion to provide a deeper theoretical understanding of how our model accounts for temporal correlation. In response, we have expanded our theoretical analysis. In our revised manuscript, we have added a dedicated section (A.3 in the Appendix) to address this topic comprehensively. We also compare UnMix-TNS with the standard test-time normalization method (TBN) in section A.3. Furthermore, we have highlighted why UnMix-TNS can help mitigate the bias introduced in the estimation of feature normalization statistics, which arises due to the time-correlated feature distribution.

Reviewer comment (Reproducibility):
Could the authors elaborate on the criteria or algorithm used to select instances from incoming test batches for updating the UnMix-TNS statistics? How do the authors ensure that this selection process does not introduce bias, which might affect the model’s generalization capability?

Author response:
We would like to clarify any confusion in the submitted paper regarding the potential for bias in instance selection. For a given batch of test feature-map instances, represented as $\mathbf{z}\in \mathbb{R}^{B\times C\times H \times W}$ (where $B$ is the number of instances, $C$ the number of channels, $H, W$ the width and height of an instance feature-map), **we do not perform any selection of the test features for updating the components statistics**. Instead, every individual feature map $z[b]$ $\in \mathbb{R}^{C\times H \times W}$ is soft-assigned to the $K$ components of UnMix-TNS. This assignment is based on the cosine similarity between the mean of each instance-wise feature map and the centers of the components (Equation 7). Based on this similarity given by $p_{b,k}$, we update the parameters $\mu_k^t$, $\sigma_k^t$ of all $K$ components (Equations 12 \& 13 in the paper). We acknowledge the importance of this clarification.

Reviewer comment (Reproducibility):
In situations where the test batches contain highly diverse or outlier instances, how does UnMix-TNS maintain the integrity of the normalization statistics? Is there a mechanism in place to detect and appropriately handle such anomalies?

Author response:
We appreciate the reviewer's concern about handling highly diverse or outlier instances in test batches. We want to emphasize that we adopt the same setting as other test-time normalization methods for a fair comparison by assuming that temporally correlated test batches are free from outliers, such as instances from novel or unknown categories. This is an important consideration, as the presence of outliers may result in inaccurate normalization statistics estimations in UnMix-TNS, an issue that is also applicable to other methods. 

Fortunately, we are already planning to extend our work to include detecting and excluding anomalies during the test-time adaptation process. The envisioned approach involves leveraging an anomaly detection score, calculated based on the distance metric of test instances to the nearest component in the UnMix-TNS and will utilize an optimal threshold to identify outliers effectively. By identifying and pruning outlier instances using this strategy, we aim to prevent the adverse impact these instances might have on adapting the model. This will significantly improve the robustness and accuracy of UnMix-TNS in dealing with diverse test batches. We have outlined this planned improvement as a key focus for our future work in the revised version of our manuscript (conclusion section).

Reviewer comment (Experiments):
The paper doesn't discuss how to choose $K$ and $\lambda$, and the additional experiments in Appendix C only study the sensitivity to $K$.

Author response:
Thank you for pointing out the need for a clearer explanation of the selection of $K$ and $\lambda$ in our paper. Regarding $K$, its choice is crucial, as emphasized in our future work. A single UnMix-TNS component can efficiently represent several similar classes (e.g., car, truck, bus), allowing for an accurate approximation of the test distribution with a lower $K$ value. This approach is particularly beneficial for datasets with distinct class similarities. However, in datasets with a large number of classes like CIFAR100-C, DomainNet-126, or ImageNet-C, aligning $K$ with the number of classes may lead to inefficiencies and an increased memory footprint. For instance, setting $K$ to 100 in CIFAR100-C tests led to reduced performance. The optimal $K$ varies with the testing scenario; larger $K$ values are necessary for mixed domains with diverse styles, while smaller $K$ values are sufficient in single or continual domains due to reduced variability. In summary, choosing $K$ in UnMix-TNS is about balancing class/domain complexity with model efficiency, and we plan to refine these guidelines further in our future research.

Concerning $\lambda$: Equations (12) and (13) are derived based on the principles of momentum batch normalization (MBN), and the hyper-parameter $\lambda$ is set according to the MBN guidelines proposed by Yong et al. (2020). MBN aims to standardize the noise level introduced by Batch Normalization (BN) across different batch sizes, ensuring comparable noise levels between small and large batches. The tuning of $\lambda$ is integral to achieving this standardization based on the specific batch size employed. For a more in-depth explanation, we have added new information in Appendix A.4 of our paper.

Reviewer comment (Reproducibility):
Algorithm 1 in Appendix B.2 is both helpful and confusing: helpful because it describes the algorithm concisely; but confusing because the shapes don't add up. For example, `instance_mean` should have shape `(B, 1, C)`, otherwise you cannot get `hat_mean` of size `(B, K, C)` through broadcasting. It would be better to include a runnable Python snippet in the Appendix, even if that's slightly more verbose.

Author response:
Thank you for your valuable feedback regarding Algorithm 1 in Appendix B.2. We understand your concerns about the clarity of the shapes used in the algorithm and apologize for any confusion caused. To address this, we have revised Algorithm 1 to ensure the correct alignment of tensor shapes.

In the updated algorithm, we have introduced the **unsqueeze** operation to adjust the dimensions of various tensors appropriately. The **instance\_mean** is now unsqueezed along dimension 1, altering its shape to (B, 1, C). This adjustment allows for the correct broadcasting when calculating the **hat\_mean** of size (B, K, C). Additionally, we have applied the unsqueeze operation to self.components\_means along dimension 0, resulting in a shape of (1, K, C). Furthermore, we have unsqueezed $p$ along the last dimension, changing its shape to (B, K, 1). By incorporating all the above steps, we can now correctly combine these tensors, ensuring the resultant **hat\_mean** has the intended size of (B, K, C). We hope that these clarifications and the revised algorithm will make the implementation more straightforward and eliminate any confusion.

Reviewer comment (Reproducibility):
In Appendix B.2, it's unclear what `torch.var_mean(x, dim=[2, 3])` means. Does `x` have dimension `(batch, channel, height, weight)`?

Author response:
Thank you for your query. We would like to clarify that in this context, the variable $x$ indeed represents a batch of intermediate features within a Convolutional Neural Network (CNN). As is typical with CNNs, the tensor $x$ possesses four dimensions: batch size, number of channels, height, and width. Therefore, when we call torch.var\_mean(x, dim=[2, 3]), the function calculates the variance and mean across the height and width dimensions of the tensor $x$. This operation effectively computes the instance-wise statistics for each feature in the batch, considering the spatial dimensions (height and width) of the CNN's intermediate feature maps.

Reviewer comment (Theory):
From my understanding, UnMix-TNS will fail when test samples gradually shift from one cluster to another, in which case the active $\mu_k$ would follow the sample, causing two clusters to eventually overlap. How would you prevent that?

Author response:
Thank you for your insightful comment regarding the potential issue of UnMix-TNS when test samples gradually shift between clusters, leading to an overlap of clusters. 

We would like to highlight that such a scenario can happen if the semantics of an input image change gradually over time, causing a gradual drift in the feature representation. To address this, UnMix-TNS is designed to update the statistics of the active component slowly, using an exponential moving average. This method of updating helps in mitigating the risk of rapid shifts in the active  $\mu_k$. Moreover, by choosing a lower value for the momentum hyperparameter $\lambda$,  we can further slow down the movement of the active $\mu_k$, making it less susceptible to quickly following the distribution drift. This careful calibration ensures that the active $\mu_k$	responds to changes in the data distribution at a controlled pace, reducing the likelihood of overlapping clusters due to rapid shifts in feature representation.

Reviewer comment (Experiments):
Not really a question, but an interesting experiment is to set $K$ to the number of corruption types, and see if $p_{b,k}^t$ can recover the corruption type for each sample correctly. Maybe each cluster will correspond to the label instead of corruption. Who knows?

Author response:
Thank you for your intriguing suggestion regarding the experiment on setting $K$ to the number of corruption types. Furthermore, conducting an experiment to test this hypothesis could be very enlightening, potentially revealing how different layers of a network interpret and represent various aspects of the data, from style and domain at the initial layers to semantic content towards the end.

In line with your suggestion, we have conducted additional experiments on CIFAR10-C under mixed-domain scenarios. We performed two sets of experiments: firstly, setting $K=15$, which corresponds to the number of corruption types, and secondly, setting $K=150$, representing the number of corruptions multiplied by the number of classes. Our findings from these experiments were quite revealing. In the first scenario, we observed an average error rate of 63.5\% across the 15 corruption types. However, in the second scenario, where $K$ was significantly higher, the error rate decreased to 40.9\%. These results underscore the importance of choosing an appropriate value for $K$ depending on the test scenario. In mixed domain scenarios, with a wide range of styles in the test distribution, a larger $K$ is essential for effectively capturing this diversity. This is in contrast to single domain scenarios, where the domain is consistent, or continual domain scenarios, with sequential domain shifts.

In summary, our experiments reinforce the notion that selecting $K$ in UnMix-TNS is a balancing act between the complexity of the classes and domains, and the practical considerations of model efficiency. Particularly in datasets with a large number of classes, setting $K$ equal to the number of classes can lead to inefficiencies, as demonstrated by our experimental results. We believe these findings will add a valuable perspective to our understanding of how $K$ influences model performance in various scenarios.

Reviewer comment (Theory):
Multiple target domains: One of the claimed contribution is that UnMix-TNS has robustness when tested on continual domain and mixed domain. However, I am very confused which part of UniMix-TNS is designed for and beneficial to these multi-domain settings, especially the mixed domain setting. Considering that there is a new batch containing images from different domains, although different images may correspond to different UnMix-TNS components, they are finally normalized with the same mean and variance according to Eq (11). It seems like UnMix-TNS still treat multiple target domains as one single target domain.

Author response:
Thank you for your insightful comment. We agree that UnMix-TNS does not incorporate a specific module tailored for mixed domain settings. However, we want to shed light on the fact that UnMix-TNS shows robustness against temporal correlations (different values of Dirichlet coefficient $\delta$) of the test images with respect to their inaccessible labels compared to other test time normalization schemes.

The core functionality of UnMix-TNS lies in its ability to retain current and past statistics of online test features through its $K$ components, functioning as a time-decaying memory. In cases where test features are label-correlated, UnMix-TNS effectively simulates an i.i.d. scenario. It accomplishes this by estimating the current mean and standard deviations from its $K$ components, each representing different groups of test features encountered previously.

Additionally, the gradual update of the $K$ components over time enables UnMix-TNS to adeptly handle continual domain settings, where domain shifts occur sequentially rather than randomly. This capability was demonstrated in our experiments, showing UnMix-TNS's effectiveness not only in single domain settings but also in continual domain scenarios. Therefore, while UnMix-TNS may not specifically target mixed domain settings, its underlying design principles and operational mechanics provide flexibility and adaptability across varying domain configurations, including mixed domain. Furthermore, as supported by our experiments, our method's adaptability and performance in mixed-domain scenarios can be further enhanced by adjusting the momentum hyperparameter $\lambda$ and appropriately increasing the value of $K$ to encompass the style diversity in the mixed domain.

 Finally, our claim regarding the mixed domain scenario has been revised in the updated manuscript.

Reviewer comment (Experiments):
Unclear experimental setting and unsatisfactory performance: The author claimed that they follow the protocol outlined in (Lim et al., 2023). However, the performance for most of the baselines are significantly worse than the results in (Lim et al., 2023). Also, the proposed UnMix-TNS fails to outperform “Source” in mixed domains. I believe the comparison to baselines only makes sense if the test-time normalization is beneficial. (I presume that “Source” means no adaptation. Please correct me if I am wrong)

Author response:
Thank you for your comment and for highlighting these concerns. Firstly, we apologize for any confusion caused by the experimental protocol reference. We followed the protocol set out in Marsden et al., 2023, not Lim et al., 2023. This error has been corrected in our revised manuscript.

Regarding the performance of baselines, it's important to note that Lim et al., 2023, primarily reported baseline performances in an i.i.d. setting. They only report the performance of their method for class imbalance for the single domain setting. Our approach with UnMix-TNS, while occasionally not surpassing the source model, does show consistent improvements over other test-time normalization schemes, particularly where these schemes may not have yielded significant enhancements. Additionally, it's important to recognize that our results are consistent with Lim et al., 2023's observation. They noted underperformance compared to the source model in class imbalance settings and indicated that relying solely on normalization techniques may not always lead to significant improvements in non-i.i.d settings.

However, UnMix-TNS stands out by offering notable improvements to test-time adaptation methods. In some cases, it even helps these methods outperform the source model in different scenarios, including mixed domain. For instance, when integrating UnMix-TNS with ROID, we observed a significant performance leap over the source model, a feat not achieved with TBN as the normalization technique. It is noteworthy that UnMix-TNS outperforms the source model in the mixed domains scenario on DomainNet-126 and CIFAR10-C datasets. While the results for mixed domains have been reported with $K=128$ in the manuscript, increasing $K$ to $K=1024$ shows even more impressive results on CIFAR100-C. Specifically, the error rate on CIFAR100-C decreases to 45.8\%, outperforming the source model by 0.7\%.

Nevertheless, we can enhance the performance of UnMix-TNS in mixed domain adaptation scenarios by tuning the $\lambda$ hyperparameter. As detailed in Appendix A.4, this adjustment, guided by momentum batch normalization (MBN) principles, ensures consistent noise levels across different batch sizes. $\lambda$ is calculated as $\lambda=1-(1-\lambda_0)^\frac{B}{B_0}$ where $\lambda_0$ and $B_0$ represent the ideal values, and $B$ is the actual batch size. In mixed domain scenarios, with more varied styles and higher noise, we suggest a modified formula $\lambda=1-(1-\lambda_0)^\frac{B}{B_0*15}$ considering $15$ as the number of corruptions. This adjustment lowered the error rate on CIFAR100-C to 45.9\%, slightly outperforming the source model.

Reviewer comment (Writing):
If the temporal correlation mainly

Author response:
It seems the point about temporal correlation was not completed. Could you please elaborate on this, so we can address it more effectively?

Reviewer comment (Experiments):
Choice of K. The appendix discusses the influence of K. However, it is still unclear how to choose K in practice. How is it related to the number of classes and number of domains? (Although the number of domains might not be exposed.)

Author response:
We appreciate your inquiry regarding the selection of $K$ in our UnMix-TNS framework. The determination of $K$ is a critical element, as we have previously mentioned in our discussion of future work. To begin, we highlight that a single UnMix-TNS component can sometimes effectively represent multiple similar classes (e.g., car, truck, bus). This means that even with a relatively low $K$ value, our model can approximate the mean and variance of the test distribution accurately. This approach is particularly beneficial when class similarities are pronounced. However, in scenarios with a large number of classes, such as CIFAR100-C, DomainNet-126, or ImageNet-C, setting $K$ equal to the number of classes can lead to inefficiencies. Specifically, it can result in an excessive memory footprint, which is impractical for many applications. For instance, in our experiments with CIFAR100-C, we observed that setting $K$ to 100 actually decreased performance.

Furthermore, the optimal value of K varies depending on the test scenarios. In mixed domain scenarios, where there is a diverse range of styles in the test distribution, a larger $K$ is necessary to capture this diversity effectively. This contrasts with single domain scenarios, where the domain remains consistent, or continual domain scenarios, where the domain shifts sequentially. In these cases, a smaller $K$ may be sufficient due to the reduced variety in the test distribution.

In summary, the choice of $K$ in UnMix-TNS is a balance between the complexity of the classes and domains and the practical considerations of model efficiency. While a higher $K$ may offer more nuanced representation in diverse settings, it also comes with trade-offs in terms of computational and memory requirements. Our ongoing research aims to refine these guidelines further and provide more concrete recommendations in various application contexts.

Reviewer comment (Presentation):
Figure 2 is pixelated. Please consider improving the dpi.

Author response:
Thank you for your feedback regarding minor weakness. We acknowledge the pixelation issue of Figure 2 and have improved the dpi to enhance its clarity in the revised manuscript. Also, in the revised manuscript (Section 2.1, second paragraph, Test-time adaptation under label temporal correlation), we have clarified that it specifically refers to the temporal correlation of labels.  We thank you for noting the font style inconsistency after Equation (1) on Page 5; we corrected this in the revised manuscript. Also, we appreciate your suggestion to replace $\sim$ with $\approx$ for clarity in denoting approximation rather than distribution. This change is implemented to enhance the precision of the expressions in the paper.

Reviewer comment (Writing):
Page 5 after Equation (1), two exp have different font style. Also I recommend changing $\sim$ to $\approx$ since $\sim$ usually means “following the distribution of”.

Author response:
Thank you for your feedback regarding minor weakness. We acknowledge the pixelation issue of Figure 2 and have improved the dpi to enhance its clarity in the revised manuscript. Also, in the revised manuscript (Section 2.1, second paragraph, Test-time adaptation under label temporal correlation), we have clarified that it specifically refers to the temporal correlation of labels.  We thank you for noting the font style inconsistency after Equation (1) on Page 5; we corrected this in the revised manuscript. Also, we appreciate your suggestion to replace $\sim$ with $\approx$ for clarity in denoting approximation rather than distribution. This change is implemented to enhance the precision of the expressions in the paper.

Reviewer comment (Theory):
In Equation (5), it seems like all $\mu_{k, c}^0$ for different $k$ distribute on the line of $\mu_c + t \sigma_c$. Is that intentional? Or it makes more sense if they do not have such low-rank structure?

Author response:
In response to the reviewer's comment on Equation (5) regarding the distribution of $\mu_{k,c}^0$, we would like to clarify that $\mu_{k,c}^0$ is a scalar quantity representing the mean of the $k^{th}$ UnMix-TNS component for channel $c$, and $t$ is a scalar sampled from a Normal Distribution $\mathcal{N}(0,1)$. Given this, the vector $\mu_k$ does not lie on a line and has no low-rank structure.

Reviewer comment (Theory):
Equation (7). Are there any insight on using cosine similarity instead of L2 distance?

Author response:
We thank the reviewer's comment. The L2 distance is susceptible to the curse of dimensionality, and as the number of channels increases, the contrast between L2 distances ($b^{th}$ instance's statistics and the $K$ statistic components) tends to diminish. Consequently, it becomes necessary to decrease the temperature $\tau$ when computing the assignment probability $p_{b,k}$ in Equation (10) to avoid updating all $K$ components with the same weight. Such a distance function, therefore, requires careful setting of the temperature for every layer. On the other hand, cosine similarity does not suffer from this problem and only requires fixing one unique and shared temperature for all layers. Thus, we have opted to retain this score.