Year: 2021

Abstract:
To conduct local inference on edge devices, it is necessary t o deploy compact machine learning models on such devices. When such a compact mod el is applied to a new environment, its inference accuracy can be degraded if the target data from the new environment have a different distribution from the source data used for model training. To ensure high inference accuracy in the new environment, it is indispensable to adapt the compact model to the target dat a. However, to protect users’ privacy, the target data cannot be sent to a centr alized server for joint training with the source data. Furthermore, utilizing the t arget data to directly train the compact model cannot achieve sufficient inference accuracy due to its low model capacity. To this end, a scheme called source-targ et unified knowledge distillation (STU-KD) is developed in this paper. It first ad apts a large source model to the target data on the edge device, and a large target model is obtained. The knowledge of the large target model is then transferred t o the compact model via knowledge distillation. Since training the large model leads to large memory consumption, a domain adaptation method called lite residu al hypothesis transfer is designed to achieve memory-efficient adaptation to the ta rge data on the edge device. Moreover, to prevent the compact model from forgett ing the knowledge of the source data during knowledge distillation, a collabo r tive knowledge distillation (Co-KD) method is developed to unify the source da ta 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 serve r cannot obtain the true model parameters of the compact model. Extensive experimen ts co ducted upon several tasks of object recognition show that STU-KD can imp rove 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 model is not impacted by incorporating secure aggregation into STU-KD.

text: 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 cannot be sent to a centralized server for joint training with the source data. Furthermore, utilizing the target data to directly train the compact model cannot 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 cannot 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.

heading: 1 INTRODUCTION
text: 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 cannot 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 cannot 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 cannot 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 cannot 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.

heading: 2 RELATED WORK
text: Cross Domain Model Compression.Many schemes (Yu et al., 2019; Chen et al., 2019; Feng et al., 2020; Dillard et al., 2020; Nguyen-Meidine et al., 2020) combine network compression with unsupervised domain adaptation (UDA) to obtain a compact model that performs well on the target data. These schemes adopt traditional UDA methods that assume thetarget data and the source data are stored in one place, which are not applicable when the targetdata and the source data are isolated. Moreover, these schemes need to train a large source model via UDA. The large memory footprint of training the large source model hinders the application of these schemes on edge devices.
On-device Adaptation. Recently, some schemes are developed to adapt an on-device mod l to the target data. TinyTL (Cai et al., 2020) assumes that the on-device model has sufficient model capacity, and adapts this model to the labeled target data with low memory footprint. However, TinyTL is not applicable to the scenario with unlabeled target data. MobileDA (Yang et al., 2020) directly trains a compact model via a UDA method called deep CORAL (Sun & Saenko, 2016). It employs a teacher model to boost the inference accuracy ofthe compact model. However, due to its low model capacity, directly training the compact model cannot achieves sufficient inference accuracy on the target data. Although a teacher model is employed, it is only trained over the source data, so the teacher model can only achieve limited improvement of the inference accuracy.
Domain Adaptation Under Data Isolation. Some schemes are designed to conduct UDA in a scenario where the source data and the target data are held bydifferent data owners. In federated domain adaptation (FDA), e.g., federated adversarial domain adaptation (Peng et al., 2020), the data owners can collaboratively train a globally shared model. In the training process, each data owner only shares its local updates rather than the target or source data. Schemes like source-free domain adaptation (SFDA) can accomplish UDA in the absence of the source data. In Li et al. (2020a), the
source model is trained over some generated target-style samples. In Liang et al. (2020), the source model is trained through information maximization loss over the target data and self-supervised pseudo-labeling. The SFDA schemes in (Kurmi et al., 2021; Liu et al., 2021) leverage generative models to generate fake source samples for UDA. However, to obtain a model with sufficient accuracy, both FDA and SFDA schemes require the source model tohave sufficient model capacity. However, such a requirement cannot be satisfied by a compact model with low model capacity.

heading: 3 METHODOLOGY
text: One centralized server andM edge devices (M ≥ 1) are considered in this paper. The labeled source data in the server are denoted as{xs, ys} wherexs represents the feature vector of a source sample andys the ground-truth label of the source sample. The source dataare sampled from the source domain denoted asPs. The unlabeled target data on the edge devices are denoted as{xt} wherext represents the feature vector of a target sample. The targetdata are sampled from the target domain denoted asPt. Each edge device needs to accomplish a classification task with K classes using a compact model. The compact model on the edge dvic m ∈ {1, ...,M} is defined as a functionfm, and its model parameters are denoted aswm. Given the above setup, source-target unified knowledge distillation (STU-KD) is developed to adapt the compact model to the target data.

