Dynamic Knowledge Distillation for Black-box Hypothesis Transfer Learning

HTL provides generalized theoretical guarantees for finding optimal transfer parameters of source hypotheses. One main task of HTL is to figure out which hypotheses are helpful given a collection of source hypotheses Wang and Hebert (2016); Kuzborskij and Orabona (2017). Kuzborskij and Orabona (2013) analyzes the stability of transferring a single hypothesis to the target domain based on least-squares with biased regularization. It assumes that the source hypothesis must be a linear predictor living in the same space of the target predictor, which prevents the use of the source hypothesis as a black-box function. Another theoretical work Perrot and Habrard (2015) addresses HTL in the context of supervised regularized metric learning by using a biased regularization with source hypothesis. Recently in Fernandes and Cardoso (2019), the HTL framework is generalized to include four learning tasks: regression, classification, learning to rank and recommender system. The authors use the structural similarity as transferable knowledge, which can be coefficients of a support vector machine for classification or the weak estimators and their associated importance for an AdaBoost model. Du et al. (2017) proposes a practical HTL algorithm which predefines a transformation function and treat it as an input of the learning algorithm.

Among general knowledge transfer works, KD Hinton et al. (2014) is an efficient approach to build smaller, faster models by ”distilling” knowledge from large, complex models. There have been studies using KD in the context of transfer learning and domain adaptation when the source domain data is not available. For example, Ao et al. (2017) utilizes a generalized distillation framework (an extension of KD) to leverage the knowledge from the source domain by using both labeled and unlabeled target data. Besides, Nayak et al. (2019) addresses the no training data problem in KD by synthesizing the crafted samples from the source model and using them as surrogates to train the target model. In spite of using KD, Ahn et al. (2019) proposes a more general knowledge transfer framework through maximizing mutual information between two networks based on the variational information maximization technique. Another method Chidlovskii et al. (2016) extends feature corruptions and their marginalization to solve the domain adaptation problem. In healthcare applications, Hong et al. (2019) addresses cross-domain knowledge transfer on healthcare data sets, by jointly optimizing the combined loss of attention imitation and target imitation during KD. Mei and Xia (2019) discusses three knowledge transfer approaches for developing risk prediction models on electronic health record data, including injecting knowledge to input features, to the objective functions and to output labels.

With a non-negative loss function $\ell:\mathcal{Y}\times\mathcal{Y}\mapsto\mathbb{R}_{+}$, we denote by $\ell(f(\mathbf{x}),\mathbf{y})$ the loss of hypothesis $f$ on instance $(\mathbf{x},\mathbf{y})$. The risk of $f^{t}$, with respect to the target probability distribution $\mathcal{D}$, and the empirical risk measured on $D^{t}$ are defined as

In order to inject knowledge from $f^{s}$ while minimizing the above risks for $f^{t}$, we formalize the HTL problem within the framework of Regularized Empirical Risk Minimization (ERM):

where $\Omega(f^{t})$ is the regularization term on the target model $f^{t}$ to penalize its inconsistency with the knowledge contained in the source model $f^{s}$, and $\alpha,\beta$ are weighting coefficients deciding what proportion of knowledge from $f^{s}$ to be transferred.

In general, the regularization term for ERM can be in various forms. For example, if the source model $f^{s}$ is a white-box function and is in the same hypothesis space as the target model, it can be defined as the distance between two functions $\Omega(f^{t})=\|f^{t}-f^{s}\|$ Kuzborskij and Orabona (2013), where $\|\cdot\|$ is a norm function defined in space $\mathcal{F}$. In the case of multiple auxiliary hypotheses, hypothesis-wise weighting should be considered in this regularization term Kuzborskij and Orabona (2017), e.g., $\Omega(f^{t})=\|f^{t}-f^{s}_{\gamma}\|$, where $f^{s}_{\gamma}=\sum_{i}\gamma_{i}f^{s}_{i}$ is an optimal combination of source hypotheses. Besides, given the source domain data, Wang et al. (2019) define regularization as the performance gap between the source domain and the target domain $\Omega(f^{t})=(\mathcal{L}_{D^{s}}^{{\Gamma}^{s}}({f}^{t})-\mathcal{L}_{D^{s}}^{{\Gamma}^{s}}(f^{s}))+(\mathcal{L}_{D^{t}}^{{\Gamma}^{t}}(f^{s})-\mathcal{L}_{D^{t}}^{{\Gamma}^{t}}({f}^{t}))$ where ${\Gamma}^{s}$ and ${\Gamma}^{t}$ are instance-wise weighting parameters for data in the source domain and the target domain, respectively. Different from the above situations, we have no access to either the parameters of the source hypothesis or the source domain data. To overcome this challenge, we will resort to knowledge distillation and define $\Omega(f^{t})$ using the distillation loss.

