Cross-device Collaborative Test-time Adaptation with Zeroth-order Optimization and Model Merging

Technical Summary: Cross-device Collaborative Test-time Adaptation with Zeroth-order Optimization and Model Merging

1. Problem Statement

Deep neural networks (DNNs) often suffer from performance degradation due to domain shifts, where the distribution of test data differs significantly from training data. While Test-time Adaptation (TTA) offers a solution by updating models on the fly using unlabeled test data, existing methods predominantly rely on backpropagation (BP). BP is resource-intensive, requiring large memory to store computational graphs, intermediate activations, and loss gradients. This makes standard TTA infeasible for resource-limited edge devices (e.g., mobile phones, IoT sensors) which frequently encounter domain shifts but lack the computational capacity for BP.

Furthermore, while Cross-device Collaborative TTA (CDC-TTA) frameworks exist to share model information between resource-abundant and resource-limited devices, simply applying Zeroth-order Optimization (ZOO)—which avoids BP by using only forward passes—to high-dimensional model weights results in slow convergence and poor performance.

2. Methodology

The authors propose a novel TTA method designed for resource-limited devices within a CDC-TTA framework. The approach integrates three core pillars: Zeroth-order Optimization, Model Merging, and a unique preprocessing strategy.

A. Framework Overview
The system operates on a heterogeneous network of devices sharing model weights via a central server.

• Server-side: Performs heavy preprocessing on shared models before distribution.
• Device-side: Executes TTA using the preprocessed models without backpropagation.

B. Server-side Preprocessing
To mitigate the high-dimensional optimization difficulty inherent in ZOO and reduce redundancy, the server applies three sub-components to the "task vectors" (the difference between fine-tuned model weights and a base model):

1. Non-influential Weight Suppression (NWS): Applies top-k% filtering to each task vector, zeroing out elements with low absolute values that do not significantly influence the merging process.
2. Redundancy Reduction (RR) via Low-rank Approximation: Recognizing that task vectors from similar domains are redundant, the method applies Singular Value Decomposition (SVD) to the concatenated task matrix. This approximates the matrix with a lower rank, effectively grouping similar information.
3. Adaptive Rank Selection (ARS): Automatically determines the optimal rank for the low-rank approximation for each model weight group based on a recovery error threshold. This minimizes the number of parameters to be optimized while preserving necessary information.

C. Device-side Test-time Adaptation
The resource-limited device receives the base model, reduced task matrices, and initial merging coefficients.

1. Model Merging: The target model is constructed by merging the base model with the reduced task matrices using a linear combination of merging coefficients.
2. Zeroth-order Optimization (ZOO): Instead of optimizing millions of model weights, the device optimizes only the low-dimensional merging coefficients. The method uses Zeroth-order Stochastic Gradient Descent (ZO-SGD) with Simultaneous Perturbation Stochastic Approximation (SPSA). It calculates pseudo-gradients by measuring the loss difference between forward passes with positive and negative perturbations applied to the coefficients.
3. Memory Suppression: To further reduce memory usage, the method employs a "random seed trick." Instead of storing random perturbation vectors, the device samples a random seed for each iteration and resamples the perturbations, ensuring memory costs remain comparable to standard inference.

3. Key Contributions

• Resource-Friendly TTA: The paper proposes a TTA method executable on devices with small memory capacities by replacing backpropagation with Zeroth-order Optimization.
• Dimensionality Reduction via Model Merging: By optimizing merging coefficients rather than raw model weights, the method drastically reduces the optimization dimension, addressing the slow convergence typically associated with ZOO in high-dimensional spaces.
• Synergistic Preprocessing: A unique preprocessing pipeline (NWS, RR, and ARS) is introduced to trim non-influential weights and reduce inter-model redundancy, enhancing the efficiency of the ZOO process.
• Quantization Compatibility: The method is applicable to quantized models where backpropagation is inherently inexecutable, demonstrating improved prediction accuracy in such scenarios.

4. Experimental Results

The authors evaluated the method on common corruption benchmarks (CIFAR10-C, CIFAR100-C, ImageNet-C) and a style-transfer benchmark (Stylized-ImageNet).

• Performance: The proposed method consistently outperformed existing BP-free TTA methods (BN stats, T3A) and the closest competitor, CoLA, across various datasets and corruption types. For instance, on ImageNet-C with ResNet-50, the proposed method achieved a lower average error rate (59.4%) compared to CoLA (61.3%).
• Memory Efficiency: The method significantly reduced memory usage compared to conventional BP-based TTA. While BP-based methods required storing large computational graphs and gradients (e.g., ~9.37 GB in the authors' illustration), the proposed method reduced this to levels compatible with inference-only operations (e.g., ~1.84 GB).
• Ablation Studies: Experiments confirmed that the combination of Model Merging, ZO-SGD, and the full preprocessing suite (NWS, RR, ARS) yielded the best results, with each component contributing to the overall performance and efficiency.

5. Significance and Claims

The paper claims that its primary significance lies in making Test-time Adaptation feasible for resource-constrained edge devices, a critical requirement for real-world applications like automated surveillance, manufacturing quality assurance, and wearable healthcare. By leveraging the CDC-TTA framework, the method transforms the challenge of high-dimensional optimization into a low-dimensional problem through model merging and preprocessing. The authors assert that this approach successfully balances the trade-off between computational resource limitations and the need for robust, adaptive models in dynamic environments, offering a practical solution where traditional BP-based TTA fails.