heading: 3.1 LITE RESIDUAL HYPOTHESISTRANSFER
text: To boost the inference accuracy of the compact model, a largesource model with sufficient model capacity, defined as a functionfT , is adopted in STU-KD. It can learn more fine-grained representations of the target data compared with the compact model. Before being deployed to theM edge devices, the large source modelfT is pretrained over the source data by minimizing an objectiv functionJT (wT ) = E(xs,ys)[L(fT (xs;wT ), ys)], wherewT denotes the parameters of the large source model, andL(·, ·) is the standard cross-entropy loss. On edge devicem ∈ {1, ...,M}, the large source model is then adapted to the target data. The chall nges of the adaptation are twofold. First, the target data are unlabeled so that fine-grained representations of the target data cannot be learned via supervised learning. As a result, unsupervised domain adaptation (UDA) is needed. Second, edge devicem does not have the source data, so many UDA methods (Kang et al., 2019; Tang & Jia, 2020; Xu et al., 2020) are not feasible for the adaptation. Thus, source-free UDA methods (Li et al., 2020a; Liang et al., 2020; Liu et al., 2021) must be employed. However, existing source-free 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.
To tackle the above challenges of the adaptation, a memory-efficient UDA method called lite residual hypothesis transfer (LRHT) is designed by enhancing the architecture of a source-free UDA method such as SHOT. LRHT adopts the same loss function as that of SHOT (denoted asJS in Figure 2), so unsupervised learning can be conducted over the unlabeled target data. The details of the loss function of SHOT are provided in Appendix A. However, the feature extractor in SHOT must be replaced with a new architecture that can be trained in a memory-efficient manner. As shown in Figure 2, 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-tunedby the LR modules. The LR modules fine-tune the outputs of the feature extractor by changing the intermediate activations of the feature extractor. Define thei-th building block of the feature extractor as a functiongi and the LR module added to the it as a functionri. Given the activationsai−1 of the (i − 1)-th building block, LR
moduleri changes the activations of thei-th building block asai = gi(ai−1) + ri(ai−1), where gi(ai−1) denotes the original activations of thei-th building block, andri(ai−1) denotes the residual activations of LR moduleri. During the adaptation, the parameters of the feature extractor nd that of the classifier are fixed, and only the LR modules are trainedto fine-tune the outputs of the feature extractor. Compared to the case of directly training the featur extractor (like SHOT), training the LR modules generates a much smaller volume of activations and demands much lower memory footprint. Thus, 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.

heading: 3.2 COLLABORATIVE KNOWLEDGE DISTILLATION
text: The next step is to transfer the knowledge of the large targetmodel to the compact model.