In this section we propose dynamic knowledge distillation as a concrete algorithm of black-box hypothesis transfer learning formulated in Eq. 2. For the sake of simplicity, we omit the index of $(\mathbf{x}_{i},\mathbf{y}_{i})$ and directly use $(\mathbf{x},\mathbf{y})$ to denote the $i$-the sample. We also use $\mathbf{p}^{t}$ and $\mathbf{p}^{s}$ to denote the output probability vector of $f^{t}(\mathbf{x})$ and $f^{s}(\mathbf{x})$, respectively. We can formulate our regularized loss function as

The first term is the standard loss for the supervised learning task corresponding to the first part in Eq. 2; The cross-entropy loss is often used for a classification problem.

where $H(\cdot,\cdot)$ is the cross-entropy function, $C$ is the number of classes, $y_{j}$ indicates whether the sample belongs to class $j$, and $p^{t}_{j}$ is the probability of the sample belonging to class $j$ which is computed by the target model.

Inspired by knowledge distillation Hinton et al. (2014), the second term is the distillation loss which aims to leverage knowledge from $f^{s}$. The source model can generate a large probability for the true class while small probabilities for other classes. The dark knowledge underlying the small probabilities can be distilled using a softmax with a high temperature $T>1$:

However in the case of a black-box hypothesis $f^{s}$, we cannot directly get the logits $z_{j}$ needed to compute the soften probability $\tilde{p}_{j}$. We tackle this issue by first reverting the probability back to the logits $z_{j}=\log(p_{j})+c$ ($c$ is a constant number) and then plug $z_{j}$ into Eq. 5. As a result, specially for the black-box source hypothesis, we can soften the original predicted probability $p_{j}$ using

where the intermediate constant $c$ is cancelled out. Thus, we can calculate the distillation loss $\mathcal{L}_{2}$ for regularization by

where $\tilde{p}^{t}_{j}$ is computed by Eq. 5 for the target model $f^{t}$, while $\tilde{p}^{s}_{j}$ is computed by Eq. 6 for the black-box source model $f^{s}$.

The most important contribution lies in that we propose an instance-wise weighting mechanism for the coefficients $\alpha$ and $\beta$ to avoid negative transfer. We adjust the weights between the standard loss $\mathcal{L}_{1}$ and the distillation loss $\mathcal{L}_{2}$ based on a consistency score $S$ between the ground-truth label $\mathbf{y}$ and the probability $\mathbf{p}^{s}$ predicted by the source hypothesis $f^{s}$. The consistency score $S$ is expected to be in the range $[0,1]$, and to be larger when $\mathbf{p}^{s}$ is closer to $\mathbf{y}$. A practical definition of such $S$ can be the exponential of the negative cross-entropy loss of $\mathbf{p}^{s}$ against $\mathbf{y}$:

Finally, the instance-wise dynamic weighting mechanism is defined as

where $\lambda$ and $\delta$ are two hyperparameters satisfying $0\leq\lambda\leq 1$, $0\leq\delta\leq 1$, and $0\leq\lambda+\delta\leq 1$. The coefficient $\alpha$ will be instance-wise adjusted in the range $[\lambda,\lambda+\delta]$. Through this mechanism, the target learning task will focus more on the standard loss when $f^{s}$ generates inconsistent predictions with the ground-truth label, and hence avoid the negative transfer problem. The overview of our dynamic knowledge distillation is illustrated in Figure 1. We also summarize the detailed training procedure of $f^{t}$ in Algorithm 1.

