S OURCE -T ARGET U NIFIED K NOWLEDGE D ISTILLATION FOR M EMORY -E FFICIENT F EDERATED D OMAIN A DAPTATION ON E DGE D EVICES
Anonymous authors
Paper under double-blind review A BSTRACT
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 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 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 sufﬁcient inference accuracy due to its low model capacity.
To this end, a scheme called source-target uniﬁed knowledge distillation (STU-KD) is developed in this paper.
It ﬁrst adapts a large model to the target data on the edge device, and then the knowledge of the large model is transferred to the compact model via knowledge distillation.
Since training the large model leads to large memory consumption, a domain adaptation method called lite residual hypothesis transfer is designed to achieve memory-efﬁcient adaptation to the target data on the edge device.
Moreover, to prevent the compact model from forgetting the knowledge of the source data during knowledge distillation, a collaborative 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.
STUKD can be easily integrated with secure aggregation so that the server cannot obtain the true model parameters of the compact model.
Extensive experiments conducted upon several tasks of object recognition show that STU-KD can improve the inference accuracy by up to 14 .
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.
I NTRODUCTION
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 space 1 , 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 architecture search (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.
Consequently, the compact model suffers low inference accuracy on the target data.
To ensure high
[Figure]
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 in Kang et al.
; Tang & Jia (2020);
Xu et al.
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 size 48 consumes nearly 172 GFLOPs and 1 . 1 GB memory space per batch, which is affordable for
Jetson Nano.
However, this type of approach cannot achieve high inference accuracy on the target data due to limited model capacity of the compact model.
To this end, a scheme called source-target uniﬁed knowledge distillation (STU-KD) is developed in this paper.
The key idea is to utilize a large model with sufﬁcient model capacity to learn ﬁnegrained representations of the target data, and then transfer the knowledge of the large model to the compact model.
As shown in step 1 in Figure 1, a large source model is loaded on the edge device and is then adapted to the target data to obtain a large target model.
The challenges of this step are two-fold.
First, the target data are unlabeled, so the ﬁne-grained representations cannot be learned via supervised learning.
As a result, unsupervised domain adaptation (UDA) is needed in step 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 device due to conﬁdentiality 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 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.
For example, if ResNet-50 is retrained with batch size 48 on Jetson Nano, the computational cost is nearly 364 GFLOPs, which is affordable for Jetson
Nano.
However, the memory consumption per batch is nearly 5 GB, exceeding the 4 GB memory space limit of Jetson Nano.
To tackle the challenges in step 1 of STU-KD, a memory-efﬁcient UDA method called lite residual hypothesis transfer (LRHT) is designed by enhancing the architecture of a source-free UDA method such as SHOT.
More speciﬁcally, the same loss function as that of SHOT is adopted, so unsupervised learning can be conducted with unlabeled target data.
However, the feature extractor in SHOT must be replaced with a new architecture that can be trained in a memory-efﬁcient 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 ﬁne-tuned by training the LR modules only while keeping the parameters of the feature extractor ﬁxed.
As a result, in LRHT the training process of the large source model involves neither the feature extractor nor the classiﬁer.
Since training the LR modules generates a much smaller volume of activations and demands much lower memory footprint, the training process of the large target model is highly memory-efﬁcient.
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-efﬁcient 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 accuracy.
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 edge device, 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 challenge here is how to consolidate these 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 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.
; Karimireddy et al.
(2020)) are not 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 the compact model on the edge device and that on the server, as shown in step 2 in Figure 1.
More speciﬁcally, the training process of the global compact model is formulated into a consensus problem over 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 related 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 on the target data by up to 14 .
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.
R ELATED W ORK
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 the target data and the source data are stored in one place, which are not applicable when the target data 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 model to the target data.
TinyTL (Cai et al., 2020) assumes that the on-device model has sufﬁcient 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 of the compact model.
However, due to its low model capacity, directly training the compact model cannot achieves sufﬁcient 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 by different 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
[Figure]
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 sufﬁcient accuracy, both FDA and SFDA schemes require the source model to have sufﬁcient model capacity.
However, such a requirement cannot be satisﬁed by a compact model with low model capacity.
M ETHODOLOGY
One centralized server and M edge devices ( M ≥ 1 ) are considered in this paper.
The labeled source data in the server are denoted as { x s , y s } where x s represents the feature vector of a source sample and y s the ground-truth label of the source sample.
The source data are sampled from the source domain denoted as P s .
The unlabeled target data on the edge devices are denoted as { x t } where x t represents the feature vector of a target sample.
The target data are sampled from the target domain denoted as P t .
Each edge device needs to accomplish a classiﬁcation task with K classes using a compact model.
The compact model on the edge device m ∈ { 1 , ..., M } is deﬁned as a function f m , and its model parameters are denoted as w m .
Given the above setup, source-target uniﬁed knowledge distillation (STU-KD) is developed to adapt the compact model to the target data.
3. L ITE R ESIDUAL H YPOTHESIS T RANSFER
To boost the inference accuracy of the compact model, a large source model with sufﬁcient model capacity, deﬁned as a function f T , is adopted in STU-KD.
It can learn more ﬁne-grained representations of the target data compared with the compact model.
Before being deployed to the M edge devices, the large source model f T is pretrained over the source data by minimizing an objective function J T ( w T ) =
E ( x s ,y s )
[ L ( f T ( x s ; w T ) , y s )] , where w T denotes the parameters of the large source model, and L ( · , · ) is the standard cross-entropy loss.
On edge device m ∈ { 1 , ..., M } , the large source model is then adapted to the target data.
The challenges of the adaptation are twofold.
First, the target data are unlabeled so that ﬁne-grained representations of the target data cannot be learned via supervised learning.
As a result, unsupervised domain adaptation (UDA) is needed.
Second, edge device m 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-efﬁcient 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 as J S 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-efﬁcient 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 ﬁne-tuned by the LR modules.
The LR modules ﬁne-tune the outputs of the feature extractor by changing the intermediate activations of the feature extractor.
Deﬁne the i -th building block of the feature extractor as a function g i and the LR module added to the it as a function r i .
Given the activations a i − 1 of the ( i − 1) -th building block, LR module r i changes the activations of the i -th building block as a i
=
g
i ( a i − 1 ) + r i ( a i − 1 ) , where g i ( a i − 1 ) denotes the original activations of the i -th building block, and r i ( a i − 1 ) denotes the residual activations of LR module r i .
During the adaptation, the parameters of the feature extractor and that of the classiﬁer are ﬁxed, and only the LR modules are trained to ﬁne-tune the outputs of the feature extractor.
Compared to the case of directly training the feature 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-efﬁcient manner.
3. C OLLABORATIVE K NOWLEDGE D ISTILLATION
The next step is to transfer the knowledge of the large target model to the compact model.
3.2. C ASE I: SINGLE EDGE DEVICE
Consider the case of a single device that contain a compact model f 1 with the model parameters w 1 .
On the edge device, the large target model is leveraged to generate 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 speciﬁcally, 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 the source data need to be uniﬁed 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 as J K , is written as
[Equation]
where z c = f 1 ( x t ; w 1 ) is the logit of the compact model, z T = f T ( x t ; w ′ T ) is the logit of the large target model with parameters w ′ T (i.e., the soft label), and p ( 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 as J C .
The expression of J C is
[Equation]
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 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 the difference, existing FL algorithms (e.g., McMahan et al.
; Karimireddy et al.
; 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:
[Equation]
where α is a balancing hyperparameter, and w 0 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.,
[Equation]
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 speciﬁcally, in the t -th iteration of Co-KD, the t -th iterate of w 1 , denoted as w 1 ( t ) , is determined on the edge device by w 1 ( t ) :
= arg min w 1 L 1 ( w 0 ( t − 1) , w 1 , λ 1 ( t − 1)) , where w 0 ( t − 1) denotes the ( t − 1) -th iteration of w 0 , and λ 1 ( t − 1) denotes the ( t − 1) -th iteration of λ 1 .
Afterwards, the edge device uploads its Lagrange vector λ 1 ( t − 1) and w 1 ( t ) to the server, and then the server determines w 0 ( t ) :=
arg min w 0 L 0 ( w 0 , w 1 ( t ) , λ 1 ( t − 1)) .
Afterwards, the server sends w 0 ( t ) to the edge device, and the edge device determines λ 1 ( t ) :
=
λ 1 ( t − 1)
+ ρ ( w 1 ( t ) − w 0 ( t )) .
Since both w 0 ( t ) and w 1 ( t ) do not have closed-form solutions, stochastic gradient descent (SGD) is utilized to numerically determine w 0 ( t ) and w 1 ( t ) .
As a convergence criterion is reached, e.g., || w 0 ( t ) − w 0 ( t − 1) || ≤ ǫ 0 where ǫ 0 is a predeﬁned threshold, Co-KD stops.
The compact model with w 0 ( t ) is regarded as the global compact model, and then is deployed to the edge device.
Since a solution obtained by ADMM satisﬁes Karush–Kuhn–Tucker (KKT) conditions, w 0 ( 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 of w 0 , because pretraining over the source data can be regarded as a regularization on the parameter w 0 , 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 sufﬁcient generalization performance, i.e., it can achieve high inference accuracy over both the target data and the source data.
3.2. C ASE II: MULTIPLE EDGE DEVICES
When a single edge device only has a limited number of target data, it is hard to learn the ﬁne-grained representations of the target data even by 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 globally shared 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 shared compact model is ﬁrst formulated into a consensus problem Ω ′ over the M 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 device m determines w m ( t ) and then sends w m ( t ) and λ m ( t − 1) to the server.
Afterwards, the server determines w 0 ( t ) and then broadcasts w 0 ( t ) to all the edge devices.
Edge device m then computes λ m ( t ) .
As the convergence criterion is reached, the compact model with w 0 ( t ) is broadcast to all the edge devices as the new compact model.
3. D ISCUSSION ON P RIVACY P RESERVATION
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 recover some information of the target data from the uploaded parameters.
This indicates that STU-KD is vulnerable 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 P Mm =1 w ′ m ( t ) to update w 0 ( t − 1) (see equation (5) in Appendix B), where w ′ m ( t ) denotes the parameters uploaded by edge device m .
To prevent the server from obtaining w ′ m ( t ) without affecting the update process of w 0 ( t − 1) , the M edge devices can add noise to their parameters before uploading, based on the following rule.
For any two edge devices i and j ( i < j ) from the M edge devices, edge device i adds a random vector n i,j to w ′
i ( t ) , and edge device j subtracts n i,j from w ′ j ( t ) (how to generate n i,j can refer to Appendix C).
By adding and subtracting these random vectors, the disturbed parameters ˆ w m ( t ) =
w ′ m ( t ) + P m<
j n m,j − P l<m n l,m are obtained on edge device m .
ˆ w m ( t ) is then uploaded to the server.
The server cannot recover w ′ m ( t ) from ˆ w m ( t ) , but it can compute P Mm =1 w ′ m ( t ) =
P Mm =
ˆ w m ( t ) , and then updates w 0 ( t − 1) .
In the case of a single edge device, the edge device can be regarded as two virtual devices by dividing the target data into two parts.
Each virtual device uses one part to train a compact model.
Afterwards, the above secure aggregation mechanism for multiple devices can be applied directly.
[Table] [Table]
E XPERIMENT 4.
[Table]
Three public datasets for domain adaptation are used to evaluate STU-KD.
Ofﬁce-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).
Ofﬁce-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 compact model 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 ﬁrst 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) (denoted as 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 to the target data, and the large target model is used to generate soft labels for the target data.
The difference 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 rest 10% 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 and 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 momentum 0 .
9 and weight decay 1 e − 3 is adopted as the optimizer.
The learning rate is set to η 0 = 0 .
The batch size on the edge device is set to 8 , and the batch size in the server is set to 32 .
Moreover, for all the experiments, the balancing hyperparameter α is set to 0 .
8 ; the penalty hyperparameters ρ is set to 0 .
STU-KD is implemented based on a library called FedML
(He et al.,
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
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.
4. E XPERIMENTAL R ESULT
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 Ofﬁce-31 dataset, ImageCLEF-DA dataset, and Ofﬁce-Home dataset are reported in Table 1, Table 2, and Table 3, respectively.
In each table, the upper half records the inference accuracy of each scheme on the all the target data; the bottom half records the inference accuracy of each scheme on test data.
The method ResNetrepresents 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-art scheme SHOT, our schemes improve the inference accuracy on the target data by 7 . 7% , 3 . 3% , and 14 .
7% on Ofﬁce-31 dataset, ImageCLEF dataset, and Ofﬁce-Home dataset, respectively.
The gain of our schemes can be explained as follows.
The large target model used in STU-KD can learn more ﬁne-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 target data, compared with the compact model that are directly trained using the target data.
Since both 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 degrade due 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, indicating that 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 training the 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 to 8 .
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 nearly 67% , 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 Ofﬁce-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 higher inference accuracy on the source data compared to STU-KD-AVG.
However, such a inference accuracy is still lower than that of STU-KD,
[Table]
and the gap is nearly 8 .
1% in average.
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-KD-SGD, 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 nearly 100 times that of STU-KD.
The reason is as follows.
To compute one gradient of the objective function in problem Ω , the edge device are required to send a local gradient to the server.
To obtain a converged model, thousands of the gradients of the objective function need to be computed during the training process, 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 accuracy 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 is due to the random seeds used in the experiments, and it can disappear when more experiments are conducted.
C ONCLUSION
In this paper, a source-target uniﬁed 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 ﬁrst loaded and adapted to the target data on the edge device, and a large target model was obtained.
Such an adaptation was enabled by a memory efﬁcient lite residual hypothesis transfer algorithm.
A collaborative knowledge distillation method that uniﬁes 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 inference 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.
A L OSS FUNCTION OF SHOT
Deﬁne the parameters of the feature extractor as w e and the loss function of SHOT as J S .
J S consists of two terms: an information maximization (IM) loss (Krause et al., 2010; Hu et al., 2017) denoted as J I and a cross-entropy (CE) loss of self-supervised learning denoted as J p .
The expression of J S can be written as
[Equation]
where β is a balancing hyperparameter.
Deﬁne the target output of a target sample x t as f T ( x t ) .
IM loss J I captures the uncertainty of the target outputs of the target data and the diversity of those target outputs by
[Equation]
where K is the number of classes in the classiﬁcation task, δ k ( a )
= exp
[Equation]
a k ) P
i exp
[Equation]
a i ) is the k -th element in the softmax output of a vector a , and ˆ p k is the k -th element in the mean output embedding ˆ p =
E x t [ δ ( f T ( x t ))] of the target data.
CE loss J p is written as
[Equation]
where ˆ y t is a pseudo-label generated by a clustering based pseudo-labeling mechanism (Caron et al., 2018; Liang et al., 2020) for x t .
B C O -KD FOR M ULTIPLE E DGE D EVICES
When a single edge device only has limited number of target data, it is hard to learn the ﬁne-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 globally shared compact model.
The M edge devices collaborate with the cloud server to train a global compact model by solving
[Equation]
where w denotes the parameters of the global compact model.
The training process of the global compact model is formulated into a consensus problem Ω ′ over the M edge device and the server.
[Equation]
where J ( m ) K is the local loss function of KD on edge device m , w m is the compact model on edge device m , and w 0 is the compact model updated in the centralized server.
Co-KD is then extended to iteratively solve problem (4) as follows.
In the t -the iteration of Co-KD, edge device m determines w m ( t ) by
[Equation]
and sends its Lagrange vector λ m ( t − 1) and w m ( t ) to the cloud server.
Next, the server updates its compact model by
[Equation]
The updated compact model is then sent to the M edge devices.
Afterwards, edge device m updates λ m ( t − 1) via
[Equation]
0 ( t ))
×... ×...
×... ×... ..
L R
L R L R
L R ..
As the convergence criterion is reached.
Co-KD stops.
The compact model with w 0 ( t ) is then broadcast to the M edge devices as their new compact models.
Similarly to the case of single edge device, SGD is utilized to determine w m ( t ) on edge device m and to determine w 0 ( t ) on the server.
More speciﬁcally, the gradient g m ( w m ( t − 1)) used to update w m ( t − 1) is computed by
[Equation]
Let w ′ m ( t ) =
λ m ( t − 1)+ ρ w m ( t ) denote the parameters uploaded by edge device m .
The gradient g 0 ( w m ( t − 1)) used to update w 0 ( t − 1) can then be represented as
[Equation]
C A DDITIONAL I NFORMATION ABOUT S ECURE A GGREGATION
To generate the same random vector on edge device i and edge device j , the same random seed is used in the random number generator on each edge device.
Such a random seed can be shared among the two users using Difﬁe–Hellman key exchange while keep the server unknown to the random seed.
D M ORE D ETAILS OF E XPERIMENTAL S ETUP
Three public datasets for domain adaptation are used to evaluate STU-KD.
Ofﬁce-31 (Saenko et al., 2010) is a widely used dataset for domain adaptation.
It has 4 , 652 images and 31 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 are 12 classes and 50 images in each class.
Ofﬁce-Home (Venkateswara et al., 2017) has 15 , 500 images and 65 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 selected 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-speciﬁc 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-speciﬁc 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 the ResNet-50 model with four lite residual modules is shown in Figure 6.
[Table]
The compact model and the large source model are trained through backpropagation, and minibatch SGD with momentum 0 .
9 and weight decay 1 e − 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 feature 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 + 10 p ) − 0 .
75 , where p is the training progress changing from 0 to 1 .
E M ORE E XPERIMENTAL R ESULTS E. H YPERPARAMETER A NALYSIS
To analyze the effects of hyperparameters on STU-KD, the experiment of STU-KD on Ofﬁce- (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 target 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 α leads to higher inference accuracy on the source data.
The reason is as follows.
The smaller is the value of α , the larger is the weight of the
CE loss in equation (1).
As a result, the global compact model needs 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. D IFFERENT N ETWORK A RCHITECTURES
To show that STU-KD is applicable to different network architectures, two additional experiments are conducted.
In the ﬁrst experiment, Proxyless-Mobile (Cai et al., 2019) is selected as the compact model, and ResNet-50 (He et al., 2016) is selected as the large source 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.
[Table] [Table]
E. V ISUALIZATION
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), as shown 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 veriﬁes the results in Table 1.
Such a result also indicates that directly using KD can lead severe catastrophic forgetting.
E. R ESULTS OF M ULTIPLE E DGE D EVICES
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 speciﬁcally, in the second experiments, the {+[Table] TO-KD 77 .77 .77 .77 .76 .77 . STU-KD 87 .
+}
{+STU-KD-SA 87 .
 [Figure] [Table] [Table]+}
[Table]
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 ﬁrst group of experiments, M is set to 5 , and the target data are randomly partitioned intoparts each of which is held by one edge device.
Three schemes are evaluated on Ofﬁce-31 dataset:
SHOT (Liang et al., 2020), TO-KD [Table] 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, ﬁve edge devices and the server collaboratively train a global compact model.
The experimental results on A → W task and that on A → 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 by 6 .
58% on A → W task, and by 6 .
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 aggregation does not affect the performance of STU-KD.
In the second group of experiments, two transfer tasks are conducted: 1) A from Ofﬁce-31 dataset is selected as the source domain, and D and W from Ofﬁce-31 are selected as the target domains, denoted as A → DW ; 2) I from ImageCLEF-DA dataset is selected as the source domain, and C and
P from ImageCLEF-DA are selected as the target domains, denoted as I → CP .
In task A → DW , the two target domains are similar; while in task I → CP , the two target domain are much different from each other.
The number of edge device M is set to 6 .
For each transfer task, the ﬁrst 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 on the 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 task A → DW and task I → 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 target domains, though STU-KD does not provide extra performance gain for the compact model by involving more data into the training process.