heading: 3.2.1 CASE I: SINGLE EDGE DEVICE
text: Consider the case of a single device that contain a compact modelf1 with the model parametersw1. On the edge device, the large target model is leveraged to generat the soft labels of the target data. The knowledge of the target data is then transferred to the compact model via knowledge distillation (KD) (Hinton et al., 2015). More specifically, the compact model is trained over the target data and the soft labels via minimizing the loss function of KD. However, no source data are involved in the KD process, which leads to the compact model gradually forgetting the knowledge of the source data. That is, catastrophic forgetting (McCloskey & Cohen,1989) occurs.
To prevent catastrophic forgetting, the target data and thesource data need to be unified to train the compact model in above process of knowledge transfer. Thus,a collaborative knowledge distillation method (Co-KD) is designed as follows. On the edge device, a compact model is trained over the target data and the soft labels via KD. The loss function of KD, denoted asJK , is written as
JK(w1) = Ext [DKL(p(zT , τ)‖p(zc, τ))],
wherezc = f1(xt;w1) is the logit of the compact model,zT = fT (xt;w′T ) is the logit of the large target model with parametersw′T (i.e., the soft label), andp(z, τ) denotes the soften probability vector with a temperatureτ of a logit z. On the server, another compact model is trained over the source data via a cross-entropy loss denoted asJC . The expression ofJC is
JC(w1) = E(xs,ys)[L(f1(xs;w1), ys)].
The challenge here is how to consolidate the two compact models into a global compact model. The setting of this challenging problem is similar to that of federated learning (FL) (McMahan et al., 2017). However, there exists one major difference. FL requir s 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 the difference, existing FL algorithms (e.g., McMahan et al. (2017); Karimireddy et al. (2020); He et al. (2020a)) are not effective for consolidating the compact models in Co-KD.
To this end, an alternating direction method of multipliers(ADMM) (Boyd et al., 2011) based learning algorithm is developed for Co-KD to consolidate the compact models. The training process of the global compact model is formulated into a consensus problemΩ over the edge device and the server:
min w0,w1 αJK(w1) + (1− α)JC(w0)
s.t.w1 = w0, (1)
whereα is a balancing hyperparameter, andw0 denotes the parameters of the compact model trained on the server. For problemΩ that has a consensus constraint, ADMM is utilized to divide problem Ω into two subproblemsΩ0 andΩ1, i.e.,
Ω0 : min w0
L0(w0,w1,λ1) := (1− α)JC(w0) + 〈λ1,w1 −w0〉+ ρ
2 ||w1 −w0||
2,and
Ω1 : min w1
L1(w0,w1,λ1) := αJK(w1) + 〈λ1,w1 −w0〉+ ρ
2 ||w1 −w0||
2,
whereλ1 denotes a Lagrange vector, andρ is a penalty hyperparameter. SubproblemΩ0 and subproblemΩ1 are then iteratively solved on the server and on the edge device, respectively. More
specifically, in thet-th iteration of Co-KD, thet-th iterate ofw1, denoted asw1(t), is determined on the edge device byw1(t) := argminw1 L1(w0(t− 1),w1,λ1(t− 1)), wherew0(t − 1) denotes the (t − 1)-th iteration ofw0, andλ1(t − 1) denotes the(t − 1)-th iteration ofλ1. Afterwards, the edge device uploads its Lagrange vectorλ1(t − 1) andw1(t) to the server, and then the server determinesw0(t) := argminw0 L0(w0,w1(t),λ1(t− 1)). Afterwards, the server sendsw0(t) to the edge device, and the edge device determinesλ1(t) := λ1(t − 1) + ρ(w1(t) − w0(t)). Since bothw0(t) andw1(t) do not have closed-form solutions, stochastic gradient descent (SGD) is utilized to numerically determinew0(t) andw1(t). As a convergence criterion is reached, e.g., ||w0(t) −w0(t − 1)|| ≤ ǫ0 whereǫ0 is a predefined threshold, Co-KD stops.w0(t) is regarded as the parameters of the global compact model that is then deploy d to the edge device.
Since a solution obtained by ADMM satisfies Karush–Kuhn–Tucker (KKT) conditions,w0(t) can converge to a local optimum of ProblemΩ. Moreover, the compact model is pretrained over the source data. It can avoid bad local optimum of poor generalization during training ofw0, because pretraining over the source data can be regarded as a regularization on the parameterw0, as studied in (Erhan et al., 2010). As a result, the global model obtained by Co-KD can always converge to a local optimum with sufficient generalization performance, i.e., it can achieve high inference accuracy over both the target data and the source data.

heading: 3.2.2 CASE II: MULTIPLE EDGE DEVICES
text: When a single edge device only has a limited number of target data, it is hard to learn the finegrained representations of the target data even by the largemod l. As a result, the compact model trained by Co-KD still suffers low inference accuracy on thetarget data. To tackle this challenging, Co-KD can be extended to unify the target data of the multipleedge devices for training a global compact model. Due to page limit, the details of Co-KD for multiple edge devices are provided in Appendix B, and only the main idea is provided below.
The training process of the global compact model is first formulated into a consensus problemΩ′ over theM edge device and the server. ProblemΩ′ is then divided into(M + 1) subproblems by following ADMM, and each edge device as well as the server iteratively solves one subproblem. In the t-th iteration, edge devicem determineswm(t) and then sendswm(t) andλm(t − 1) to the server. Afterwards, the server determinesw0(t) and then broadcastsw0(t) to all the edge devices. Edge devicem then computesλm(t). As the convergence criterion is reached, the compact model with parametersw0(t) is broadcast to all the edge devices as the new compact model.

heading: 3.3 DISCUSSION ONPRIVACY PRESERVATION
text: In STU-KD, the edge device does not upload the target data to the centralized server, providing a basic level of privacy preservation. However, recent studies (Melis et al., 2019; Zhu et al., 2019; Zhao et al., 2020) show that it is possible for the server to rec v r some information of the target data from the uploaded parameters. This indicates that STU-KD isvulnerable to such recovery attacks.
One advantage of STU-KD is that it can be easily integrated with secure aggregation (Bonawitz et al., 2017) to resist the recovery attacks. In the case of multiple devices, the server only requires ∑M m=1 w ′ m(t) to updatew0(t − 1) (see equation (5) in Appendix B), wherew ′
m(t) denotes the parameters uploaded by edge devicem. To prevent the server from obtainingw′m(t) without affecting the update process ofw0(t − 1), theM edge devices can add noise to their parameters before uploading, based on the following rule. Forany two edge devicesi andj (i < j) from theM edge devices, edge devicei adds a random vectorni,j tow′i(t), and edge devicej subtractsni,j from w′j(t) (how to generateni,j can refer to Appendix C). By adding and subtracting these random vectors, the disturbed parametersŵm(t) = w′m(t) + ∑ m<j nm,j − ∑ i<m ni,m are obtained on edge devicem. ŵm(t) is then uploaded to the server. The server cannot recoverw′m(t) from ŵm(t), but it can compute ∑M m=1 w ′ m(t) = ∑M m=1 ŵm(t), and then updatesw0(t− 1).
In the case of a single edge device, the edge device can be regarded s two virtual devices by dividing the target data into two parts. Each virtual deviceus s one part to train a compact model. Afterwards, the above secure aggregation mechanism for multiple devices can be applied directly.

