TinyTTA: Efficient Test-time Adaptation via Early-exit Ensembles on Edge Devices

Thanks for all your positive comments. Please see below our response to the specific weakness and questions.

Q1: Unclear how to use activation memory to split the network. How did you analyze the activations for grouping the models into sub-modules? Did you apply some principle here, or was it an ad-hoc grouping?

A1: Per-layer memory profiling is based on TinyTL[1], as discussed in lines 141-148 and 253-254. The grouping is post-hoc by visualizing the similar memory usage of adjacent layers as discussed in line 146. We will further elaborate this in Appendix D.

Q2: Proposed early-exists is not a novel technique but is well-known in the literature. What extent self-ensembling and early-exists is really new or just a re-brand of existing methods?

A2: We acknowledge that early-exit networks have been explored in the literature, primarily for inference efficiency. However, they have never been explored under data distribution shifts and for adaptation purposes. Our work is the first to introduce early exits for on-device test-time adaptation in edge devices, enhancing both inference efficiency and model accuracy with data distributional shifts. Specifically, our approach innovatively integrates weight standardization (WS) within the heads of early exits and designs these heads for subsequent modules. This ensures adaptation relies solely on WS in the heads, maintaining both inference speed and model accuracy with data distributional shifts. Moreover, our approach uniquely combines self-ensembling (frozen after fine-tuning), early exits, and WS normalization. This innovative design significantly improves both inference speed and accuracy. Additionally, our method brings novelty in hardware deployment by introducing the TinyTTA Engine, enabling the first on-device test-time adaptation with unlabeled data.

We will include more references on early-exit networks in the Related Work section and further clarify this novelty in Section 3.2.

Q3: While the evaluation is generally good, it does not highlight the memory overhead of introducing early exists.

A3: We now conducted experiments to explore the overhead in memory and the number of parameters by adding new prediction heads. The results are below:

According to the results of the above table, we observe the following:

• Adding new prediction heads results in a minimal memory increase. For instance, MobileNet shows an overhead ranging from 0.01 MB to 0.17 MB (with parameters ranging from 2.7K to 43.30K), which is negligible compared to the 512 MB MPU memory.
• Different models exhibit varying overheads. EfficientNet shows the highest increase in both memory (up to 0.75 MB) and parameters compared to MobileNet and RegNet.
• Overall, compared to the 512 MB MPU memory deployed, the overhead remains relatively small. We will extend this discussion in Appendix D.

Q4: Citations can be improved: There are a few Arxiv papers that have been published already.

A4: We will thoroughly review our paper and update our citations using DBLP to ensure high-quality references.

Q5: The authors mention data distributional shifts multiple times in the paper. I suggest to clarify this a bit more, since it is unclear to me if we look at shift in the label space or data space.

A5: The shift occurs in the data space. As illustrated in Fig. 2, for an image classification task with an image as the input, the leftmost cat image shows a small distribution shift, as it still clearly depicts the cat; in this case, the severity is 1. Gradually, in the second and third images, the cat’s image becomes noisier and harder to discern with the human eye, corresponding to severities of 3 and 5, respectively. We will further clarify these distributional shifts in the data space in Section 3.

Q6: What do you mean when you write, "activations consume much more memory compared to weights"?

A6: Activations consume much more memory than weights in neural networks because the value of activations need to be stored for every layer during backpropagation, especially when using large batch sizes and deep networks. For instance, consider a single CNN layer with an input image (e.g., CIFAR100C) size of 224x224x3 and a batch size of 1. The input activations alone would be around 1 (batch size) * 224 * 224 * 3 (channels) ≈ 1.8 million values, and the output activations after a convolutional layer with 64 filters would be 224 * 224 * 64 (channels/filters)≈ 34.1 million values. In contrast, the convolutional layer weights would only be 64 (channels/filters) * 3 * 3 * 3 (kernel size) ≈1,728 values. This indicates that the sheer volume of activations, especially in layers producing large feature maps, leads to significantly higher memory consumption compared to the relatively small number of weights. We will clarify this in Appendix A.

[1] TinyTL: Reduce Memory, Not Parameters for Efficient On-Device Learning. NeurIPS 2020

Dear reviewer. Thank you for reading our rebuttal! We believe that our response addresses your raised weaknesses (more justification and ablation study) and questions (more experimental results). If you agree that our response addresses the weaknesses and questions, please consider raising your score. If you have any outstanding concern, please let us know so that we can do our best to address them.

