Mosaicking to Distill: Knowledge Distillation from Out-of-Domain Data

Introduction

Knowledge distillation (KD) has emerged as a popular paradigm for model compression and knowledge transfer, attracting attention from various research communities . The goal of KD is to train a lightweight model, known as the student, by imitating a pre-trained but more cumbersome model, known as the teacher, so that the student masters the expertise of the teacher. In recent years, KD has demonstrated encouraging results over various machine learning applications, including but not limited to computer vision , data mining , and natural language processing

Nevertheless, the conventional setup for KD has largely relied on the premise that, data from at least the same domain, if not the original training data, is available to train the student. This seemingly-mind assumption, paradoxically, imposes a major constraint for conventional KD approaches: in many cases, the training data and even their domain for a pre-trained network are agnostic, due to for example confidential or copyright reasons. Hence, the in-domain prerequisite significantly limits the applicable scenarios of KD, and precludes taking advantage of the sheer number of publicly-available pre-trained models, many of which with unknown training domain , to carry out massive knowledge transfer.

In this paper, we aim at the ambitious goal of conducting KD using only out-of-domain (OOD) data, which, in turn, enables us to greatly relax the conventional prerequisite and thereby largely strengthens applicability of KD. Unarguably, OOD-KD is by nature a highly challenging task, since the domain discrepancy will inevitably impose a major obstacle towards the proper functioning of the pre-trained teacher. In fact, if we are to conduct naive KD on the raw OOD data, the resulting student model, as will be demonstrated in our experiments, fails to provide any performance guarantee on the target domain. This phenomenon signifies the limited generalization capability learned from OOD data, which is unsurprising.

To this end, we propose a novel assembling-by-dismantling approach, termed as MosaicKD, that allows us to take advantage of OOD data to conduct KD. Our motivation stems from the fact that, even though data from different domains exhibit divergent global distributions, their local distributions, such as patches in images, may however resemble each other. This observation further inspires us to leverage the local patterns, shared by the OOD and target-domain data, to resolve the domain shift problem in OOD-KD. As such, the core idea of MosaicKD is to synthesize in-domain data, of which the local patterns imitate those from real-world OOD data, while the global distribution, assembled from local ones, is expected to fool the pre-trained teacher. As shown in Figure 1, the shared local patterns are extracted from OOD data and re-assembled into in-domain data. Intuitively, this process is analogous to mosaic tiling, where tesserae are utilized to compose the whole art piece.

Specifically, in MosaicKD, we frame OOD-KD problem as a novel four-player min-max game involving a generator, a discriminator, a student, and a teacher, among which the former three are to be learned while the last one is pre-trained and hence fixed. The generator, as those in prior GANs, takes as input a random noise vector and learns to mosaic synthetic in-domain samples with locally-authentic and globally-legitimate distributions, under the supervisions back-propagated from the other three players. The discriminator, on the other hand, learns to distinguish local patches extracted from the real-world OOD data and from the synthetic samples. The entire synthetic images are fed to both the pre-trained teacher and the to-be-trained student, based on which the teacher provides category knowledge for data synthesis and the student mimics the behavior of the teacher so as to carry out KD. The four players collaboratively reinforce one another in an adversarial fashion, and collectively accomplish the student training.

In short, our contribution is the first dedicated attempt towards OOD-KD, a highly practical yet largely overlooked problem, achieved through a novel scheme that mosaics in-domain data. The synthetic samples, generated via a four-player min-max game, enjoy realistic local structures and sensible global semantics, laying the ground for a dependable knowledge distillation from the pre-trained teacher. We conduct experiments over classification and semantic segmentation tasks across various benchmarks, and demonstrate that MosaicKD yields truly encouraging results much superior to those derived by its state-of-the-art competitors on OOD data.

Related Work

Knowledge distillation. Knowledge Distillation aims to craft a lightweight student model from cumbersome teachers , by either transferring network outputs or intermediate representations . In the literature, knowledge distillation and its variants largely rely on the premise that original training data is available during distillation, which is vulnerable in real-world applications due to privacy or copyright reasons . Recently, data-free knowledge distillation has attracted attention from various research communities, which trains student model only with synthetic data. However, due to the difficulty in data synthesis without real-world samples, data-free KD usually leads to a degraded student on complicated tasks . Another solution to relax the conventional prerequisite on training data is to use some OOD samples. Unfortunately, it is found that naive knowledge distillation on OOD usually fails to learn a comparable student model from teachers .