heading: 4 EXPERIMENT

heading: 4.1 EXPERIMENTAL SETUP
text: Three public datasets for domain adaptation are used to evaluate STU-KD.Office-31(Saenko et al., 2010) has three domains: Amazon (A), DSLR (D), and Webcam (W). ImageCLEF-DA has three domains: ImageNet ILSVRC 2012 (I ), Caltech-256 (C), and Pascal VOC 2012 (P).Office-Home (Venkateswara et al., 2017) has four domains: Artistic images (Ar ), Clip Art (Cl), Product images (Pr), and Real-World images (Rw).
The following schemes are compared with STU-KD: the compactmodel without domain adaptation, source hypothesis transfer (SHOT) (Liang et al., 2020), and target-only knowledge distillation (TO-KD). Here TO-KD is constructed by replacing collaborative knowledge distillation (Co-KD) in STU-KD with knowledge distillation (Hinton et al., 2015). In addition, four variants of STUKD are also evaluated. The first three variants are constructed by replacing the ADMM based learning algorithm in STU-KD with federated average (McMahan et al., 2017) (denoted as STUKD-AVG), distributed SGD (Zinkevich et al., 2010) (denotedas STU-KD-SGD), and SCAFFOLD (Karimireddy et al., 2020) (denoted as STU-KD-SCAFFOLD). The forth variant is formed by integrating STU-KD with secure aggregation (Bonawitz et al.,2017) (denoted as STU-KD-SA). For SHOT, its compact model is directly trained over the target data. For STU-KD and its variants, a large target model is obtained by adapting a large source model t the target data, and the large target model is used to generate soft labels for the target data. Thedifference is that the compact model of TO-KD is only trained over the target data and the soft labels, while the compact models of the rest schemes are trained over the target data, the soft labels, and the source data.
For each transfer task,90% of the source data are used for training; the rest10% of the source data are used as the test data. The inference accuracy on the source data is determined by testing each scheme on the test data. All the target data are used for training a d test, as is done in (Long et al., 2017; 2018; Liang et al., 2020). The inference accuracy on the target data is determined by testing each scheme on all the target data. In all the experiments, ResNet-18 (He et al., 2016) is selected as the the compact model. ResNet-50 is selected as the large source model for STU-KD. The compact model and the large source model are trained through backpropagation and mini-batch SGD with momentum0.9 and weight decay1e−3 is adopted as the optimizer. The learning rate is set to η0 = 0.001. The batch size on the edge device is set to8, and the batch size in the server is set to 32. Moreover, for all the experiments, the balancing hyperparameterα is set to0.8; the penalty hyperparametersρ is set to0.3. STU-KD is implemented based on a library called FedML (He etal.,
2020b) in PyTorch (Paszke et al., 2019). Each experiment is executed for three times, and the mean of the inference accuracy is reported. All the experiments are executed on a server with one i910900k CPU, one GeForce RTX 3090 GPU, and 64 GB RAM. More details of the experimental setup are given in Appendix D.

