STAG: Structural Test-time Alignment of Gradients for Online Adaptation

TL;DR

STAG addresses the challenge of test-time adaptation under extreme information scarcity by leveraging the classifier's intrinsic geometry. It derives class-wise structural anchors from self-structural entropy and aligns the online entropy gradient with these anchors through a closed-form, cosine-similarity regularizer, incurring near-zero overhead. The approach is architecture-agnostic and can be integrated with existing TTA objectives, delivering consistent gains across CNNs and Transformers on both image classification and continual semantic segmentation, especially in challenging online regimes. These results highlight the practical value of exploiting classifier geometry for stable, real-time adaptation in the wild.

Abstract

Test-Time Adaptation (TTA) adapts pre-trained models using only unlabeled test streams, requiring real-time inference and update without access to source data. We propose StructuralTest-time Alignment of Gradients (STAG), a lightweight plug-in enhancer that exploits an always-available structural signal: the classifier's intrinsic geometry. STAG derives class-wise structural anchors from classifier weights via self-structural entropy, and during adaptation analytically computes the predicted-class entropy gradient from forward-pass quantities, aligning it to the corresponding anchor with a cosine-similarity loss. This closed-form design incurs near-zero memory and latency overhead and requires no additional backpropagation beyond the underlying baseline. Across corrupted image classification and continual semantic segmentation, STAG provides broadly applicable performance gains for strong TTA baselines on both CNN and Transformer architectures regardless of the underlying normalization scheme, with particularly large gains under challenging online regimes such as imbalanced label shifts, single-sample adaptation, mixed corruption streams and long-horizon continual TTA.

Figures (8)

• Figure 1: The Overall Framework of STAG. (a) Phase 1: Structural Anchor Derivation. We derive structural anchors $\mathbf{a}_k$ by calculating the gradient of self-structural entropy from the classifier weights. (b) Phase 2: Online Gradient Alignment. During test-time adaptation, STAG analytically computes the online entropy gradient $\mathbf{g}_t$ from test samples and aligns it with the corresponding structural anchor $\mathbf{a}_{\hat{y}_t}$.
• Figure 2: Validation of classifier weights as structural anchors. (a) Softmax assignment probability heatmap for 20 selected ImageNet classes (ResNet50-GN). (b) Metric consistency across architectures: ResNet50-BN, ResNet50-GN and VitBase-LN. $H$, $\psi$, and Acc denote self-entropy, mean diagonal probability (MDP), and structural accuracy, respectively.
• Figure 3: Top-1 accuracy running curves across three different architectures on ImageNet-C Gaussian noise (severity 5)
• Figure 4: The t-SNE visualization of the feature space on ImageNet-C Frost (severity 5)
• Figure 5: Semantic segmentation performance (mIoU, %) between TENT and its STAG-integrated version on the Cityscapes-to-ACDC continual test-time adaptation tasks.