Domain adaptation and generalization. Most learning algorithms strongly rely on the premise that the source data for training and the target data for testing are independent and identically distributed , ignoring the OOD problem that is frequently encountered in real-world applications. In the literature, the OOD problem is usually addressed by domain generalization (DG) or adaptation (DA) . Adaptation is a popular technique for aligning the source and target domain , which typically requires the target domain to be accessible during training. In recent years, Domain adaptation has been extended to open-set settings where the label space of training and testing data are different . In comparison, domain generalization is similar to domain adaptation, but does not requires the information from the target domain . Domain generalization trains a model on the source domain only once and directly applied the model on the target domain . Despite the success of DA and DG in supervised learning, OOD problem is still under-studied in the context of knowledge distillation.

Generative adversarial networks (GAN). Generative adversarial network is initially introduced by Goodfellow et al. for image generation , where a generator is trained to fool a discriminator in an adversarial minimax game. In recent years, several works have been proposed to improve the performance of GANs, from the perspective of image quality , data diversity , and training stability . Besides, Some explorations have been taken to make GAN training more efficient with limited data via transfer learning or augmentation . In this work, we study data synthesis in OOD settings, where the original training data is unavailable and thus conventional GAN technique can not be diretly deployed.

Out-of-Domain Knowledge Distillation

Problem Definition (OOD-KD). Given a teacher model $T(x;\theta_{t})$ obtained from training domain $\mathcal{D}=\{\mathcal{X},\mathcal{Y},P_{\mathcal{X}\times\mathcal{Y}}\}$, the goal of OOD-KD is to craft a student model $S(x;\theta_{s})$ only leveraging out-of-domain data from $\mathcal{D}^{\prime}=\{\mathcal{X}^{\prime},\mathcal{Y}^{\prime},P_{\mathcal{X}^{\prime}\times\mathcal{Y}^{\prime})}\}$, where $\mathcal{X}^{\prime}\neq\mathcal{X}$ and $\mathcal{Y}^{\prime}\neq\mathcal{Y}$.

In OOD-KD, due to the domain divergence between OOD data and original training data, some important patterns may be the missing and the corresponding knowledge on these patterns might not be appropriately transferred from teachers to students. To address the OOD problem, we propose a novel assembling-by-dismantling approach to craft in-domain samples from out-of-domain ones, which effectively alleviates the domain gap between the transfer set and unavailable training set, making KD applicable on out-of-domain data.

Proposed Method

In the absence of original training data $X$, directly minimizing the risk on an OOD set $X^{\prime}$ would be problematic due to the diverged data domain. In this work, we introduce a generative method for OOD-KD, dubbed as MosaicKD, where a generator $G(z;\theta_{g})$ is trained to synthesize a more helpful distribution $P_{G}$ for student learning. Specifically, MosaicKD is developed upon the distributionally robust optimization (DRO) framework which has been widely used to tackle domain shift . Given a pre-defined distance metric $d(\cdot,\cdot)$ for distributions, the basic form of DRO framework can be formalized as the following:

As mentioned above, the searching space built upon OOD data is insufficient for establishing an reliable upper bound for optimization. To address this problem, MosaicKD introduces a new way to construct the searching space based on local patches. Our motivation stems from the fact that, patterns of natural images are often organized hierarchically, where high-level patterns are assembled from low-level ones. Although the domain of original training data $X$ and OOD data $X^{\prime}$ are diverged, their local patterns may still resemble each other. For example, the patterns of “fur” can be shared by different animal species from varied domains. Note that each images is assembled from local patches, we propose an assembling-by-dismantling strategy to re-organizes shared local patches and synthesize in-domain data for training.

