Contrastive Model Invertion for Data-Free Knolwedge Distillation

1 Introduction

Recent advances in deep learning have lead to a vast number of publicly available pre-trained models, which covers a variety of scenarios such as recognition [Krizhevsky et al., 2012] and detection [Ren et al., 2015]. Albeit their striking performance, most of these pre-trained models are cumbersome in size and problematic to be deployed into capacitysensitive edge devices. To resolve this issue, knowledge distillation (KD) was proposed to craft a lower-latency model by learning dark knowledge from pre-trained teacher models. It soon became a flourishing topic due to its effectiveness and engineering simplicity, and has been adopted to both Euclidean and non-Euclidean data [Hinton et al., 2015; Yang et al., 2020b]

∗Corresponding author.

Despite the promising results achieved, existing KD methods largely rely on massive training data to transfer the knowledge from pre-trained teacher models to students. However, in many cases, the training data is not released together with the pre-trained models, due to privacy or transmission reasons, which makes these methods inapplicable. Data-free KD is thus proposed to address this problem. The vital step of data-free KD is model inversion, for which the goal is to recover training data from pre-trained teacher models. With the derived synthesized data, student models can be learned with ease through directly leveraging the data-driven KD methods.

Model inversion per se has been studied for a long time for various purposes. For example, [Mahendran and Vedaldi, 2015] investigated model inversion for better understanding the deep representations. [Fredrikson et al., 2015] studied model inversion attack to infer sensitive information. Recently, the study of model inversion resurges as data-free KD gains more attention [Yin et al., 2020]. Specifically, datafree KD imposes higher demands on model inversion, due to the following reasons. First, the derived data should follow the same distribution as the original training data, since otherwise the student model can barely inherit knowledge from the teacher with these out-of-distribution data. Second, the derived data should be of rich diversity, so that the student model is able to master the comprehensive knowledge with these diverse data.

Existing inversion methods, unfortunately, remain incompetent in meeting such demands. The work [Chen et al., 2019], for instance, inverted the classification model by fitting the "one-hot" prediction distribution. [Yin et al., 2020], on the other hand, synthesized images by regularizing the distribution of intermediate feature maps using statistics stored in batch normalization layers of the teacher model. Both methods rely on some assumptions on the true data distribution and optimize each instance independently by fitting the prior distribution. As no constraints are explicitly imposed to encourage data diversity, these methods suffer from the mode collapse problem, where the generated instances turn out to be highly similar to each other. The work of [Micaelli and Storkey, 2019] proposed to generate more discriminant data for training by mining harder or adversarial examples. Although some performance gain is achieved for data-free KD, the generated data often lack natural image statistics and are easily identified as unrealistic.

In this paper, we attempt to alleviate the mode collapse problem in data-free KD, through the lens of promoting data diversity. From experiments, we observe that with the same amount of data, higher data diversity indicates stronger instance discrimination. Inspired by this phenomenon, we first propose a definition of data diversity based on instance discrimination, and then propose the Contrastive Model Inversion (CMI) for tackling mode collapse while keeping the distribution of generated data closer to real data distribution. In this way, the generated data becomes more diverse and realistic.

Specifically, in CMI we introduce another contrastive learning objective, where the positive image pair comprises a cropped image and whole image of the same data instances, while the negative image pair comprises two different data instances. By encouraging positive pairs of images to be close to each other and negative pairs of images to be isolated apart, CMI significantly improves the image diversity and realism, and therefore promotes the performance of data-free KD. Experiments of pre-trained models on CIFAR-10, CIFAR-100, and Tiny-ImageNet demonstrate that CMI not only ensures more visually plausible instances than the state of the arts, but also achieves significantly superior performance when the generated data are used for knowledge distillation.

Our contributions are therefore summarized as follows.

• We propose a definition of data diversity, which enables us to incorporate the diversity into the objective explicitly to improve the diversity of generated data.
• We propose a novel Contrastive Model Inversion approach to handle mode collapse in data-free KD while enforcing the distribution of generated data more to imitate the real one.
• Extensive experiments are conducted to validate the superiority of CMI with respect to the state of the arts.

