Temporal Test-Time Adaptation with State-Space Models

Authors: Mona Schirmer, Dan Zhang, Eric Nalisnick
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: test-time adaptation, state-space models, probabilistic modelling, dynamical systems

Abstract:
Distribution shifts between training and test data are inevitable over the lifecycle of a deployed model, leading to performance decay. Adapting a model on test samples can help mitigate this drop in performance. However, most test-time adaptation methods have focused on synthetic corruption shifts, leaving a variety of distribution shifts underexplored. In this paper, we focus on distribution shifts that evolve gradually over time, which are common in the wild but challenging for existing methods, as we show. To address this, we propose STAD, a probabilistic state-space model that adapts a deployed model to temporal distribution shifts by learning the time-varying dynamics in the last set of hidden features. Without requiring labels, our model infers time-evolving class prototypes that act as a dynamic classification head. Through experiments on real-world temporal distribution shifts, we show that our method excels in handling small batch sizes and label shift.

Decision: reject
Reviewer scores: 6, 3, 3, 6

Reviews and discussion:
<thread>
creation date:
1729204046266

thread_type:
official_review

<thread_object>
creator: reviewer_2f87

<content>
summary:
This paper proposes a method STAD designed to adapt a model’s classification head during inference to handle temporal distribution shifts, which are underexplored. The paper focuses on test-time adaptation, proposing to track the evolution of hidden features using SSMs. STAD dynamically updates class prototypes to reflect the evolving data distribution. Extensive experiments demonstrate that STAD performs well in temporal distribution shifts, label shifts, and small batch size settings. The paper evaluates its approach on several real-world datasets, showing robustness and performance gains compared to baseline methods.

strengths:
1. The paper is generally well-written with clear explanations of both the problem and the proposed solution.
2. The authors provide a strong theoretical foundation for the use of probabilistic SSMs to address temporal test-time adaptation, with detailed mathematical formulations of the dynamics model.
3. The empirical analysis is comprehensive, with evaluations on multiple datasets, including both real-world temporal shifts and synthetic corruption shifts.

weaknesses:
1. The paper does not clearly describe the sensitivity of STAD to hyperparameters which could impact the method’s robustness.
2. While STAD adapts the last layer to changing distributions, the reliance on class prototypes could be problematic when distributions are overfitted to the target domains and cause catastrophic forgetting.
3. The gradual distribution shift assumption is pretty strong especially during test time, which undermines the significance of the method. For example, ATTA [1] also targets distribution shifts during test time, which also includes some similar distribution shifts constrains. Given these constrains, they reason that a few labeled data is necessary to make the setting reasonable in real world scenarios. As what I believe, the gradual distribution shift assumption is a major limitation of this work. While it may be hard to make this method work perfectly on non-trivial domain shift datasets such as PACS and VLCS, there are two important experiments to do. (1)  Building boundaries on the ability limitations and applicability of this work on non-trivial domain shift datasets. (2) Validating that it is **not necessary** to use labels in the gradual shift setting by showing that STAD's performance is similar to the methods using labels (as Oracle), such as active online learning and ATTA.

[1] Gui, Shurui, Xiner Li, and Shuiwang Ji. "Active Test-Time Adaptation: Theoretical Analyses and An Algorithm." In The Twelfth International Conference on Learning Representations.

questions:
How can we quantify the gradual distribution shifts?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730425397117

thread_type:
official_review

<thread_object>
creator: reviewer_hzgz

<content>
summary:
This paper addresses the test-time adaptation problem under gradual distribution shifts. The proposed method models the gradual distribution shift in the representation space using linear state-space models, implemented through the von Mises-Fisher distribution. Experimental results on multiple real-world datasets demonstrate the effectiveness of the proposed algorithm.

strengths:
This work conducts experiments on several real-world datasets, enhancing the practical applicability of TTA algorithms in real-world settings.

By modeling gradual distribution shifts using linear state-space models, this work provides fresh insights in the field of TTA.

weaknesses:
The paper lacks discussion and empirical comparisons with related work, particularly in the field of gradual unsupervised domain adaptation, which addresses very similar problems (e.g., [1, 2]). As a result, the contribution of this paper may be overstated.