To validate that our algorithm can overcome domain shift, we specially create a modified version of MNIST which consists of a source and a target domain. Specifically, we can omit all samples of certain digits to curate the source domain. For example, when omitting digit 3, we obtain a source domain with 49362, 4507 and 8990 samples in training, validation and test set, respectively. Meanwhile, to curate the target domain, we randomly sample 10% of all examples of MNIST, which gives us a target domain with 5500, 500 and 1000 samples in training, validation and test set, respectively. We call this modified MNIST dataset as “m-MNIST”. There are three key properties of m-MNIST: 1) there are overlaps between source domain and target domain, 2) target domain has much fewer samples than source domain, 3) source domain sees fewer classes than target domain. Therefore, m-MNIST can be regarded as a simulation of domain data shift, more specifically the prior shift, in domain adaptation.

We first train a CNN model with 2 convolution layers on the source domain of m-MNIST removing digits 3, to serve as the source model $f^{s}$ to be transferred. This model achieves an accuracy of 98.59% on the source test set. However, as can be seen in Table 1, it only obtains an accuracy of only 88.82% on the target test set due to its misclassification on digits 3. In fact, we can prove the source model is still very valuable since it provides a wealth of knowledge regarding to the other nine digits. We choose a one-layer MLP model as the target model $f^{t}$, and the TD model reaches an accuracy of 86.50%. If we applies skdHTL with $\lambda=0.5,T=2$, we get a negative transfer with a worse accuracy of 82.68%. Nevertheless, we can improve the accuracy to 88.62% using dkdHTL with $\lambda=0.1,\delta=0.9,T=2$.

Furthermore, if we omit more digits in the source domain, e.g., omitting both digits 3 and 7, the accuracy of the SH model will be further decreased to 77.44%. Given this imperfect source model, skdHTL leads us to a worse transfer result with an accuracy of 75.83% . Nevertheless, using dkdHTL we can still achieve a target model of accuracy as high as 88.42%. Besides, we also explore the setting where a two-layer CNN model is used as the target model $f^{t}$. The TD model and skdHTL model give us a target model of accuracy 94.06% and 79.15% respectively, which are better than the setting with an MLP as $f^{t}$ as expected. In the meantime, dkdHTL achieves the best accuracy of 94.25%.

Following the third setting in Table 1, we also investigate the effects of hyperparameters $\lambda$, $\delta$, and $T$ in dkdHTL. According to Table 2, we find that dkdHTL achieves the best accuracy of 96.68% with $\lambda=\delta=0.5,T=3$. More generally, a larger $\lambda$, which makes the target model focus more on the target data $D^{t}$, gives a better accuracy. Meanwhile, a larger $\delta$, which lets the target model learn more from the source model $f^{s}$ when it is consistent with ground-truth, gives a better accuracy as well. The temperature $T$ shows relatively mild effects on the accuracy.

Office-31 Saenko et al. (2010) is a standard dataset for visual transfer learning. It has three domains: Amazon (A), Webcam (W), and DSLR (D). Each domain contains 31 unbalanced classes with a total of 4,110 images. Choosing different pairs of the source domain and the target domain, we construct six HTL tasks: A$\to$W, D$\to$W, W$\to$D, A$\to$D, D$\to$A, and W$\to$A. The character on the left side of $\to$ represents the source domain from which the hypothesis $f^{s}$ is generated, while the right side means the target domain.

Following state-of-the-art transfer learning/domain adaptation works such as Zhang et al. (2019), we use ResNet-50 He et al. (2016) with parameters fine-tuned from the model pre-trained on ImageNet Russakovsky et al. (2015) for both $f^{s}$ and $f^{t}$. $f^{s}$ is first trained on the corresponding source domain in advance, and then used as a black-box function during the training of $f^{t}$.