Thanks for all your positive comments. Please see below our responses.

Q1: Other type of data other than vision data?

A1: We tested on the Musan Keywords Spotting audio dataset [2] using a pretrained MicroNets [1] (86% accuracy on Speech Commands V2), which includes 35 speech commands with real-world noises, as below:

We observed:

• A 33% performance drop in the pretrained model under distribution shifts.
• TinyTTA achieved the highest accuracy of 0.61 (8% improvement).

Q2: Implementation Details.

A2: As in line 683, we will provide the source code of TinyTTA Engine. We will further elaborate on Appendix C.

Q3: Non-TTA method?

Q3:We implemented a new baseline using Test-Time Training (TTT) [3] with self-supervised rotation classification, using SGD (momentum 0.9, lr 1e-5), augmentation size 20 and batch size 1. Memory is tested on a Raspberry Pi Zero 2W.

• TinyTTA outperforms TTT by ~50% across four datasets and models.
• TTT requires, on average, double the memory compared to TinyTTA.

Q4: The experiments focus on specific type of MCU and MPU. Generalizeable to other chips?

A4: Yes, we selected these two platforms as they represent a range of devices. Specifically, since our MCU is based on Core M processors, TinyTTA can be directly deployed on ARM Cortex-M-based MCUs. We also chose one of the smallest Raspberry Pi to ensure that other Pis can also run TinyTTA. Further clarification will be provided in Appendix C.1.

Q5: Ablation study of early exit w/ and w/o model updates.

A5: The accuracy of the model w/o updates is discussed in Fig. 5. The memory usage (MB) w/o updates is discussed below.

• The memory usage w/ and w/o adaptation is very similar.
• TinyTTA generally requires limited memory to perform on-device TTA. We will update Fig. 7 and Section 5.4.

Q6: Computational overhead of self-ensembles and early exits?

A6: Self-ensembles are conducted offline (line 465). We analysed the overhead in memory and the number of parameters of the early exits:

Adding new prediction heads results in minimal memory increase, e.g., MobileNet's overhead is 0.01 MB to 0.17 MB, negligible compared to 512 MB MPU memory. We will extend this in Appendix D.

Q7: Motivation for using entropy thresholding? Sensitivity of hyperparameters?

A7: Entropy thresholding, as used in many TTA methods like [4], avoids high-entropy, less reliable samples to maintain TTA performance. Hyperparameters, including the entropy threshold, are determined post hoc [3] and discussed in Appendix G. We will add a table of layer exits and other hyperparameters in Appendix D.

Q8: Energy consumption compared baseline methods?

A8: We now compared the latency and energy consumption on Raspberry Pi zero 2 W MPU using CIFAR10C:

• TinyTTA has an inference time of 0.22 seconds per sample and energy consumption of 0.122 mWh, showing high efficiency.
• TinyTTA outperforms baselines, reducing latency by 12% (1500 seconds) and energy consumption by 12% (0.83 Wh) compared to EATA.

[1] MicroNets. MLSys 2021
[2] Importantaug. ICASSP 2022
[3] Test-Time Training with Self-Supervision for Generalization under Distribution Shifts. ICML 2020
[4] Efficient Test-Time Model Adaptation without Forgetting. ICML 2022

Thanks for all your positive comments. Please see below our response to the specific weakness and questions.

Q1: Lack of discussion on design choices: why can't the "fine-tune bias only" technique from TinyTL [28] be used in test-time adaptation? What is its performance compared to only fine-tuning the early exit header proposed in this work?

A1: TinyTL [28] needs labels for incoming samples to adapt the model, whereas TTA typically handles situations where only samples are available without labels. This is more practical in real-world settings, as providing labels is often challenging due to unforeseen noise and environmental factors. We now conducted "fine-tune bias only" via entropy minimization and compared it with adjusting the exits. Specifically, during TTA, only biases are allowed to be updated. The results are below:

Based on the experiment, we can observe that:
(1) Adjusting bias alone could not achieve reliable TTA performance.
(2) TinyTTA is relatively stable across four datasets and models.
We consider that the results are related to the characteristics of TTA, which aims to align data distribution shifts by adjusting the mean and variance. Adjusting the bias alone is insufficient to maintain reliable performance across different datasets and conditions. We will incorporate these new results in Appendix E under a new heading: "Comparison with Updating Bias Only.".