2 RelatedWorks

Model Inversion (MI) aims at re-constructing inputs from the parameter of a pre-trained model, which is originally proposed to understand the deep representation of neural networks [Mahendran and Vedaldi, 2015]. Given a function mapping φ(x) and the input x, a standard model inversion problem can be formalized as finding a x′ to achieve the smallest d(φ(x), φ(x′ )), where d(·, ·) is a error function, e.g., mean squared error. This paradigm, known as model inversion attack [Wu et al., 2016], is widely used in several areas such as model security [Zhang et al., 2020] and interpretability [Mahendran and Vedaldi, 2015]. Recently, the inversion techniques have shown their effectiveness in knowledge transfer [Lopes et al., 2017; Yin et al., 2020], leading to data-free knowledge distillation.

Data-Free Knowledge Distillation recently emerges as a popular task within the realm of knowledge distillation [Ye et al., 2019; Shen et al., 2019; Yang et al., 2020a; Ye et al., 2020], and aims to learn a student model from a cumbersome teacher without accessing real-world data [Lopes et al., 2017; Ma et al., 2020], so as to achieve model compression [Yu et al., 2017]. The contributions of existing data-free works can be roughly divided into two categories: adversarial training and data prior. Adversarial training is motivated by robust optimization, where worst-case samples are synthesized for student learning [Micaelli and Storkey, 2019; Fang et al., 2019]. Data prior provide another perspective for data-free KD, where the synthetic data are forced to satisfy a pre-defined prior, such as total variance prior [Mahendran and Vedaldi, 2015] and Batch normalization statistics [Yin et al., 2020].

Contrastive Learning has achieved tremendous progress in field of self-supervised learning [Chen et al., 2020; He et al., 2020]. Its core idea to treat each sample as a distinct category and learn how to distinguish them [Wu et al., 2018; Liu et al., 2021]. In this work, we re-visit the contrastive learning framework from another perspective, where its ability of instance discrimination is used to model the data diversity in model inversion.

3 Method

3.1 Preliminary

Model inversion, as a vital step for data-free knowledge distillation, aims to recover training data X ′ from a pre-trained teacher model ft(x; θt) as an alternative to the inaccessible original data X. In this part, we discuss three typical inversion techniques:

BN regularization is originally introduced in [Yin et al., 2020] to regularize the distribution of X by making an Gaussian assumption about features from network. The regularization is usually represented as the divergence between the feature statistics N(µl(x), σ2 l (x)) and Batch normalization statistics N (µl , σ2 l ) as follows:

Class prior is usually introduced for class-conditional generation, which makes an "one-hot" assumption about the network predictions on x ∈ X ′ [Chen et al., 2019]. Given a pre-defined category c, it encourages to minimize the cross entropy loss:

Adversarial Distillation is motivated by robust optimization, where the x is forced to produce large disagreement between teacher ft(x; θt) and student fs(x; θs) [Micaelli and Storkey, 2019; Fang et al., 2019], i.e., maximize a Kullback-Leibler divergence term:

Unified Framework Combining the above mentioned techniques will lead to a unified inversion framework [Choi et al., 2020] for data-free knowledge distillation:

where α, β and γ are balance terms for different criteria. However, model inversion is actually a ill-conditioned problem where the synthetic data satisfying the above criteria are not unique. As no diversity constraints are explicitly imposed in this framework, the conventional inversion method may tend to be "lazy" and repeatedly synthesize duplicated samples. To overcome this problem, we propose a diversityaware inversion technique, namely Contrastive Model Inversion (CMI).

3.2 Contrastive Model Inversion