where $P_{X^{\prime}}$ refers to the distribution of OOD data and $P_{z}$ refers to the prior distribution of the latent variable $z$. $C(x^{\prime})$ and $C(G(z))$ refers to the cropped patches from OOD data and generated data. The main difference between Eqn. (3) and the objectives in vanilla GANs lies in the patch-level discrimination, where MosaicKD only regularizes local patterns to be authentic, leaving the global structure unrestricted. As mentioned above, global patterns can be assembled from local ones, MosaicKD assembles these patches to synthesize in-domain data through label space aligning.

Label Space Aligning. As no inter-patch restriction is introduced in Eqn. (3), the generator may only produce images with meaningless global semantic, although their local patterns are plausible. In this step, we turn to the teacher model for more information for in-domain data synthesis. In KD, the teacher model is trained on the original training data $X$, whose output is a conditional probability $T(x;\theta_{t})=p(y|x,\theta_{t})$, which corresponds to the confidence that $x$ belongs to $y$-th category. To align the label space, a naive method is to maximize the confidence of teacher predictions, i.e., minimizing the entropy term $H(p(y|x,\theta_{t}))$. However, prior works have shown that such a simple probability maximization may only lead to some “rubbish samples” without too much useful visual information for student training. To address this problem, we propose a regularized objective to align label space , which is formalized as:

In Equation 4, the first term refers to the discrimination loss as mentioned in Equation 3, which regularizes the local patterns to be authentic. The second term refers to the entropy loss for confidence maximization, which works on full images and assembles local patterns to synethsize desired categories. This objective will be simultaneously optimized with Equation (3) to keep the authenticity of local patches.

DRO in MosaicKD. As aforementioned, Equation 3 regularizes the local patterns to be authentic and Equation 4 aligns the label space of synthetic data to that of training data. They collaboratively construct a new searching space for DRO framework as follows:

Regularization $\mathcal{R}(G,D,T)$ force the generator to leverage local patterns of OOD data for data synthesis, which leads to a special searching space defined on all possible schemes of patch assembling. Different to the conventional DRO, MosaicKD uses a small radius for robust optimization, where the target domain can be covered by the searching space. We relax the regularization of Equation 5 to obtain an optimizable DRO objective for training, formalized as follows:

2 Optimization

Patch Discriminator. For training efficiency, the discriminator in equation 3 can be implemented as a Patch GAN with carefully designed receptive fields and patch overlap. Specifically, we stack several convolutional layers to build a fully convolutional network, whose output is a score map instead of a single true-or-fake scalar. Each score unit accepts a $L\times L$ local patches for discrimination. We apply an additional stride downsampling with step size $s$ on the score map to control the overlap between patches. A large step size $s$ will lead to more independent patches, which effectively reduce the structure restrictions in OOD images.

Full Algorithm. The full algorithm of MosaicKD is summarized in Alg. 1, where a generator $G(z;\theta_{g})$, a discriminator $D(x;\theta_{d})$, a fixed teacher model $T(x;\theta_{t})$ and a student $S(x;\theta_{s})$ are collectively optimized under the guidance of $\mathcal{L}_{local}$, $\mathcal{L}_{align}$ and $\mathcal{L}_{DRO}$.

Experiments

Datasets. The proposed method is evaluated on two mainstream vision tasks, i.e., image classification and semantic segmentation. Four datasets are considered in our experiments as in-domain training set, including CIFAR-100 , CUB-200 , Stanford Dogs and NYUv2 . For OOD-KD, we substitute original training data with OOD data, including CIFAR-10 , Places365 , ImageNet and SVHN .

Evaluation metrics. For image classification, accuracy and Frechet Inception Distance (FID) are used to evaluate different methods. FID indicates the divergence of two datasets, which was originally used to assess the synthesis quality of GANs :

where $(\mu_{1},\Sigma_{1})$ and $(\mu_{2},\Sigma_{2})$ are the mean and covariance statistics of generated and original samples. For semantic segmentation, we use mean Intersection of Unions (mIoU) as the performance metric. More details about datasets, training protocol, and metrics can be found in supplementary materials.

2 Results of Knowledge Distillation

