Dual-Path Adversarial Lifting for Domain Shift Correction in Online Test-time Adaptation

TL;DR

The paper tackles performance degradation of vision transformers under domain shift during inference by proposing online test-time adaptation. It introduces Dual-Path Adversarial Lifting, which augments each ViT layer with a domain shift token and uses a prediction network alongside an update network in an adversarial, dual-path framework to separate domain noise from class information. The update path employs smooth optimization (Sharpness-Aware Minimization) for robust class-token updates, while the prediction path uses non-smooth optimization to better capture domain-specific shifts, supported by a theoretical discussion. Experiments across ImageNet-C, ImageNet-R, ImageNet-A, VisDA-2021, and Office-Home demonstrate consistent improvements over state-of-the-art fully online TTA methods, underscoring practical impact for real-time cross-domain robustness.

Abstract

Transformer-based methods have achieved remarkable success in various machine learning tasks. How to design efficient test-time adaptation methods for transformer models becomes an important research task. In this work, motivated by the dual-subband wavelet lifting scheme developed in multi-scale signal processing which is able to efficiently separate the input signals into principal components and noise components, we introduce a dual-path token lifting for domain shift correction in test time adaptation. Specifically, we introduce an extra token, referred to as \textit{domain shift token}, at each layer of the transformer network. We then perform dual-path lifting with interleaved token prediction and update between the path of domain shift tokens and the path of class tokens at all network layers. The prediction and update networks are learned in an adversarial manner. Specifically, the task of the prediction network is to learn the residual noise of domain shift which should be largely invariant across all classes and all samples in the target domain. In other words, the predicted domain shift noise should be indistinguishable between all sample classes. On the other hand, the task of the update network is to update the class tokens by removing the domain shift from the input image samples so that input samples become more discriminative between different classes in the feature space. To effectively learn the prediction and update networks with two adversarial tasks, both theoretically and practically, we demonstrate that it is necessary to use smooth optimization for the update network but non-smooth optimization for the prediction network. Experimental results on the benchmark datasets demonstrate that our proposed method significantly improves the online fully test-time domain adaptation performance. Code is available at \url{https://github.com/yushuntang/DPAL}.

Figures (5)

• Figure 1: Representative examples of attention weights of the last layer in ViT-B/16 for different corruptions. The class token without adaptation focuses on almost the whole image due to the domain shift. The class token after our adaptation focuses on the object that is close to the clean image attention. There are also some peaky outlier values as explained in darcet2023vision.
• Figure 2: An overview of the proposed Dual-Path Adversarial Lifting method. During inference in the target domain, the domain shift token $P_l$, the prediction network $\Phi_{l}$, and the update network $\Psi_{l}$ are updated given each mini-batch testing samples. The dual-path lifting transformer (Left). The details of the lifting block in each layer (Right).
• Figure 3: Target domain accuracy for different prediction network smoothness factors. As the smoothness increases ($\rho$), the target accuracy in ImageNet-C decreases, indicating that smoothing similarity loss for the prediction network leads to sub-optimal generalization.
• Figure 4: The t-SNE visualization for the domain shift features in the first layer (top) and last layer (bottom) of ViT. Different colors represent different domains. It is able to learn domain-specific knowledge for different domains across layers.
• Figure 5: Ablation study in ImageNet-C at the highest severity. The DST represents the domain shift token.