Temporal Test-Time Adaptation with State-Space Models

TMLR Paper5244 Authors

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.

Submission Length: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission:

During the rebuttal, we have made the following changes to the paper (marked in blue):

- Modified Algorithm 1 to focus on a high-level overview of STAD and moved it to the main paper (Section 3.2).
- Moved the E-step of STAD-vMF from Appendix B.2 to the main paper (Section 3.3).
- Added the update equations for the EM algorithm of STAD-Gauss to Appendix B.1.
- Added Algorithm 2 and Algorithm 3, which summarize the EM updates for STAD-Gauss and STAD-vMF to Appendix B.1 and B.2 respectively.
- Expanded Section 6 with a discussion on the types of distribution shifts that STAD can effectively and less effectively handle.
- Revised wording in the introduction to clarify that we build on probabilistic SSMs, rather than structured SSMs such as Mamba.

Assigned Action Editor: ~Vincent_Dumoulin1

Submission Number: 5244