CIFAR-100. Table 1 reports the results of knowledge distillation on CIFAR-100 dataset. Here we use CIFAR-10, ImageNet, Places365 and SVHN as OOD data to evaluate MosaicKD for OOD settings. We compare the proposed MosaicKD to various baselines, including data-free KD methods (DAFL , ZSKT , DeepInv. , DFQ ) and OOD-KD methods naively adapted from state-of-the-art KD approaches (BKD , Balanced , FitNet , RKD , CRD and SSKD ).

As shonw in Table 1, despite the mismatched distributions, conventional KD approaches still learn some useful but incomplete knowledge from OOD data (i.e., yielding significantly superior performance to random guessing), which indicates the existence of shared patterns between OOD data and training data. Further, some exploration was taken to evaluate importance of category balance and representation transfer for OOD-KD. First, we balance the OOD data by re-sampling the scarce categories according to teacher’s predictions. However, results show that balancing the OOD data may not help students learn correct class information, because most samples in OOD data are just misclassified outliers. In the context of OOD settings, the balance operation may lead to over-fitting on outliers, which may further degrade the student performance on the test set. As mentioned before, OOD data and original data may share some local patterns, which can be extracted by shallow layers of networks. We apply four representation transfer approaches, i.e., FitNet, RKD, CRD and SSKD to study the their role in OOD-KD. Compared with RKD that focus instance relation, we found that response-based methods like fitnet can transfer more helpful information in OOD settings, where the student directly imitates the teacher’s intermediate outputs of teachers. In general, transferring low-level features sometimes can be helpful for OOD-KD. However, note that CRD works on the high-level representation extracted from the penultimate layer, transferring these knowledge may be inappropriate for OOD-KD because high-level features may be unrelated to target tasks.

In this work, we handles the OOD-KD problem as a generative problem, instead of directly using OOD data for training. The proposed method leverage local patterns of OOD data for data synthesis, where some task-related patterns will be “assembled” from shared local patches. Results show that these re-assembled data can effectively transfer knowledge from teachers to students. In Table 1, we also extend our method to different types of OOD sets. We found that the performance of MosaicKD is related to the degree of domain divergence between OOD data and original data. For example, ImageNet is an object recognition dataset while Places365 is a scene classification dataset. Results show that, for the target data CIFAR-100, MosaicKD can achieve better performance on ImageNet compared to Places365.

Fine-grained Classification. To further study the effectiveness of our approach, we conduct knowledge distillation on fine-grained datasets as shown in Table 3, using Places365 as OOD data. OOD-KD on fine-grained data is a challenging problem, because different categories is visually similar. Results show that our method achieves superior performance compared to baselines methods.

Semantic Segmentation. Semantic segmentation can also be viewed as a classification task, where the network is trained to predict the category of each pixel. We apply our method to the NYUv2 dataset, following the protocol in . The teacher network is a deeplab v3 network with resnet-50 backbone. The student is a freshly initialized deeplabv3-MobileNetv2 model. Our method can effectively improve the knowledge transfer on OOD data and achieve competitive results even compared to vanilla KD settings.

3 Quantitative Analysis

Data balance and FID. Figure 3 provides some statistical information of OOD data and generated samples, including the category balance predicted by teachers and the per-class FID scores. The category is ranked according to their percentages. Note that the original CIFAR-10 dataset only contains 10 categories, which is very limited compared with the 100 categories of CIFAR-100. As illustrated in Figure 3 (a), some CIFAR-100 categories are missing in CIFAR-10. Besides, the large FID between OOD data and original training data also indicates that, even though some samples are categorized to some classes by the teacher, they may still belong to outliers. by contrast, our method successfully balances different CIFAR-100 categories and alleviates the domain gap (lower class FID), especially for unbalanced categories.

The influence of patch size. Patch size plays an essential role in our method, which determines the flexibility of data synthesis. The optimal patch size actually depends on the divergence between OOD data and in-domain data. For OOD data with large domain discrepancy, small patches are usually required due to the limited local similarity. As shown in Table 4, we evaluate our method with different patch sizes and report the test accuracy of student models as well as the patch FID. According to the Table 4, we find that for OOD data like CIFAR-10 and Places365, a large patch size (e.g., 18) can be used for student learning. However, for SVHN dataset, a smaller patch size would be more appropriate, as SVHN severely diverges from CIFAR-100.