heading: 4.2 EXPERIMENTAL RESULT
text: The main results of a single edge device are shown in this subsection. Due to page limit, the results of different hyperparameter settings, different network architectures, and the visualization results are provided in Appendix E.1, Appendix E.2, and Appendix E.3, respectively. The results of multiple edge devices are provided in Appendix E.4. The results on Office-31 dataset, ImageCLEF-DA dataset, and Office-Home dataset are reported in Table 1, Table 2, and Table 3, respectively. In each table, the upper half records the inference accuracy of eachscheme on the all the target data; the bottom half records the inference accuracy of each scheme ont st data. The method ResNet-18 represents the compact model without domain adaptation. Firstly, the inference accuracy on the target data is evaluated. TO-KD and STU-KD (our schemes) outperform all the existing schemes on all the transfer tasks. Compared with the state-of-the-artscheme SHOT, our schemes improve the inference accuracy on the target data by7.7%, 3.3%, and14.7% on Office-31 dataset, ImageCLEF dataset, and Office-Home dataset, respectively. The gain ofour schemes can be explained as follows. The large target model used in STU-KD can learn more fine-grained representations of the target data than the compact model. Thus, the compact model that inherits the knowledge from the large target model can achieve higher inference accuracy on the targ t data, compared with the compact model that are directly trained using the target data. Sinceboth STU-KD and TO-KD utilize the large target model, their performance is comparable. When considering the inference accuracy on the source data, the performance of SHOT and TO-KD degradedue to catastrophic forgetting, compared with the compact model without adaptation. For STU-KD, its performance is much closer to that of the compact model without adaptation, indicatinghat STU-KD indeed alleviates or even prevents catastrophic forgetting catastrophic forgetting. Note that in some cases (e.g., D→A, C→I, and Ar→Cl), STU-KD achieves higher inference accuracy on the source data than the compact model without adaptation. The reason can be that by trainingthe compact model over both the target data and the source data, the compact model learns more general representations of the source data, which boosts its inference accuracy on the source data.
The peak memory footprint of each scheme is shown in Figure 3,where SHOT (ResNet-50) represents training the ResNet-50 model on the edge device via SHOT. The memory footprint is determined when the batch size is set to8. For both TO-KD and STU-KD, their peak memory footprints are caused by training the large model via lite residual hypothesis transfer (LRHT). Compared to the case of directly training the large source model on the edge device, STU-KD reduces the memory footprint by nearly67%, indicating that LRHT indeed saves the memory space and enables the adaptation of the large model on the edge device.
Three variants of STU-KD (STU-KD-AVG, STU-KD-SGD, and STU-KD-SCAFFOLD) are evaluated on Office-31 dataset, and the results are shown in Table 4. Although STU-KD-AVG is comparable to STU-KD in terms of inference accuracy on the target data, it suffers poor inference accuracy on the source data. As shown in Figure 4, STU-KD-AVG only converges to a point with a low inference accuracy on the source data. STU-KD-SCAFFOLD employs an advanced FL algorithm (i.e., SCAFFOLD (Karimireddy et al., 2020)), so it achieves higherinference accuracy on the source data compared to STU-KD-AVG. However, such a inference accuracyis still lower than that of STU-KD,
Table 4: Inference accuracy (%) onOffice-31dataset for the variants of STU-KD (target data and source data)
Method A→D A→W D→A D→W W→A W→D Avg.
STU-KD-AVG 93.6± 0.5 89.1± 0.2 73.5± 0.1 97.3± 0.1 74.2± 0.3 98.7± 0.3 87.7 STU-KD-SGD 94.0± 0.1 88.9± 0.1 67.6± 8.3 97.7± 0.2 73.8± 0.3 99.1± 0.1 86.8 STU-KD-SCAFFOLD 93.8± 0.6 89.1± 0.2 73.8± 0.1 97.4± 0.3 74.3± 0.1 98.5± 0.2 87.8 STU-KD 94.0± 0.2 88.8± 0.1 72.3± 0.3 97.5± 0.1 74.6± 0.1 99.1± 0.1 87.7 STU-KD-SA 94.2± 0.2 88.8± 0.4 72.2± 0.5 97.4± 0.5 74.6± 0.2 99.0± 0.2 87.7
STU-KD-AVG 41.0± 0.4 50.8± 0.9 67.3± 2.1 100.0± 0.0 56.7± 0.7 86.7± 6.4 67.1 STU-KD-SGD 68.4± 8.1 83.1± 0.9 92.7± 1.1 100.0± 0.0 92.5± 2.0 97.5± 1.2 89.0 STU-KD-SCAFFOLD 62.9± 8.3 81.8± 1.4 88.0± 4.0 100.0± 0.0 82.5± 8.0 92.1± 7.7 84.6 STU-KD 81.1± 0.8 83.5± 0.8 97.3± 1.1 100.0± 0.0 94.2± 1.8 98.8± 0.0 92.5 STU-KD-SA 84.3± 0.2 83.7± 0.3 97.3± 1.1 100.0± 0.0 92.1± 0.7 98.8± 0.0 92.7
Figure 3: The peak memory footprint of STU-KD and the comparison methods.
Figure 4: The inference accuracy curves of STU-KD and its variants onW → A task.
Figure 5: The communication cost of STU-KD and its variants onW → A task.
and the gap is nearly8.1% in average. Moreover, the communication cost of STU-KD-SCAFFOLD is twice that of STU-KD, since SCAFFOLD requires uploading both local model updates and a set of control parameters. The comparison on these FL based variants verify that FL is not effective in STU-KD to consolidate the compact models from the edge device and the server. As for STU-KDSGD, its performance is the closest to that of STU-KD among these three variants. However, its communication cost is too large. As shown in Figure 5, the communication cost of STU-KD-SGD is nearly100 times that of STU-KD. The reason is that the edge device needsto frequently sends its local gradients to the server to compute the gradients of theobj ctive function, which leads to the large communication cost. The results of STU-KD-SA is also reported in Table 4. Compared with STU-KD, STU-KD-SA achieves the comparable inference accura y on both the source data and the target data, verifying that secure aggregation does not affect the performance of STU-KD. Note that the slightly better performance of STU-KD-SA than STU-KD isdue to the random seeds used in the experiments, and it can disappear when more experimentsare conducted.