Q2: Why is there a "lack of support for normalization layers on MCUs"? Since the authors have developed their own MCU library, why can't the normalization layer be added to the library?

A2: Normalization layers are designed to work with mini-batches of data. However, due to limited memory resources, MCUs typically only allow a single batch of data. As such, normalization layers can technically be added to libraries for MCUs. In practice, the norm layer and the convolutional layer operations are combined into a single convolutional layer operation in order to save the computation and memory as in [1,2]. We will emphasize this further in the paper at Appendix A “Modulating and Finetune TTA.”

Q3: Limited experiments: As the authors mentioned in the Conclusion section, this work only targets image data. Thus, it is unclear whether the design in this work can be generalized to other applications. For example, will only updating the early exit header be sufficient for other applications?

A3:We conducted a different real-world distribution-shifted audio data modality experiment using a pretrained model on MicroNets [1], which was trained on Speech Commands V2 with 86% accuracy. This dataset contains 35 keywords such as “yes,” “no,” “forward,” etc. We tested the model on the Musan Keywords Spotting test dataset [2], which includes 35 speech commands under real-world noises such as dial tones, fax machine noises, car idling, thunder, wind, footsteps, rain, and animal noises. The setting aims to adapt the pretrained speech command model to real-world scenarios. TinyTTA parameters are: learning rate (lr) = 1e-5, batch size of 1, the SGD optimizer with a momentum of 0.9, and self-ensemble of early exit layers at [3, 5, 7]. The results are as follows:

Based on the experiment, we can observe that:
(1) The pretrained model could experience a performance drop of 33% in distribution shift settings.
(2) TinyTTA improved accuracy by 8% over the baseline, showing strong resilience to various noises.
(3) TinyTTA achieved the highest accuracy of 0.61, significantly outperforming other methods (the highest among the state-of-the-art techniques is CoTTA with 0.23).

Q4: Lack of details on the TinyTTA Engine: Since the algorithm is not entirely novel.

A4: Our approach uniquely combines self-ensembling (frozen after fine-tuning), early exits, and WS normalization. This innovative design significantly improves both inference speed and accuracy. Additionally, our method brings novelty in hardware deployment by introducing the TinyTTA Engine, enabling the first on-device test-time adaptation with unlabeled data. We will provide more details in Appendix C including details of backpropagation, operators, layerwise update strategy, and dynamic memory allocation.

Q5: Will you open-source the TinyTTA engine?

A5: Yes. As stated in line 683, The TinyTTA Engine Code will be made fully publicly available upon acceptance.

Q6: Is there any real-world case to show that updating the model via backpropagation, instead of simply switching some modes, makes a significant difference for tiny models?

A6: In realistic scenarios of TTA, we do not have knowledge of the given target domain. Hence, it is difficult to switch to the right model suitable for the target domain. Additionally, note that our target hardware consists of MCUs with extremely limited storage, typically at most 1 MB. Even if we have knowledge of the target domain for a stream of data, we can only store up to 2-3 tiny models (refer to Table 1 for the memory and storage requirements of a single model). We will clarify this in Appendix A.

[1] Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference. CVPR2018
[2] TensorFlow Lite Micro. MLSys 2021.

Dear reviewer. Thank you for reading our rebuttal! We believe that our response addresses your raised weaknesses (more justification and ablation study) and questions (more experimental results). If you agree that our response addresses the weaknesses and questions, please consider raising your score. If you have any outstanding concern, please let us know so that we can do our best to address them.

We appreciate the insightful comments. Please see below our responses.

Q1: In Fig.3, initial layers occupy much higher memory than later layers. Is it better to avoid pass through some of the initial layers and active more later layers?

A1: Memory usage primarily concerns activations, which store the outputs at each layer. This storage is essential during backpropagation, as computing gradients requires retaining both the outputs and gradient values for each node in each layer. However, we only update the Weight Standardization (WS) layer in the heads at early exits and freeze the submodules after self-ensembling (cf. Fig 2). This process does not require backpropagation to the submodules, thus eliminating the need for activation memory. Consequently, it does not increase memory usage.

Additionally, avoiding the initial layers is not feasible because these layers capture crucial information from the input. Skipping them would result in incomplete information being learned, leading to degraded performance. We will further clarify this in Appendix A, “Modulating and Finetune TTA”.

Q2: The authors mention "After training, only the submodules and early exits will be deployed on MCUs." It seems that the method first do self-ensemble offline which suffers from extra training cost.

