Source-Target Unified Knowledge Distillation for Memory-Efficient Federated Domain Adaptation on Edge Devices

To conduct local inference on edge devices , it is necessary to deploy compact machine learning models on such devices . When such a compact model is applied to a new environment , its inference accuracy can be degradedif the target data from the new environment have a different distribution fromthe source data used for model training . To ensure high inference accuracy in thenew environment , it is indispensable to adapt the compact model to the target data . However , to protect users ’ privacy , the target data can not be sent to a centralized server for joint training with the source data . Furthermore , utilizing the target data to directly train the compact model can not achieve sufficient inferenceaccuracy due to its low model capacity . To this end , a scheme called source-target unified knowledge distillation ( STU-KD ) is developed in this paper . It first adapts a large source model to the target data on the edge device , and a large targetmodel is obtained . The knowledge of the large target model is then transferred to the compact model via knowledge distillation . Since training the large modelleads to large memory consumption , a domain adaptation method called lite residual hypothesis transfer is designed to achieve memory-efficient adaptation to the targe data on the edge device . Moreover , to prevent the compact model from forgetting the knowledge of the source data during knowledge distillation , a collabor tive knowledge distillation ( Co-KD ) method is developed to unify the source data on the server and the target data on the edge device to train the compact model.STU-KD can be easily integrated with secure aggregation so that the server can not obtain the true model parameters of the compact model . Extensive experiments co ducted upon several tasks of object recognition show that STU-KD can improve the inference accuracy by up to14.7 % , as compared to the state-of-the-art schemes . The results also reveal that the inference accuracy of the compact modelis not impacted by incorporating secure aggregation into STU-KD . 1 INTRODUCTION . Many computer vision ( CV ) applications , such as mobile robots , require local inference on edge devices because of the requirements on data privacy and low latency . To enable local inference on edge devices , it is necessary to deploy compact machine learning models on such devices . For example , considering an edge computing device Jetson Nano with 472 GFLOPS GPU and 4 GB memory space1 , which is commonly used for image recognition , a ResNet-18 model ( He et al. , 2016 ) is preferred over a ResNet-50 model , as the inference time of the former case ( 26 ms ) is much smaller than the latter one ( 64 ms ) ( Yang et al. , 2020 ) . Such compact machine learning models can be obtained by manual design based on experts ’ experience ( Sandler et al. , 2018 ; Howard et al. , 2019 ; Lin et al. , 2020 ) or by some automated machine learning ( AutoML ) techniques , e.g. , network compression ( Ning et al. , 2020 ; Li et al. , 2020b ) and neural architecturesearch ( He et al. , 2018 ; Liu et al. , 2019 ) . When an edge device with a compact model works in a new environment , the unlabeled target data collected from the new environment can have a different distribution from the labeled source data that are used to train the compact model , i.e. , domain shift ( Gretton et al. , 2009 ) occurs . Con- 1https : //developer.nvidia.com/embedded/jetson-nano-developer-kit sequently , the compact model suffers low inference accuracy on the target data . To ensure high inference accuracy in the new environment , it is necessary to adapt the compact model to the target data . A typical approach is to send the target data to a cloud server where the model is trained via unsupervised domain adaptation ( UDA ) ( e.g. , the methods inKang et al . ( 2019 ) ; Tang & Jia ( 2020 ) ; Xu et al . ( 2020 ) ) and then deployed back to the device . However , it leads to loss of data privacy . To avoid this issue , another type of approach is to train the compact model locally over the target data via UDA methods . It is doable on an edge device . For example , training ResNet-18 with batch size48 consumes nearly172 GFLOPs and1.1 GB memory space per batch , which is affordable for Jetson Nano . However , this type of approach can not achieve hgh inference accuracy on the target data due to limited model capacity of the compact model . To this end , a scheme called source-target unified knowledgedistillation ( STU-KD ) is developed in this paper . The key idea is to utilize a large model with sufficient model capacity to learn fine-grained representations of the target data , and then transfer its knowledge to the compact model . As shown in step 1 in Figure 1 , a large source model is loaded on the edgedevice and is then adapted to the target data to obtain a large target model . The challenges ofthis step are two-fold . First , the target data are unlabeled , so the fine-grained representations canot be learned via supervised learning . As a result , unsupervised domain adaptation ( UDA ) is needed instep 1 . Second , the edge device does not have the source data , since the volume of the source data can be too large to be stored on the edge device , or the source data can not be exposed to the edge devicue to confidentiality of these data . Without the source data , many UDA methods ( Kang et al. , 2019 ; Tang & Jia , 2020 ; Xu et al. , 2020 ) are not applicable for the adaptation of the large source model . Thus , source-free UDA methods ( Li et al. , 2020a ; Liang et al. , 2020 ; Liu et al. , 2021 ) must beemployed . However , existing sourcefree UDA methods , e.g. , source hypothesis transfer ( SHOT ) ( Liang et al. , 2020 ) , require retraining of the large source model , which leads to large memory consumption on the edge device . For example , if ResNet-50 is retrained with batch size48 on Jetson Nano , the computational cost is nearly 364 GFLOPs , which is affordable for Jetson Nano . However , the memory consumption per batch is nearly5 GB , exceeding the4 GB memory space limit of Jetson Nano . To tackle the challenges in step 1 of STU-KD , a memory-efficient UDA method called literesidual hypothesis transfer ( LRHT ) is designed by enhancing the architecture of a source-free UDA method such as SHOT . More specifically , the same loss function as that of SHOT is adopted , so unsupervised learning can be conducted with unlabeled target data . However , the featureextractor in SHOT must be replaced with a new architecture that can be trained in a memory-efficint manner . Thus , the new architecture is designed by adding lite residual ( LR ) modules ( Cai et al.,2020 ) to the feature extractor such that its outputs can be fine-tuned by training the LR modules only while keeping the parameters of the feature extractor fixed . As a result , in LRHT the training process of the large source model involves neither the feature extractor nor the classifier . Since training the LR modules generates a much smaller volume of activations and demands much lower memory fotprint , the training process of the large target model is highly memory-efficient . In other words , by using LRHT the large source model can be adapted to the target data to obtain a large target model in a memory-efficient manner . In step 2 of STU-KD ( in Figure 1 ) , the large target model is utilized to generate soft labels for the target data . The knowledge of the target data is then transferred to the compact model by training it over the target data and the soft labels via knowledge distillation ( KD ) ( Hinton et al. , 2015 ) . However , no source data are involved in the KD process , leading to the compact model gradually forgetting the the knowledge of the source data , i.e. , catastrophic forgetting ( McCloskey & Cohen , 1989 ) occurs . As a result , when the edge device encounters data from the source domain after the KD process , the compact model can suffer low inference accury . For example , consider a mobile robot whose compact model is trained via the above KD process . When it moves to an environment similar to that of the source domain , its compact model can not recognize the objects in that environment with a high accuracy . To avoid catastrophic forgetting , it is necessary to train the compact model considering both the target data and the source data . Thus , a collaborative knowledge distillation ( Co-KD ) method is designed as follows . On the edgeevice , a compact model is trained over the target data and the soft labels via KD . On the server , another compact model is trained over the source data via supervised learning . The challengehere is how to consolidate these two compact models into a global compact model . The setting of this c allenging problem is similar to that of federated learning ( FL ) ( McMahan et al. , 2017 ) . However , there exists one major difference . FL requires that different nodes have the same type of loss functions , while Co-KD has to use the cross-entropy loss on the server and the loss function of KD ( i.e. , Kullback–Leibler ( KL ) divergence in the state-of-the-art methods ) on the edge device . Due to such a difference , existing FL algorithms ( e.g. , McMahan et al . ( 2017 ) ; Karimireddy et al . ( 2020 ) ) arenot effective for consolidating the compact models in Co-KD . To this end , an alternating direction method of multipliers ( ADMM ) based learning algorithm is designed for Co-KD to consolidate thecompact model on the edge device and that on the server , as shown in step 2 in Figure 1 . More specifically , the training process of the global compact model is formulated into a consensus problemover the edge device and the server . The consensus problem is then divided into two subproblems by ADMM . The edge device trains its compact model by iteratively solving the subproblem related to the target data , and the server trains its compact model by iteratively solving the subproblem relat d to the source data . In the second step of STU-KD , the edge device needs to upload the parameters of the compact model to the server . It is possible for the server to recover some information of the target data from these parameters ( Zhu et al. , 2019 ) . To ensure privacy of the target data , secure aggregation ( Bonawitz et al. , 2017 ) is adopted in STU-KD so that the server can not obtain the true parameters of the compact model . Extensive experiments are conducted to evaluate the performance of STU-KD . The results show that STU-KD increases the compact model ’ s inference accuracy onthe target data by up to14.7 % while maintaining high inference accuracy on the source data , as compared to the state-of-art methods . Moreover , after employing secure aggregation , the performance of STU-KD is not affected .

Summary: A compact model deployed to a device may not work well if this device has a different data distribution. This work proposes to load a large pretrained model onto a device and then adapt it to the target data on the device. As directly training the full large model is too memory-heavy, this work proposes to adapt the large model's knowledge by training only part of its parameters on the local device data. To then transfer global model knowledge to the compact model this work proposes a collaborative knowledge distillation.