We run the experiments on different partitions of Office-31 to evaluate the performance when the training data is large enough as well as when the training data is not sufficient. The first training-validation-test partition is 8:1:1 for each domain, and the second partition is 5:1:4, which has much fewer training instances than the first one. Classification accuracy for different tasks with a large training set as well as a small training set are reported in Table 3 and Table 4, respectively. Note that the performances of TD on all tasks decrease on small training set comparing to the large one, whereas the performances of SH don’t strictly align with the size of the test set. For skdHTL models, hyper-parameters are set as $\lambda=0.5$, $\delta=0$ and $T=2$, following the common KD practice. Two types of dkdHTL models with different hyperparameter settings are evaluated: 1) dkdHTL1 with $\lambda=0.5$, $\delta=0.5$ and $T=2$, and 2) dkdHTL2 with $\lambda=0.7$, $\delta=0.3$ and $T=4$.

From Table 3, we can observe that the two dkdHTL models consistently outperform almost all baselines. Comparing with skdHTL, the dynamic knowledge distillation is more suitable in transfer learning tasks. With a smaller training partition in the target domain, the overall performances become degraded due to the insufficiency of training samples, as shown in Table 4. For three adaptation tasks $W\to D$, $D\to A$, $W\to A$, dkdHTL2 still achieves the optimal performance (99.00%, 84.49%, 84.04%), whereas it is inferior to the TD model in tasks $A\to W$ and $A\to D$. It is reasonable since different tasks may have different optimal hyperparameter settings, due to the heterogeneity of transfer tasks. It is worth noting that the task $D\to A$ is the most difficult task out of all six tasks, and dkdHTL2 always gets the highest classification accuracy in this task. Overall, domain adaptation experiments on Office-31 validate the effectiveness of dkdHTL.

Based on MIMIC-III Johnson et al. (2016), a freely accessible critical care database, a recent work Harutyunyan et al. (2019) constructed benchmark machine learning datasets. In particular, an “In-Hospital Mortality” (IHM) task was defined to predict whether an ICU patient will die at discharge given the first 48 hours observation of the ICU stay.

In our experiment, we first split the MIMIC-III data into two domains according to the type of ICU admission. As can be seen in Figure 2, the source domain contains patients with admission type “Emergency/Urgent”, while the target domain contains patients with admission type “Elective”. The source domain has 13242, 2903 and 2880 samples in the training, validation and test set, respectively. Meanwhile the target domain has 2877, 606, and 712 samples in the training, validation and test set, respectively. Each sample has 48 timestamps of 76 features, with a label indicating mortality at discharge.

Note that there exists significant class imbalance in this IHM task. In detail, the mortality rate in the source domain is 14.30% and that in the target domain is only 5.46% (this also reflects the domain difference). Therefore with regard to model evaluation, we consider three measurements including accuracy, auROC (area under the Receiver Operating Characteristic curve) and auPRC (area under the Precision Recall Curve). To keep consistence with Harutyunyan et al. (2019), we adopt bidirectional RNN with LSTM (Long-Short Term Memory) units to implement all networks. The source model and target model have the same RNN architecture, both with 16 hidden units in the bidirectional LSTM layer.

We train a source model on the source domain, which generates accuracy of 88.68%, auROC of 84.52% and auPRC of 48.72%, on its test set. Note that this result is very close to the performance reported by Harutyunyan et al. (2019) which is trained on the whole dataset. According to the results reported in Table 5, the source model is a good teacher model on the target domain, with accuracy of 94.24%, auROC of 87.26% and auPRC 36.39%. In contrast, the TD model achieves the worst performance considering auROC of 84.06% and auPRC of 29.31%. The skdHTL model with $\lambda=0.5,\delta=0,T=2$ improves the performance of TD model to auROC of 86.46% and auPRC of 33.10%. The dkdHTL model with $\lambda=0.1,\delta=0.1,T=4$ brings us a better target model with auROC of 86.70% and auPRC of 34.08%, and with smaller standard deviations. This hyperparameter setting for dkdHTL is chosen from $\lambda,\delta\in[0.1,0.3,0.5,0.7,0.9],T\in[2,3,4]$ by grid search. We find that the dkdHTL model is slightly inferior to the SH model in terms of auROC and auPRC. It is because there is inherently a gap between the student model and the black-box teacher model in knowledge distillation. In fact, if the detailed parameters $\mathbf{w}_{s}$ of $f^{s}$ is available, stronger regularization terms such as $\|\mathbf{w}-\mathbf{w}_{s}\|$, can be added to our loss function in order to force the target model to mimic the source model more closely.