Overview With a pre-trained teacher model ft(x; θt), the goal of contrastive model inversion is to produce a set of x ∈ X ′ with strong diversity, with which comprehensive knowledge can be extracted from teachers. In this section, we develop an interesting definition for data diversity and, based on this, introduce the proposed Contrastive Model Inversion (CMI). Our motivation is intuitive: under the constraint of the same data amount, higher diversity usually indicates stronger instance distinguishability. To this end, we model the data diversity with an instance discrimination problem [Wu et al., 2018] and construct an optimizable objective with contrastive learning.

Definition of Data Diversity Given a set of data X ′ , An intuitive description of data diversity would be "how distinguishable are the samples from the dataset", which reveals a positive correlation between the diversity and instance distinguishability. Thus if we have a certain metric d(x1, x2) to estimate the distinguishability for an instance pair {x1, x2}, then we can develop a clear definition for data diversity as the following:

where d(x1, x2) will be applied to all possible (x1, x2) pairs from X. There are various ways to define the metric d(·, ·), leading to different diversity criteria. For example, the pretrained model ft(x; θt) is actually an embedding function that maps data x to a high-level feature space, where a naive metric can be defined as d(x1, x2) = ‖ft(x1)-ft(x2)‖, which is known as perception distance [Li et al., 2003]. However, this distance may be problematic for diversity estimation due to the following issues: 1) the function ft is actually not explicitly trained for measuring the similarity between instances, where the meaning of euclidean distance is unknown. 2), the embedding ft(x) may encode structured information about inputs, which can not be captured by this metric. 3) this distance metric is unbounded, and we are not to figure out how large it should be to indicate a good diversity. In this case, maximizing such a distance metric on ft may only lead to abnormal results. Thus, a more appropriate embedding space are required to construct a meaningful metric for distinguishability. In what follows, we propose a learning-based metric for data diversity, which is established by solving a contrastive learning objective.

Data Diversity from Contrastive Learning contrastive learning is original proposed to learn useful representations from data in a self-supervised manner, where an instancelevel discrimination problem is established by treating each sample as a distinct category [Wu et al., 2018]. Through contrastive learning, a network can learn how to distinguish different instances, which exactly coincides with our requirements for the metric d(·, ·). Based on this, we introduce another network h(·) as an instance discriminator upon the teacher network ft that accepts feature ft(x) as input and project it into a new embedding space. For simplification, we use v = h(x) to represent v = (h ◦ ft)(x) because the teacher network are fixed. In the new embedding space of h(·), we use simple cosine similarity to describe the relationship between data pair x1 and x2 as the following:

Then, an instance discrimination problem can be formulated in the form of contrastive learning framework [Chen et al., 2020], where each instance will be randomly transformed into different views and should be matched correctly. For each instance x ∈ X ′ , we construct a positive view x+ by random augmentation and treat other instances x- as negative ones. The contrastive learning loss is formalized as follows:

where τ denotes the temperature. Discriminator h(·) can learn how to distinguish different samples by push positive pairs closer and pull negative pairs apart, which provide a "contrast" metric for any (x1, x2) pair. To further understand how it related to the data diversity, let's define the metric d(x1, x2)=- exp(sim(x1,x2,h)/τ) exp(sim(x1,x1+,h)/τ) and rewrite Eqn. 5:

where the constant Z(x-) refers to the amount of negative samples for each instance xi. Thus we can directly maximize the diversity Ldiv by minimizing the contrastive loss Lcr.

Model Inversion In the previous part, we have bridged the data diversity with the contrastive learning objective, which can be directly optimized to make data more diverse. This section integrates contrastive learning into model inversion to form our final algorithm, namely contrastive model inversion.

As illustrated in Figure 3, our approach is consists of four components: a generator g(·; θg), a teacher network ft(·; θt), an instance discriminator h(·; θh) and a memory bank B. The discriminator is a simple multi-layer perception as used in [Chen et al., 2020], which accept the representations in the penultimate layer as well as the global pooling of intermediate features as the input.