A2: The self-ensemble model is trained offline, so there is no additional training cost on the device. The training procedure is discussed in Appendix C.2. We will include a heading for the first paragraph as 'Pre-training of Self-ensemble' and further clarify this point.

Q3: Since source data should be unavailable. How does the authors train the self-ensemble networks?

A3: The self-ensemble networks are trained offline, so we still use the source data, the same as related work in the literature [2]. However, once deployed on the device, the adaptation using TinyTTA does not have access to the source data, adhering to the standard TTA pipeline and the same configuration as in [2]. We will clarify this point in Section 3.1.

Q4: What is the difference compared to [1]?

A4: We discussed [1] TinyEngine in Appendix C.4. Specifically, TinyEngine focuses on on-device training (with labeled data), whereas TinyTTA Engine (ours) focuses on on-device test-time adaptation (with unlabeled data). TinyEngine pre-determines the layers and channels to be updated before deployment statically (i.e., as a binary file), executing these updates at runtime. In comparison, TinyTTA is dynamic during inference towards exits in submodules, enabling the exiting of high-entropy samples for reliable TTA. To enable TinyEngine for TTA, the only viable solution is using TENT [3] to finetune with entropy minimization. To this end, we compared TinyEngine (dubbed as TE) using TENT on a Raspberry Pi Zero 2W, using batch size 1, with TinyTTA (ours) in terms of accuracy. All configurations are the same as in Appendices B and C. The results are below:

We can make the following observations:

• Powered by the TinyTTA Engine, TinyTTA generally performs stably as it allows for dynamically exiting high-entropy samples.
• TE is unable to perform stable TTA across all datasets. We will update the paper in Appendix C.4.

Q5: Which part of the model will be trained on device?

A5: After deployment, as shown in Fig. 2, only the exits will be updated on-device, while the remaining parts of the model are frozen, ensuring both high TTA accuracy and low memory usage.

Q6: Why does activation memory change per layer, and no weight memory in Fig 3, compared to "fine-tuning per-layer memory usage"?

A6: The primary memory usage for activations is determined by the size of the last output tensor of each layer, essentially storing each layer’s outputs. Since each layer’s output shape is different, their activation memory will accordingly be different. The weight memory is very small (a few KBs) compared to the activation memory. The weight memory usage for modulating TTA, i.e., the change of two parameters, Scale (γ) and Shift (β), is relatively small and not visible in Fig 3.

Consider a single CNN layer with an input image size of 224x224x3 and a batch size of 1. The input activations alone would be around 1 (batch size) * 224 * 224 * 3 (channels) ≈ 1.8 million values, and the output activations after a convolutional layer with 64 filters would be 224 * 224 * 64 (channels/filters)≈ 34.1 million values. In contrast, the convolutional layer weights would only be 64 (channels) * 3 * 3 * 3 (kernel size) ≈1,728 values. This indicates that the sheer volume of activations, especially in layers producing large feature maps, leads to significantly higher memory consumption compared to the relatively small number of weights. We will clarify this in Appendix A.

Q7: The authors said "to entirely omit usage of normalization layers ...". Isn't WS normalization a normalization layer?

A7: We discussed how to deploy Weight Standardization (WS) normalization in lines 189-190 and Fig 4. Specifically, WS will be applied within the CNN exit layer (i.e., a new CNN layer which was introduced by TinyTTA Engine during deployment) to avoid using batch normalization layers. We will further clarify this in Section 3.3.

[1] On-device training under 256kb memory." NeurIPS 2022
[2] EcoTTA: Memory-Efficient Continual Test-time Adaptation via Self-distilled Regularization, CVPR 2023
[3] Tent: Fully test-time adaptation by entropy minimization, ICLR 2021

Dear reviewer. Thank you for reading our rebuttal! We believe that our response addresses your raised weaknesses (more justification and ablation study) and questions (more experimental results). If you agree that our response addresses the weaknesses and questions, please consider raising your score. If you have any outstanding concern, please let us know so that we can do our best to address them.

Dear reviewers,

This is the final reminder for the reviewer-author discussions. It will end on August 13 11:59 AoE, and then we will start AC-reviewer discussions.

If you have already concluded the discussions with the authors, thank you so much!

If you have not responded to the author rebuttal yet, please do so immediately. We have been waiting for your response.

Best,
Your AC