heading: 5 CONCLUSION
text: In this paper, a source-target unified knowledge distillation (STU-KD) scheme was developed to adapt the compact model on the edge device to the target data while protecting users’ privacy. In the scheme, a large source model was first loaded and adapted to the target data on the edge device, and a large target model was obtained. Such an adaptation wasenabled by a memory efficient lite residual hypothesis transfer algorithm. A collaborative knowledge distillation method that unifies the target data and the source data for training the compact model was developed to transfer the knowledge of the large target model to the compact model while preventing catastrophic forgetting. Secure aggregation was also employed in STU-KD to enhance privacy preservation. Extensive experiments showed that STU-KD can achieve the highest infere c accuracy over the target data while maintaining the original high inference accuracy over the source data, compared with stateof-the-art schemes. STU-KD is now evaluated on public datasets. How to incorporate STU-KD into a real-world application can be an interesting topic for future work.

heading: A L OSS FUNCTION OFSHOT
text: Define the parameters of the feature extractor aswe and the loss function of SHOT asJS . JS consists of two terms: an information maximization (IM) loss (Krauseet al., 2010; Hu et al., 2017) denoted asJI and a cross-entropy (CE) loss of self-supervised learning de oted asJp. The expression ofJS can be written as
JS(we) = JI(we) + βJp(we), (2)
whereβ is a balancing hyperparameter. Define the target output of a target samplext asfT (xt). IM lossJI captures the uncertainty of the target outputs of the targetdata and the diversity of those target outputs by
JI(we) = −Ext [
K∑
k=1
δk(fT (xt)) log δk(fT (xt))] +
K∑
k=1
p̂k log p̂k,
whereK is the number of classes in the classification task,δk(a) = exp(ak)∑ i exp(ai) is thek-th element in the softmax output of a vectora, and p̂k is the k-th element in the mean output embedding p̂ = Ext [δ(fT (xt))] of the target data. CE lossJp is written as
Jp(we) = Ext [L(fT (xt), ŷt)],
whereŷt is a pseudo-label generated by a clustering based pseudo-labeling mechanism (Caron et al., 2018; Liang et al., 2020) forxt.

heading: B CO-KD FOR MULTIPLE EDGE DEVICES
text: When a single edge device only has limited number of target data, it is hard to learn the fine-grained representations of the target data even for the large model.As a result, the compact model trained by Co-KD still suffers low inference accuracy on the target data. To tackle this challenging, Co-KD can be extended to unify the target data of the multiple edge devices for training a global compact model. TheM edge devices collaborate with the cloud server to train a global compact model by solving
min w
1
M
M∑
m=1
αJ (m) K (w) + (1− α)JC(w), (3)
wherew denotes the parameters of the global compact model. The training process of the global compact model is formulated into a consensus problemΩ′ over theM edge device and the server.
min w0,w1,...,wM
1
M
M∑
m=1
αJ (m) K (wm) + (1− α)JC(w0)
s.t.wm = w0, m = 1, ...,M,
(4)
whereJ (m)K is the local loss function of KD on edge devicem, wm is the compact model on edge devicem, andw0 is the compact model updated in the centralized server. Co-KD is then extended to iteratively solve problem (4) as follows. In thet-the iteration of Co-KD, edge devicem determines wm(t) by
wm(t) = argmin wm
α
M J (m) K (wm) + 〈λm(t− 1),wm −w0(t− 1)〉+
ρ 2 ||wm −w0(t− 1)|| 2,
and sends its Lagrange vectorλm(t− 1) andwm(t) to the cloud server. Next, the server updates its compact model by
w0(t) = argmin w0 (1− α)JC(w0) +
M∑
m=1
〈λm(t− 1),wm(t)−w0〉+ ρ
2 ||wm(t)−w0||
2.
The updated compact model is then sent to theM dge devices. Afterwards, edge devicem updates λm(t− 1) via
λm(t) := λm(t− 1) + ρ(wm(t)−w0(t))
As the convergence criterion is reached. Co-KD stops. The compact model withw0(t) is then broadcast to theM edge devices as their new compact models.
Similarly to the case of single edge device, SGD is utilized to eterminewm(t) on edge devicem and to determinew0(t) on the server. More specifically, the gradientgm(wm(t−1)) used to update wm(t− 1) is computed by
gm(wm(t− 1)) = α
M ∇wmJ
(m) K (wm(t− 1)) + λm(t− 1) + ρ(wm(t− 1)−w0(t− 1)).
Letw′m(t) = λm(t−1)+ρwm(t) denote the parameters uploaded by edge devicem. The gradient g0(wm(t− 1)) used to updatew0(t− 1) can then be represented as
g0(w0(t− 1)) = ∇w0JC(w0(t− 1))− M∑
m=1
(λm(t− 1) + ρ(wm(t)−w0(t− 1)))
= ∇w0JC(w0(t− 1)) +Mρw0(t− 1)− M∑
m=1
w′m(t).
(5)

