Technical Summary: USE – A Unified Self-Ensembling Framework for Test-Time Prompt Tuning

1. Problem Statement

Vision-Language Models (VLMs), such as CLIP, have demonstrated strong generalization in zero-shot settings but remain vulnerable to distribution shifts in downstream tasks. While Test-Time Adaptation (TTA) offers a solution by leveraging unlabeled test samples, existing methods often suffer from inconsistencies between their optimization and inference stages.

Specifically, the pioneering Test-Time Prompt Tuning (TPT) method optimizes textual prompts by minimizing marginal entropy over averaged predictions of augmented views. However, TPT and its variants typically rely on a "weak" augmentation (the original image) for final inference, while the optimization stage utilizes an average of "strong" augmentations. This discrepancy creates a misalignment between the pseudo-labels used for optimization and the final prediction strategy. Furthermore, existing approaches often underweight or exclude the original test image during the construction of pseudo-labels, despite evidence that the original image can be more reliable than certain augmented views in zero-shot classification.

2. Methodology

The paper proposes USE (Unified Self-Ensembling), a framework that unifies the optimization and inference stages through a consistent self-ensembling strategy.

2.1 Revisiting TPT as Self-Training

The authors first reinterpret TPT's objective. By introducing an auxiliary distribution $q$ (a stop-gradient version of the average prediction), the marginal entropy minimization objective is decomposed into a Reverse Cross-Entropy (RCE) term and a KL divergence term. Since the KL term is optimized when $q$ is fixed, the objective effectively reduces to minimizing RCE. This reveals that TPT implicitly learns from self-generated pseudo-labels, where the pseudo-label is constructed by averaging predictions from selected low-entropy augmented views.

2.2 Self-Ensembling (SE) Strategy

To improve the quality of these pseudo-labels, the authors introduce a Self-Ensembling (SE) strategy. Empirical analysis shows that while strong augmentations are useful, the weak augmentation (the original test image) often holds distinctive value and should not be discarded.

The SE strategy dynamically aggregates predictions from both weak ($p_w$) and strong ($p_s$) augmentations. It employs an entropy-based weighting coefficient $\beta$ to adaptively adjust the contribution of the test image:

1. Ranking: The relative self-entropy rank $\delta$ of the weak augmentation is computed against the set of strong augmentations.
2. Weighting: A linear function rescales $\delta$ into a weight $\beta$ (constrained between 0.3 and 0.7). If the weak augmentation has lower entropy (higher confidence) than most strong augmentations, it receives a higher weight.
3. Pseudo-Label Construction: The final pseudo-label $q(x)$ is computed as:
   $$q(x) = \beta \cdot p_w(x) + (1 - \beta) \cdot p_s(x)$$

2.3 The Unified Framework (USE)

USE ensures consistency between the two stages of TTA:

• Optimization Stage: The textual prompt is optimized using the RCE loss, where the target pseudo-label $q(x)$ is generated via the SE strategy.
• Inference Stage: Unlike TPT, which predicts solely on the weak augmentation, USE applies the same SE strategy to generate the final prediction $\hat{p}(x)$. This unifies the objective of pseudo-label estimation and final inference.

Additionally, the authors introduce a skip technique to reduce computational cost. If the prediction from the weak augmentation is identical to the selected strong augmentations, the optimization step is skipped, as the prediction is deemed reliable.

Notably, the SE strategy can function independently as an optimization-free TTA method, requiring no backpropagation.

3. Key Contributions

• Reinterpretation of TPT: The paper reveals that TPT implicitly learns from self-generated pseudo-labels, framing the marginal entropy objective as a Reverse Cross-Entropy loss against a stop-gradient target.
• Self-Ensembling (SE): A simple, effective strategy that adaptively weights weak and strong augmentations to generate more reliable pseudo-labels. It serves as a standalone, optimization-free TTA method.
• Unified Framework (USE): A TTA framework that aligns the optimization objective with the inference objective by applying the SE strategy in both stages, ensuring consistency.
• Plug-and-Play Capability: SE is shown to be a lightweight module that can be integrated into existing TTA methods to consistently enhance performance without significant computational overhead.

4. Experimental Results

The authors evaluated USE and SE on ImageNet, its Out-of-Distribution (OOD) variants (ImageNet-A, V2, R, Sketch), and ten fine-grained datasets (e.g., DTD, Flowers102, EuroSAT) using ViT-B/16 and ResNet-50 backbones.

• Performance:
  • Optimization-Free: SE achieved the best or second-best results among optimization-free methods, outperforming CLIP, MTA, and ZERO. On the ResNet-50 backbone, SE achieved an average accuracy of 47.80% on OOD datasets, surpassing many optimization-based methods.
  • Optimization-Based: USE consistently outperformed baselines like TPT, C-TPT, R-TPT, and STS. On fine-grained datasets with ViT-B/16, USE improved the average accuracy from 64.84% (TPT) to 65.35%.
  • Robustness: When initialized with CoOp (a stronger prompt initialization), both SE and USE maintained top-tier performance, demonstrating robustness to different pretrained configurations.
• Integration: Integrating SE as a plug-in module into existing methods (TPT, C-TPT, R-TPT, STS) consistently yielded accuracy gains. For instance, adding SE to C-TPT increased its IN&O accuracy from 61.09% to 64.27%.
• Efficiency:
  • SE is highly efficient, requiring only 0.059 seconds per sample and 1.06 GB of memory, comparable to 64 forward passes of CLIP but with higher accuracy.
  • USE, utilizing the skip technique, reduced the average testing time to 0.144 seconds per image. Over 60% of test samples were skipped without accuracy degradation.

5. Significance and Claims

The paper claims that the primary significance of USE lies in its unification of optimization and inference objectives. By ensuring that the pseudo-labels used for adaptation are constructed using the same strategy as the final inference, USE eliminates the performance gap often found in TTA methods.

The authors emphasize that:

1. Consistency is Key: The performance gains are largely attributed to the consistency provided by the SE strategy across both stages.
2. Simplicity and Effectiveness: The SE strategy is a simple, entropy-based weighting mechanism that does not require complex training objectives or additional parameters, yet it yields substantial performance improvements.
3. Versatility: The framework is versatile, functioning effectively as both a standalone optimization-free method and a consistent enhancement for optimization-based methods.

The paper concludes that while the method introduces some computational overhead due to generating multiple augmentations (a common trait of episodic TTA), the use of the skip technique and the high accuracy gains make it a practical and effective approach for improving VLM robustness against distribution shifts.