The core idea of CMI is to incrementally synthesize some new samples that can be easily distinguished from the historical ones in the memory bank. Thus, the model inversion process is tackled in a "case-by-case" strategy, which means in each timestamp T , only one batch of data will be synthesized by the generator. Specifically, at the beginning of timestamp T, we re-initialize the generator and iteratively optimize its latent code z as well as the parameters θg for synthesis. In this case, the generator is only responsible for a small part of the data distribution. Compared with the pixel updating strategy used in [Yin et al., 2020] that update different pixels independently, the "case-by-case" generator can provide stronger regularization on pixels because they are produced from the shared weights θg. During synthesis, random augmentation will be applied on the synthetic images to produce a local view x and a global view x+ for contrastive learning. However, note that a single batch will be insufficient for training discriminators training. Thus we let the history images stored in memory bank B also participate in the learning process. Now, the objective for contrastive model inversion

Algorithm 1: Contrastive Model Inversion.

Input: A pretrained teacher ft(·; θt) Output: Image Bank B 1 B ← ∅ 2 initialize discriminator h(·; θh) 3 for number of batches do 4 initialize generator g(·; θg) 5 z ← N (0, 1) 6 for number of update iterations do 7 x ← g(z; θg) 8 xB ← sample(B) 9 L←αcr · Lcr(x ∪ xB, h) +βinv · Linv(x) 10 z ← z - η∇zL 11 θg ← θg - η∇θg L 12 θh ← θh - η∇θh L 13 end 14 B ←B ∪ x 15 end 16 return B

can be formalized as the following:

where Linv(·) refers to the widely used inversion criterion in Equation 4, which is only applied on the global view of images and the term Lcr refers to the proposed contrastive loss for data diversity. Note that Lcr takes both synthetic batch g(z; θg) and history data from B into account, where the history will provide useful guidance for the current image synthesis. During contrastive learning, we stop the gradient on global view and only allow backpropagation on local ones as done in [Chen and He, 2020]. We found that this operation can provide more clear gradient for local pattern synthesis. The proposed algorithm is summarized in Alg. 1.

3.3 Decision Adversarial Distillation

With the inverted data X , it is easy to train the student with the Kullback-Leibler divergence. However, the synthetic maybe not optimal for knowledge transfer, where some important patterns are missing. Adversarial distillation is a popular technique to improve student performance, which incorporates the student into image synthesis and maximizes the disagreement between teacher and student using Eqn. 10. However, large disagreement may not always correspond to valuable samples because they can be just some outliers. In this work, we pay more attention to those boundary samples and introduce a decision adversarial loss:

The function 1{·} is an indicator to enable adversarial learning when the teacher and student produce the same prediction on x, otherwise disable it. Unlike the unbounded loss term in Eqn. 10, our decision adversarial loss will make x close to the decision boundary, which can provide more information about teacher networks.

4 Experiments

4.1 Settings

In this work, we study the effectiveness of CMI on several network architectures, including resnet [He et al., 2015], vgg [Simonyan and Zisserman, 2014] and wide resnet [Zagoruyko and Komodakis, 2016]. Three standard classification datasets, CIFAR-10, CIFAR100 [Krizhevsky et al., 2009] and Tiny-ImageNet [Le and Yang, 2015] are utilized to benchmarks existing data-free KD approaches. CIFAR-10 and CIFAR-100 both contain 50,000 samples for training and 10,000 samples for validation, whose resolution is 32 × 32. Tiny-ImageNet consists of 100,000 training images of resolution 64 × 64, as well as 10,000 test images, which is more difficult than CIFAR. Details about datasets and training can be found in supplementary materials.

4.2 Benchmarks on Knowledge Distillation

Table 1 provides a benchmark for existing data-free knowledge distillation methods. T. and S. refers to the scratch accuracy of teachers and students on the original training data. We compare our approach with the following baselines: DAFL [Chen et al., 2019], ZSKT [Micaelli and Storkey, 2019], ADI [Yin et al., 2020], DFQ [Choi et al., 2020] and LS-GDFD [Luo et al., 2020]. They all follow the unified framework discussed in Sec. 3.1. Our approach and ADI both adopt a case-by-case strategy [Yin et al., 2020] for model inversion, while the other baselines tackle it as a generative problem. The main difference between our approach and baselines methods lies in the usage of contrastive learning, i.e., the additional loss term Lcr. Results show that our approach achieves superior or comparable performance compared with most baseline methods. Besides, note that the wrn40-2 contains about 10× fewer parameters than the resnet-34 models, where less prior information about training data is stored. In this case, CMI can still achieve satisfying performance compared with baseline methods.