heading: C ADDITIONAL INFORMATION ABOUT SECUREAGGREGATION
text: To generate the same random vector on edge devicei and edge devicej, the same random seed is used in the random number generator on each edge device. Suh a random seed can be shared among the two users using Diffie–Hellman key exchange while ke p the server unknown to the random seed.

heading: D MORE DETAILS OF EXPERIMENTAL SETUP
text: Three public datasets for domain adaptation are used to evaluate STU-KD.Office-31(Saenko et al., 2010) is a widely used dataset for domain adaptation. It has4, 652 images and31 classes collected from three domains: Amazon (A), DSLR (D), and Webcam (W). ImageCLEF-DA has three domains: ImageNet ILSVRC 2012 (I ), Caltech-256 (C), and Pascal VOC 2012 (P). For each domain, there are12 classes and50 images in each class.Office-Home (Venkateswara et al., 2017) has15, 500 images and65 classes from four distinct domains: Artistic images (Ar ), Clip Art (Cl), Product images (Pr), and Real-World images (Rw).
For STU-KD, a variant of ResNet-50 (He et al., 2016) is selectd as the large model. Concretely, a standard ResNet-50 model pretrained over ImageNet (Russakovsky et al., 2015) is employed as the base module. Its original fully connected (FC) layer is then replaced with a bottleneck layer and a task-specific FC layer. Moreover, a batch normalization (BN) layer (Ioffe & Szegedy, 2015) is put after the FC layer in the bottleneck layer, and a weight normalization (WN) layer (Salimans & Kingma, 2016) is put after the task-specific FC layer (Liang et al., 2020). Following the same procedure, a variant of ResNet-18 is constructed and is selected as the compact model for all the schemes. Moreover, four lite residual modules are added to the ResNet-50 model for lite residual hypothesis transfer. The architecture of theResNet-50 model with four lite residual modules is shown in Figure 6.
The compact model and the large source model are trained throug backpropagation, and minibatch SGD with momentum0.9 and weight decay1e−3 is adopted as the optimizer. The learning rate is set toη0 = 0.001 for the lite residual modules in the large model and the featur extractor in the compact model. A widely used learning rate scheduler (Ganin & Lempitsky, 2015; Long et al., 2017; 2018) is also adopted, i.e.,η = η0(1 + 10p)−0.75, wherep is the training progress changing from 0 to 1.

