Weight Average Test-Time Adaptation (WATT)

• WATT is a test-time adaptation method that averages weights from diverse input variations to achieve smoother, flatter optima and robust generalization.
• It involves generating varied test data inputs, adapting using self-supervised losses, and combining adaptation trajectories, thereby reducing overfitting and error variance.
• Empirical results demonstrate that WATT significantly improves performance in 3D point cloud classification, vision-language models, and online TTA environments.

Weight Average Test-Time Adaptation (WATT) is a family of test-time adaptation (TTA) methods that enhance the robustness of neural networks to distribution shift by averaging model weights from multiple adaptation trajectories performed on test data. WATT systematically combines different sources of predictive variation—arising from data augmentations, input sampling, or prompt templates—to improve generalization, explicitly leveraging the weight averaging principle as a regularizer. The framework has been studied in several contexts including 3D point cloud classification, vision-LLMs (VLMs) such as CLIP, and universal online TTA for streaming and temporally correlated environments (Bahri et al., 2024, Osowiechi et al., 2024, Marsden et al., 2023).

1. Theoretical Underpinnings and Motivation

WATT builds on the insight that model parameters resulting from adaptation on varied input stochasticity tend to reside in distinct basins of the loss landscape. By averaging these adapted weights, WATT induces a bias toward flatter optima, which are associated with improved generalization and stability under domain shift (Bahri et al., 2024, Osowiechi et al., 2024). This mechanism also acts as a regularizer, mitigating the risk of overfitting or catastrophic forgetting that can affect conventional self-training or entropy-minimization adaptation protocols.

An additional theoretical justification arises from robust risk minimization: averaging weights over multiple input variations effectively minimizes the worst-case empirical loss within a parameter neighborhood, providing a practical proxy for robust optimization (Bahri et al., 2024). In settings where stochastic or adversarial corruptions yield high-variance adaptation trajectories, WATT enables error cancellation by smoothing out noisy or over-confident updates.

2. Core Algorithmic Structure

While instantiations of WATT vary by application domain, three universal steps characterize the workflow:

1. Input Variation Generation Multiple stochastic variants of the test data are created. In point clouds, these can be spatial samplings via Farthest Point Sampling (FPS) and K-Nearest Neighbors (KNN) (Bahri et al., 2024). For CLIP, diverse prompt templates yield alternative textual descriptions (Osowiechi et al., 2024). In universal TTA, test batches may correspond to temporally correlated or domain-mixed streams (Marsden et al., 2023).
2. Per-Variation Adaptation For each data variation, the model is adapted using a self-supervised loss. Entropy minimization (TENT), pseudo-label cross-entropy, or other objectives updating only normalization parameters (γ, β) are standard (Bahri et al., 2024, Osowiechi et al., 2024, Marsden et al., 2023). Let $X_m$ denote the m-th data variant and $\theta$ model parameters, the loss is typically:

$$L_{ent}(\theta; X) = -\frac{1}{|X|} \sum_{x \in X} \sum_{c=1}^C p_c(x; \theta) \log p_c(x; \theta)$$

1. Weight Averaging Final adapted parameters are obtained by averaging the per-variation weights:

$$\theta_{final} = \frac{1}{M} \sum_{m=1}^M \theta_m^{(T)}$$

Here $M$ is the number of variations, $\theta_m^{(T)}$ are weights after $T$ adaptation steps.

Variants include exponential moving average blending with source weights (momentum $\alpha$) (Marsden et al., 2023), or parallel/sequential averaging over prompt-induced adaptation paths in CLIP (Osowiechi et al., 2024).

3. Specific Instantiations and Application Domains

Variant	Input Variations	Adaptation Loss	Parameters Averaged
3D Point Cloud WATT (Bahri et al., 2024)	FPS+KNN stochastic patches	Entropy minimization (TENT)	BatchNorm/LayerNorm
WATT for CLIP (Osowiechi et al., 2024)	Diverse text prompt templates	Pseudo-label cross-entropy (transductive)	Visual encoder LayerNorm
Universal TTA (ROID) (Marsden et al., 2023)	Test batch streaming (temporal)	Self-training (SLR+SCE), diversity weighting	Norm param. EMA with source