In Fig. 2, we further visualize the synthetic data from DeepInversion (ADI without adversarial distillation), DFQ, and our approach. The results show that our approach is able to synthesize much more plausible and diverse data. Both DeepInversion and DFQ suffer from model collapse problems. For example, the synthesized frogs in DeepInversion are similar in color and background. It is worth noting that the generator can usually produce better visual quality than pixel update, although it may suffer from more severe mode collapse, as shown in the images from DFQ. In contrast, our approach is able to recover more diverse and plausible results, with better local details such as eyes and textures.

4.3 Ablation Study

Module Cut-off Our ablation study begins by investigating the contribution of different module introduced in CMI, including the contrastive loss and the generator described in Sec.3.2, as well as the decision adversarial loss. This experiment is organized in a "cut-off" schema, where each module will be separately turn off to see its effectiveness. As shown in table 2, turning off both the contrast loss and generator will lead to poor performance, but it seems cut-off the generator may lead to more severe degradation. We would like to argue that this is reasonable. As mentioned in Section 3.2, a generator is used to regularize pixels of synthetic data, and a crop operation is applied to construct local and global views. In this case, if we replace the generator with pixel update as used in [Yin et al., 2020], then only the pixels inside the crop region will be updated, leading to an imbalanced optimization. In our method, decision adversarial distillation is also an important technique for knowledge distillation. However, our method without adversarial distillation can still achieve the best performance (87.71) than the state-of-the-art (DFQ, 86.14 on CIFAR-10).

Influence of Contrast In our approach, we can control the data diversity by changing the balance term αcr. In this experiment, we take a further investigation on the influence of different αcr. Fig. 3 illustrates the relation between student performance and αcr with different data amount in both high contrast and low contrast settings. We find that an increasing αcr can bring significant benefits for student learning. In order to investigate how αcr affects the data distribution, we calculate the Frechet Inception Distance (FID) for the synthetic data, which is widely used to evaluate the generation quality in GANs. As shown in table 3, the FID score is estimated on three layers, corresponding to different feature level. The 1-st pooling layer contains more low-level features, while the last pooling layer mainly focuses on highlevel semantics. There are two observation in this experiment: first, increasing αcr can indeed achieve a better FID score, i.e., a better distribution; Second, we find significant FID improvements in the shallow layer, which means our approach can effectively improve the synthesis of low-level details.

Design of Instance discriminator Another topic in ablation study lies in the design of an instance classifier. We consider three design of instance discriminator as shown in table 4. First, we remove the discriminator and directly estimate the contrast in the feature space of the teacher model. This operation would be problematic because this space is not explicitly trained for similarity metric as described in Sec. 3.2. The second design is a linear discriminator, which is used to select a subset feature for contrastive learning. We compare it with a non-linear one used in our approach. We find that a stronger instance discriminator will be helpful for the student learning.

5 Conclusion

In this work, we propose a novel model inversion approach to guarantee the diversity of synthetic data, which can bring significant benefits for downstream distillation tasks. Extensive experiments are conduct in this paper not only to show the effectiveness of our approach but also explore how it works. model inversion is an interesting topic in the context of knowledge distillation, which may provide some insight for understanding what is learned by a network.

Acknowledgements

This work is supported by National Natural Science Foundation of China (U20B2066, 61976186), Key Research and Development Program of Zhejiang Province (2018C01004), the Fundamental Research Funds for the Central Universities (2021FZZX001-23) and Alibaba-Zhejiang University Joint Research Institute of Frontier Technologies.