Ablation study. In this section, we conduct ablation to understand the role of patch learning further. We consider the following settings: (a) MosaicKD (b) mosaicKD without patch learning (c) mosaicKD without discriminator (d) mosaicKD wihout adversarial training. As shown in Table 5, we find that full image discrimination sometimes even lead to worse results than MosaicKD without discrimination. Figure 4 visualizes the synthetic data with or without patch learning. Results show that the generator without patch learning is trapped by the label space of OOD data, failing to synthesize in-domain categories like trees and apples.

Discussion

Relation to unlabeled knowledge distillation. In the literature of knowledge distillation, a slice of works also use unlabeled data for student learning. However, they either make an i.i.d. assumption about unlabeled data and original training , or assume that there are sufficient in-domain samples inside unlabeled set . In this work, we allow the unlabeled set to be fully OOD, which is more challenging than existing unlabeled settings.

Relation to data-free knowledge distillation. The commonality between the data-free algorithm and MosaicKD lies in that they both solve the KD problem through data synthesis. However, data-free KD leverages some simple priors such as category confidence and gaussian assumptions , which ignores the structural details in natural images. By contrast, MosaicKD achieves data synthesis in an assembling-by-dismantling manner, where natural patterns can be utilized to improve synthesis quality. The training curves of different methods can be found in Figure 5.

Conclusion

In this work, we propose a novel approach, termed as MosaicKD, that enables us to take advantage of only the OOD data for KD. MosaicKD follows a assembling-by-dismantling scheme, where a synthetic sample is generated by locally resembling the real-world OOD data while globally fooling the pre-trained teacher. This is technically achieved through a handy yet effective four-player min-max game, in which a generator, a discriminator, and a student network are learned in the presence of a pre-trained teacher. We validate MosaicKD over classification and semantic segmentation tasks across various benchmarks, and showcase that it yields results significantly superior to the state-of-the-art techniques on OOD data.

Acknowledgements and Disclosure of Funding

This work is supported by National Natural Science Foundation of China (U20B2066, 61976186), Key Research and Development Program of Zhejiang Province (2020C01023), the Major Scientific Research Project of Zhejiang Lab (No. 2019KD0AC01), the Fundamental Research Funds for the Central Universities, Alibaba-Zhejiang University Joint Research Institute of Frontier Technologies, Start-Up Grant from National University of Singapore (R-263-000-E95-133), and MOE AcRF TIER 1 FRC Research Grant (R-263-000-F14-114).

References

Appendix

In this document, we provide details and supplementary materials that cannot fit into the main manuscript due to the page limit. Specifically, we provide optimization details of MosaicKD in Sec. 9.1, experimental settings in Sec. 9.2, and more experimental results in Sec. 9.3.

In this work, we deploy a generator to synthesize the transfer set for knowledge distillation. Nevertheless, GANs are known to suffer from mode collapse and fail to produce diverse patterns. To this end, we leverage both OOD data and synthetic ones to train our student models, so that the generator does not need to synthesize all samples for KD. Besides, an additional balance loss is deployed to alleviate mode collapse during training, defined as:

where $p(y|x,\theta_{t})$ is the probability prediction after softmax, and $P_{G}$ denotes the distribution of generated samples. Minimizing Eq. (9) will enforce the class to be balanced during the synthesizing process.

1.2 Objectives of MosaicKD.

As shown in the main manuscript, MosaicKD aims to solve a distributionally robust optimization (DRO) problem as follows:

where $\mathcal{R}(G,D,T))\leq\epsilon$ defines the search space, i.e., a ball space with radius $\epsilon$ centered at an distribution satisfying $\mathcal{R}(G,D,T))=0$. The specific form of center distribution is unknown, but we can still train a generator $G$ to approximate it. Note that Eq. (5) is intractable due to the non-differentiable condition on the search space. With the help of lagrange duality, we can re-express the inner part of Eq. (5) as follows:

where $\lambda$ is Lagrangian multiplier and $\lambda\epsilon$ is a constant term. If $\mathcal{R}(G,D,T))\leq\epsilon$, we choose $\lambda=0$, i.e., no restriction on $\mathcal{R}(G,D,T))$, to obtain the minimal cost. If $\mathcal{R}(G,D,T))>\epsilon$, then a large $\lambda$ should be applied as a penalization. According to the derivation of Eq. (11), we obtain a relaxed version of the intractable Eq. (5), expressed as follows:

1.3 GAN Training and JS Divergence.

Following the conventions of prior works, we write the GAN training objective as follows,

As proposed in , for a fixed generated $G$ and a given data distribution $P_{data}$, the optimal discriminator $D$ is achieved when

We then replace the discriminator in Eq. (13) with the optimal one $D^{*}$, which leads to the following optimization for generator $G$:

Therefore, as mentioned in the manuscript, we optimize generative adversarial networks to minimize the regularization term $R(G,D,T)$, which is equivalent to optimizing the JS divergence between patch distributions.

2 Experimental Settings

The proposed method is evaluated on two mainstream vision tasks, i.e., image classification and semantic segmentation, over four labeled datasets for teacher training and four OOD data for student learning, as summarized in Table 6. Note that CIFAR-100, ImageNet, and Places365 may contain in-domain categories. We craft OOD subset from the full ImageNet and Places365 datasets by selecting samples with low prediction confidence, as described in Algorithm 9.2. These OOD subsets can be viewed as out-of-domain data for CIFAR-100. Besides, we resize the OOD data to the same resolution as in-domain data, e.g., $32\times 32$ for CIFAR-100, $64\times 64$ for fine-grained datasets, and $128\times 128$ for NYUv2.

In this work, all teacher models are trained using the in-domain datasets listed in Table 6 with cross entropy loss. We use SGD optimizer with $\{lr=0.1,weight\_decay=1e-4,momentum=0.9\}$ and train each model for 200 epochs, with cosine annealing scheduler. In knowledge distillation, student models are crafted using unlabeled datasets, where only the soft targets from teachers are utilized. We use the same training protocols as the teacher training and report the best student accuracy on test sets. We use Adam for optimization, with hyper-parameters $\{lr=1e-3,\beta_{1}=0.5,\beta_{2}=0.999\}$ for the generator and discriminator.

The architecture of GAN for CIFAR-100 dataset is illustrated in Tables 8 and 8. For CUB-200 ($64\times 64$) and NYU ($128\times 128$), we add more convolutional layers and upsampling or sampling layers to generate high-resolution images.

3 More Experimental Results

Given a fixed patch size, the overlap between patches plays an important role in patch learning. The overlap is controlled by interval sampling in the patch discriminator. Note that the discriminator produces a prediction map to predict each small region on the original image, which means that distant predictions should share less information. We add a prediction stride to the final discrimination to control the patch overlap. Table 9 shows the student accuracy obtained with different patch overlaps, where a larger stride corresponds to a smaller overlap. The results show that increasing stride does not benefit the students’ accuracy. Note that we use the patch GAN architecture for patch learning, which contains internal stride operations within the discriminator. These stride operations already provide an appropriate overlap for patch learning. Besides, a larger stride also means fewer training samples, which may be harmful to the GAN training.

3.2 DRO Regularization

In MosaicKD, the search space is regularized by $\mathcal{L}_{local}$ and $\mathcal{L}_{align}$, which enforces the generated samples to be locally authentic and globally legitimate. We take a further study on the above regularization to show their significance for MosaicKD. As illustrated in 6, we visualize the generated samples with different regularizations. In Figure 6(a), no regularization is applied on the generator, and we naively maximize the teacher’s confidence, which will lead to some inferior samples . In Figure 6(b), the discriminator makes decisions on full images, and, to some extent, the generator will be trapped by the class semantic of OOD data, i.e., synthesizing a car-like apple or a horse-like maple. Figure(c) showcases the synthetic samples of MosaicKD, which reveals the correct semantic of task-related classes.

3.3 ImageNet Results

Table 10 provides the student’s accuracy on $32\times 32$ ImageNet dataset with 1000 categories. We use Places365 as the OOD data and resize all samples to $32\times 32$ for training. Results show that our approach is indeed beneficial for the OOD-KD task.