In each case, only normalization parameters are typically adapted and averaged, keeping memory and compute overhead minimal (<2% for point cloud backbones (Bahri et al., 2024)).

4. Empirical Validation and Results

Substantial empirical gains from WATT have been demonstrated on both synthetic and real-world distribution shifts.

• 3D Point Clouds (ModelNet40-C, ShapeNet-C, ScanObjectNN-C):
  • Point-MAE: accuracy improved from 72.1% (TENT) to 75.0%.
  • PointNet: 60.4% to 61.3%.
  • DGCNN: 71.6% to 74.1%.
  • Consistent improvements under hard corruptions (occlusion, lidar dropout) and lower prediction variance (Bahri et al., 2024).
• CLIP Zero-Shot Vision-Language Benchmarks (CIFAR-10/100, VisDA-C, OfficeHome, PACS):
  • ViT-B/32, CIFAR-10-C: baseline 59.22% → WATT-S 73.82%.
  • Both parallel and sequential averaging modes yield robust accuracy increases, with sequential averaging preferred on stronger shifts.
  • Performance is maintained even at batch size = 1 (Osowiechi et al., 2024).
• Universal, Streaming TTA (ROID):
  • Recurrent exponential averaging with source prevents catastrophic forgetting on temporally correlated domain shifts.
  • Ablations show $\alpha \in [0.99, 0.9975]$ yields optimal stability; lower $\alpha$ degrades generalization, higher values limit adaptation (Marsden et al., 2023).

5. Key Hyperparameters, Overhead, and Practical Considerations

Canonical settings across domains include:

• Number of Variations ($M$ or $H$): 6–8 typically suffices; diminishing returns beyond this point (Bahri et al., 2024, Osowiechi et al., 2024).
• Adaptation Steps ($T$/$L$ per variation): One step per input variation is usually sufficient due to normalization-only adaptation (Bahri et al., 2024, Osowiechi et al., 2024).
• Learning Rate ($\eta$): $10^{-3}$ (standard), $10^{-4}$ for severe domain shifts (Bahri et al., 2024, Osowiechi et al., 2024, Marsden et al., 2023).
• Averaging Momentum ($\alpha$): ROID recommends $\alpha$ five to ten times larger than $\eta$; typically $\alpha=0.99$ for $\eta=10^{-3}$ (Marsden et al., 2023).
• Overhead: Additional runtime is marginal (e.g., 26ms per batch for WATT vs. 21ms for TENT in PointNet with $M=6$); memory usage is dominated by normalization layers, resulting in minimal increase (<2%) (Bahri et al., 2024).

WATT is robust to small batch sizes, enhancing applicability to low-latency and online adaptation scenarios (Osowiechi et al., 2024).

6. Comparative Assessment and Related Methods

WATT generalizes the principle of weight averaging from stochastic weight averaging (SWA) in training (Osowiechi et al., 2024), but is unique in combining adaptation under multiple data or prompt stochasticities at test time. In contrast to output ensembling or adaptation on a single data trajectory, weight averaging consolidates information across heterogenous data-induced model landscapes, yielding improved resilience to shift.

Related frameworks integrate WATT as a regularization component:

• ROID (Universal TTA): Blends source and adapted weights at every step to inhibit forgetting (Marsden et al., 2023).
• CLIPArTT, TENT, TPT: WATT consistently outperforms these baselines across VLM and point cloud evaluations, particularly under strong corruptions (Osowiechi et al., 2024, Bahri et al., 2024).

Ablation studies confirm that weight averaging offers a superior trade-off compared to output or embedding averaging, is synergistic with diversity weighting and prior correction, and can be composed with other TTA ingredients (Osowiechi et al., 2024, Marsden et al., 2023).

7. Future Directions and Significance

WATT exemplifies a flexible and resource-efficient solution for real-world test-time adaptation under diverse shift regimes. Its theoretical grounding in robust optimization and flat-minimum seeking, combined with empirical effectiveness across vision and point cloud tasks, suggests broad applicability. Potential future extensions include generalizing weight averaging to meta-ensembling strategies, adaptation in continual and federated learning, and integrating stochastic augmentations tailored to other modalities or architectures (Bahri et al., 2024, Osowiechi et al., 2024).