The gradual shift assumption represents a relatively simple form of distribution shift in dynamic environments, as its total variation is small. This simplicity limits the contribution of the work.

The proposed method appears to combine existing approaches, which limits its technical novelty.

References:

[1] Kumar, Ananya, Tengyu Ma, and Percy Liang. "Understanding self-training for gradual domain adaptation." International Conference on Machine Learning. PMLR, 2020.
[2] Wang, Haoxiang, Bo Li, and Han Zhao. "Understanding gradual domain adaptation: Improved analysis, optimal path, and beyond." International Conference on Machine Learning. PMLR, 2022.

questions:
How does the proposed method compare to self-training algorithms, e.g., Ref [1], on real-world tasks?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730626221506

thread_type:
official_review

<thread_object>
creator: reviewer_pomj

<content>
summary:
This paper introduces a novel Test-Time Adaptation (TTA) method, STAD, which leverages a state-space model to learn time-evolving feature distributions. It also proposes a new TTA setting that is more reflective of real-world scenarios.

strengths:
1. The authors propose a realistic TTA setting and conduct extensive experiments across multiple datasets to validate the method's effectiveness.
2. A new TTA method, STAD, is introduced, combining state-space models with a solid theoretical foundation.

weaknesses:
1. The method modifies and updates only the linear classifier, raising concerns about effectively handling covariate shifts. As shown in Table 4, adapting only the classifier on CIFAR-10C performs significantly worse than adapting the feature extractor. 

2. Experimentally, while a realistic TTA setting is proposed, comparisons with other similar settings and related methods are lacking. For instance, the following related works are not adequately compared:

   1. UniTTA: Unified Benchmark and Versatile Framework Towards Realistic Test-Time Adaptation. arXiv 2024.
   2. Towards real-world test-time adaptation: Tri-net self-training with balanced normalization. AAAI 2024.
   3. Robust test-time adaptation in dynamic scenarios. CVPR 2023.
   4. Universal Test-Time Adaptation Through Weight Ensembling, Diversity Weighting, and Prior Correction. WACV 2024.

   Additionally, only a small portion of common TTA datasets (CIFAR-C, ImageNet-C) are addressed, with limited focus on CIFAR-10C.

questions:
See the weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730659464291

thread_type:
official_review

<thread_object>
creator: reviewer_qpbz

<content>
summary:
The paper studied test-time adaptation, a setting where model parameters are updated based on incoming test features. It proposes STAD, a method to track gradual distribution shifts which only updates the last layer in a given network, based on the EM algorithm. The authors report results on several tasks.

strengths:
The paper is well-written and gives a principled justification for the proposed method. It considers real-world shifts and the proposed method seems to outperform the baselines generally.

weaknesses:
The paper does not appear to have any major weaknesses, but I am curious about the importance of the subfield (temporal test-time adaptation). Why is this the right setting to study (as opposed to others for distribution shift) and how does it compare to methods involving periodic retraining and unsupervised domain adaptation?

questions:
LIne 37 - why can't test-time training operate in this paradigm? what is the precise difference between TTA and TTT?
46 - Doesn't FMoW have labels? Why is this dataset suitable for the proposed setting, which claims that one of its benefits is not needing labels?
192 - missing a comma inside the parentheses?

How important is the linear Gaussian transition model assumption? Could your method make use of labels if they were made available?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732228371643

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Rebuttal

comment:
We thank all reviewers for the time and expertise they have invested in their reviews and their insightful comments and suggestions.