heading: E MORE EXPERIMENTAL RESULTS
text: E.1 HYPERPARAMETERANALYSIS
To analyze the effects of hyperparameters on STU-KD, the expriment of STU-KD on Office-31 (Table 1) data is repeated with different settings of the hyperparametersα andρ. The results are shown in Table 5. In terms of the inference accuracy on the targe data, the overall performance of STU-KD with different settings are comparable, indicating that the inference accuracy on the target data is not much sensitive to the settings of the hyperparameters. Now consider the inference accuracy on the source data. It is obvious that smallerα l ads to higher inference accuracy on the source data. The reason is as follows. The smaller is the valuofα, the larger is the weight of the CE loss in equation (1). As a result, the global compact modelne ds to achieve a smaller error over the source data after the adaptation, which can lead to higher inference accuracy on the source data. As for ρ, the results show that the larger is the value ofρ, the large is the inference accuracy on the source data, though the effect ofρ is not as strong as that ofα. The above results indicate that the inference accuracy on the source data is more sensitive to the setting of the hyperparameters than the inference accuracy on the target data. Properly reducing the value ofα or increasing the value of ρ can achieve better inference accuracy on the source data.
E.2 DIFFERENTNETWORK ARCHITECTURES
To show that STU-KD is applicable to different network architectures, two additional experiments are conducted. In the first experiment, Proxyless-Mobile (Cai et al., 2019) is selected as the compact model, and ResNet-50 (He et al., 2016) is selected as the largsource model. The results are shown in Table 6. In the second group, ResNet-34 is selected as the compact model, and the ResNet-101 is selected as the large source model. The results are shown in Table 7.
According to Table 6 and Table 7, TO-KD and STU-KD (our schemes) outperforms other schemes in terms of the inference accuracy on the target data, and these two schemes achieve comparable inference accuracy on the target data. Moreover, STU-KD outperforms TO-KD in terms of the inference accuracy on the source data. Such results are consistent with that of Table 1, Table 2, and Table 3, indicating that STU-KD is applicable to different network architectures for both the compact model and the large source model.
In the case of a single device, the representations learn by ResNet-18, SHOT, TO-KD, and STUKD are visualized by t-SNE (Van der Maaten & Hinton, 2008), asshown in Figure 7. As expected, the representations learned by STU-KD are the most discriminative for both target domain and the source domain. Although TO-KD achieves comparable inference accuracy on the target data to that achieved by STU-KD, the target representations learned by TO-KD is less discriminative than that of STU-KD. In terms of the source representations, SHOT outperforms TO-KD, which verifies the results in Table 1. Such a result also indicates that directly using KD can lead severe catastrophic forgetting.
E.4 RESULTS OFMULTIPLE EDGE DEVICES
Two group of experiments are conducted in the case of multiple edge devices. First, the target data on the multiple edge devices come from the same target domain. Second, the target data on the multiple edge devices come from different target domains. More specifically, in the second experiments, the
edge devices are divided into two groups. The devices from the same group have the target data from the same domain, while the target domains of the two groups are different.
In the first group of experiments,M is set to5, and the target data are randomly partitioned into5 parts each of which is held by one edge device. Three schemes are v luated on Office-31 dataset: SHOT (Liang et al., 2020), TO-KD, and STU-KD. For SHOT and TO-KD, each edge device trains its local compact model independently, without interactions with the centralized server. For STUKD, five edge devices and the server collaboratively train a global compact model. The experimental results onA → W task and that onA → D are reported in Table 8 and Table 9, respectively.
As shown in Table 8 and Table 9, TO-KD outperforms SHOT in terms of the inference accuracy on the target data. The gain comes from the large model capacity of large model. By unifying the target data from all the edge devices, STU-KD further improves the average inference accuracy on the target data by6.58% on A → W task, and by6.43% on A → D task. Besides, STU-KD leverages the source data to train the global compact model,and thus achieves the highest inference accuracy on the source data. After employing secure aggregation, STU-KD-SA achieves comparable performance to that of STU-KD, indicating that secure aggreation does not affect the performance of STU-KD.
In the second group of experiments, two transfer tasks are conducted: 1)A from Office-31 dataset is selected as the source domain, andD andW from Office-31 are selected as the target domains, denoted asA → DW; 2) I from ImageCLEF-DA dataset is selected as the source domain,andC and P from ImageCLEF-DA are selected as the target domains, denoted asI → CP. In taskA → DW, the two target domains are similar; while in taskI → CP, the two target domain are much different from each other. The number of edge deviceM is set to6. For each transfer task, the first three edge devices come from the one target domain, and the rest three edge devices come from another target domain. Two schemes are used to trained the compact models onthe edge devices. First, STU-KD is only applied to the edge devices from the same target domain, denoted as STU-KD (separated). Second, STU-KD is applied to all thee six edge devices to learn a global compact model, denoted as STU-KD (joint). The results of taskA → DW and taskI → CP are shown in Table 10 and Table 11, respectively.
Comparing the performance of these two schemes, it is indicated that applying STU-KD to the case of different target domains does not highly degrade the performance of the obtained compact model. That is, STU-KD is applicable to the case of different targetdomains, though STU-KD does not provide extra performance gain for the compact model by involving more data into the training process.

Review summaries:
This paper considers a new and complex setting involving domain adaptation, federated learning, and knowledge distillation: Under the premise of protecting privacy, one needs to deploy a compact model from a source central server to target client devices and requires the model to learn new knowledge with target unlabeled client data while remembering knowledge of source data on the central server. Generally, the authors propose a source-target unified knowledge distillation scheme. Within this scheme, the authors propose solutions to tackle corresponding difficulties with this setting. Specifically, to avoid the low inference accuracy due to low model capacity, a large source central model is adapted to target clients with SHOT (ICML'20). A lite-residual hypothesis transfer method is proposed to keep memory-efficient adaptation on target clients. A collaborative knowledge distillation method is proposed to defy catastrophic forgetting of source knowledge. To protect the privacy of target clients, a secure aggregation method is used. The authors validate the effectiveness of the proposed scheme on three domain adaptation datasets. 

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.