We are pleased to hear that our paper offers “fresh insights in the field of TTA” (HZGZ), has “a strong/solid theoretical foundation” (2f87, PoMj) and provides “principled justifications” (QpBz) for our “novel” method (PoMj). Reviewers appreciated our focus on temporal distribution shifts, which are “underexplored” (2f87), “more reflective of real-world scenarios” (PoMj) and helpful in “enhancing the practical applicability of TTA algorithms in real-world settings” (HZGZ). We are delighted that our “empirical analysis is comprehensive” (2f87) and we could demonstrate “robustness and performance gains compared to baseline models” (2f87, QpBz) via “extensive experiments” (2f87, PoMj). Lastly, our paper is “well-written” (QpBz,2f87), provides “clear explanations of both the problem and the proposed solution” (2f87) with “detailed mathematical formulations” (2f87).
# Clarification of the Setting
The only common point raised in the reviews (reviewers QpBz, HZGZ, 2f87) relates to the link between our precise setting and adjacent settings, particularly the link to gradual unsupervised domain adaptation (reviewer HZGZ) and TTA with supervision (reviewer 2f87).

These two settings have strictly more information at hand than test-time adaptation. Gradual domain adaptation assumes simultaneous access to the source and target data. Active test-time adaptation assumes it is possible to query labels at test time. In real-world settings, these assumptions may easily not be met. Privacy concerns or model access via an API can prevent access to the source. In dynamic environments where instant predictions are required, an intermediate annotating process may not be feasible. These restrictions have motivated the long line of test-time adaptation research (see [1,2] for surveys).

In our study, we aim to replicate the simple yet widespread and practical setting where a model is trained and then deployed to the real world without prior knowledge of the distribution shifts it may encounter. When the shift is too strong, retraining with supervision (e.g., with active TTA) is clearly necessary. However, when the shift is mild, there is an opportunity to keep the model running. This is the setting we address. We believe we achieve this effectively by using the well-established Wild-Time Benchmarks (NeurIPS 2022, >80 citations, used by, e.g., [3]), which were explicitly designed to capture real-world distribution shifts and realistic deployment scenarios.

To bring more clarity in distinguishing our problem setting from adjacent settings, we have expanded the related work section (Section 4) with a discussion of domain generalization and domain adaptation. We have also expanded the setting comparison in Table 1.

# Changes to the Paper
We have made minor changes to the paper to address the reviewers' comments:
- We have **added RoTTA as an additional baseline** (Tables 2, 3, and 4). Results confirm the experimental trends already observed: RoTTA fails to improve on all temporal distribution shifts.
- We have **added a discussion of adjacent fields** domain generalization and unsupervised domain adaptation to the related work (Section 4) and the setting comparison (Table 1).
- We have **added a more detailed discussion of additional TTA work** in Appendix C.
- We have **added another dataset** from domain generalization (Appendix E.3).
- We have **expanded the sensitivity study** in Appendix E.6.

Our paper now includes experiments on **4 types of distribution shifts** (temporal, reproduction, corruption, and style transfer) across **7 datasets** comparing against **8 baselines**. We further evaluate robustness on 12 different batch sizes for 2 label shift settings, in addition to an ablation study and qualitative results. Furthermore, we anchor our paper in past work by linking it to four related fields with additional in-depth discussion in the Appendix. We are positive about our theoretically grounded approach being well evaluated and embedded in past work.  

[1] Liang, Jian, Ran He, and Tieniu Tan. "A comprehensive survey on test-time adaptation under distribution shifts." International Journal of Computer Vision (2024): 1-34.

[2] Xiao, Zehao, and Cees GM Snoek. "Beyond Model Adaptation at Test Time: A Survey." arXiv preprint arXiv:2411.03687 (2024).

[3] Zhao, Hao, et al. "On pitfalls of test-time adaptation." ICML 2023. (see https://github.com/LINs-lab/ttab/tree/main)
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces STAD, a method leveraging state-space models for test-time adaptation under gradual distribution shifts. Addressing temporal shifts is of interest for robust learning. The paper provides comprehensive experiments and the results are promising. However, reviewers identified critical issues such as missing comparisons with related work and over-reliance on a simplified shift assumption. I find the amount of missing references to be concerning. There is also a lack of formal mathematical guarantees on the method (e.g. when the method provably works, generalization bounds etc). While I recommend rejection for the current submission, I believe this is a promising contribution and I recommend authors to address these issues. That is, (1) carefully incorporate the related literature, (2) provide stronger justification on when/why the stated assumptions are necessary, and (3) explore whether formal guarantees can be provided on the method e.g. what level of temporal shift it can tolerate or when it provably fails.