Pseudo-spherical Knowledge Distillation

1 Introduction

The knowledge distillation (KD) aims to train a network by utilizing the information given by other networks. In particular, smaller networks benefit from knowledge distillation from a bigger network. The KD objective is consists of standard classification loss and a distillation loss, which maximizes the coincidence between the probabilistic outputs of student and teacher networks. Subsequently, many distillation methods proposed distilling auxiliary information such as intermediate features or attention maps. Yet, those methods do not marginally outperform KD. Many recent works focused on transferring the structural knowledge of teachers by using contrastive learning [Tian et al., 2019; Xu et al., 2020; Chen et al., 2021]. However, they showed that attaching the KD objective to contrastive objectives significantly boosts the performance, avowing that the contrastive methods and KD are complementary to each other.

∗Contact Author

While the original KD uses cross-entropy loss for the distillation objective, we propose alternative distillation objectives that maximize the alignment between the probabilistic outputs of the teacher and student network. Our method is built on a scoring rule, which quantitatively measures the quality of a predictive distribution with respect to a reference distribution. Thus, we aim to optimize the student network by maximizing the scoring rule with respect to the probability of teacher network. Among various scoring rules, we focus on pseudospherical scoring rules, which is a representative family of proper and homogeneous scoring rules, that is suitable for optimization on softmax operator-based probabilistic outputs. To that end, we propose pseudo-spherical knowledge distillation, which utilizes pseudo-spherical scoring rule for distillation objective. We present an efficient implementation of the new distillation objectives and provide an in-depth analysis of the objectives. Through experiments on CIFAR-100 and ImageNet model compression benchmarks, we demonstrate the effectiveness of our new distillation methods. Especially, by combining contrastive distillation methods, we achieve stateof-the-art performance in CIFAR-100 benchmarks.

2 RelatedWork

Knowledge Distillation. Since the introduction of model compression [Zeng and Martinez, 2000; Buciluˇa et al., 2006], [Hinton et al., 2015] proposed knowledge distillation which effectively transfers the general and comprehensive knowledge trained from a teacher network to a student network. It allows the student network to learn the knowledge of the teacher network by minimizing the cross-entropy loss between their probabilistic outputs. Later studies focused on transferring auxiliary information other than probabilistic outputs. For example, [Romero et al., 2014] proposed FitNets, where the student learns from the feature maps of the intermediate layers. [Zagoruyko and Komodakis, 2016a] proposed attention transfer (AT), where the activated parts of the feature maps are transferred to the student. However, [Tian et al., 2019] showed that those methods do not easily outperform knowledge distillation, even have inferior performance when the student and teacher network have different architectural styles.

Recent works focus on passing knowledge that captures the structural correlation between the representations by using contrastive learning objectives.

[Tian et al., 2019] proposed contrastive representation distillation (CRD), where the contrastive objectives are used to maximize the mutual information between teacher and student representations. [Xu et al., 2020] proposed SSKD, which applies self-supervised methods to refine richer representational knowledge of a teacher network and transfer it to a student network. Also, [Chen et al., 2021] exploits Wasserstein distance to perform contrastive learning on both global and local scales.

Scoring Rules. In statistical decision theory [Dawid, 1998], the scoring rule measures the utility of a predictive distribution by assigning a numerical score on the events that materialize. [Gneiting and Raftery, 2007] elaborates the proper scoring rule, which necessitates the existence of the maximum. Thus, the proper scoring rules provide attractive loss and utility functions for estimation problems. Furthermore, the homogeneity of a scoring rule has been studied for unnormalized density estimation. [Takenouchi and Kanamori, 2017] proposed to learn unnormalized statistical models on discrete space by using the homogeneous scoring rule. Recently, [Yu et al., 2021] proposed to train an energy-based model based on the pseudo-spherical scoring rule, with application to generative modeling on high dimensional data.

3 Preliminaries

3.1 Backgrounds on Knowledge Distillation

For a K-class classification problem, let zT , zS ∈ RK be logits of teacher and student network. Then the standard KD train the student network by minimizing following objective:

where y is a one-hot label and α, β > 0 are balancing weights. In addition to the conventional classification loss Lcls, the distillation loss Ldist encourages the probability of student network to be similar to the probability of teacher network. The probability outputs are computed by a softened softmax operator, which is defined by the following:

where τ > 0 is a temperature that regulates the sharpness of the probability distribution. Let pT = στ(zT ) and pS =

στ (zS ) be the probability outputs of each teacher and student network, then the distillation objective is a cross-entropy loss between them:

where H(p, q) = - ∑k pk log qk is a cross-entropy loss, and τ 2 is multiplied to match the magnitude of gradient [Hinton et al., 2015]. However, we question the optimality of the crossentropy objective in distilling probability outputs. To that end, We propose novel distillation objectives by using the scoring rule, an alternative measure that accounts for the probabilistic discrepancy of two distributions.

3.2 Proper and Homogeneous Scoring Rule

The scoring rule evaluates the quality of a probabilistic forecast by computing a score based on the predictive distribution and the events that substantiate. Formaly, let PK be a space of all K-categorical distributions, i.e., for any p ∈ PK, ∑k∈[K] pk = 1. For a q ∈ PK, let us denote the scoring rule by S(k, q) for each k ∈ [K]. Then given a reference distribution p ∈ PK, the expected score S(p, q) is defined by

Definition 1. [Gneiting and Raftery, 2007] A scoring rule S : [K] ×PK → R is proper if the corresponding expected score satisfies:S(p, q) ≤ S(p, p), ∀p, q ∈ PK, (5)

and it is strictly proper the equality holds if and only if p = q.

Therefore, the proper scoring rule encourages the forecaster to match their predictions to the true beliefs. Given a (strictly) proper scoring rule S, it induces generalized entropy

Then one can show that HS(p) is a (strictly) convex function of p, and it induces a Bregman divergence

Also, the reverse holds from the following proposition:

Proposition 1. A scoring rule S : [K]×PK →R is (strictly) proper iff S(p, p) is a (strictly) convex function of p.

Since our goal is to match the probability of a student network to the probability of a teacher network, the proper scoring rule is an attractive distillation objective from the following:

if S is a strictly proper scoring rule.

Moreover, we introduce homogeneous scoring rule which particularly suits for softmax based probability distributions.

Definition 2. [Parry et al., 2012] A scoring rule is homogeneous if

Suppose a distribution is given by a softmax operator on a vector z ∈ RK , i.e., q = στ(z). Then a (positive) scalarmultiplication on q does not change the inference of a classifier since

for any λ > 0. Therefore, the homogeneity of a scoring rule reduces the search space for q, which makes the optimization more suitable. Moreover, if a scoring rule S is homogeneous, it suffices to compute the score with respect to the unnormalized distribution q = exp(z). The following examples demonstrate the popular choices for proper scoring rules.

Example 1. (log scoring rule) A log scoring rule is defined by Slog(k, q) := log qk. Then the corresponding expected score with respect to p is given as Slog(p, q) = ∑k∈[K] pk log qk, where it is equivalent to the maximum likelihood estimation (MLE). Remark that the log scoring rule is strictly proper, but is not homogeneous.

Example 2. (spherical scoring rule) The classical spherical scoring rule [Friedman, 1983] is defined by

where ∥q∥22 = ∑k q2 k . Then for any p, q ∈ PK, we have

where ⟨p, q⟩ is a inner product and ∠(p, q) is a angle between p and q. Therefore, the spherical scoring rule is strictly proper since it is maximized iff p = q. Also, since the scalar multiplication on q does not change the angle, it is homogeneous.

3.3 Pseudo-spherical Scoring Rule

We introduce pseudo-spherical scoring rule [Gneiting and Raftery, 2007], which is a family of proper and homogeneous scoring rules that generalizes spherical scoring rule.

Definition 3. Given γ > 0 and q ∈ PK, the pseudo-spherical scoring rule is defined as following:

where ∥q∥γ+1 := (∑k qγ+1 k ) 1 γ+1 .

Then the expected pseudo-spherical score with respect to reference distribution p ∈ PK is defined as

Alike spherical scoring rule in Example 2, the pseudospherical scoring rule is strictly proper and homogeneous.

Proposition 2. [Gneiting and Raftery, 2007] The pseudospherical scoring rule Sγ(p, q) is strictly proper and homogeneous for γ > 0.

4 Pseudo-spherical Knowledge Distillation

We present pseudo-spherical knowledge distillation (PSKD), where the knowledge of a teacher is distilled to a student by maximizing the pseudo-spherical scoring rule.

4.1 Design of objective

Since Sγ(p, q) is homogeneous, if we let q = q/Z be unnormalized distribution of q with some constant Z > 0, we have

Moreover, if q = στ (z) for some vector z ∈ RK , we have

where τ > 0 is a temperature. However, maximizing Sγ(p, q) contains fraction term, the gradient based optimization is intractable. Therefore, we define the PSKD objective Lγ(p, q) by taking a logarithm to make it feasible. We present two different approaches. First, by taking the 1 γ log(·) inside the expected score, we have

or by taking log outside of the expected score, we have

where we negate the terms to fit in minimization problem. The negative of latter objective is called γ-score [Takenouchi and Kanamori, 2017], which is also a strictly proper and homogeneous scoring rule as it is the composition of strictly increasing log function and pseudo-spherical scoring rule. [Fujisawa and Eguchi, 2008] also called the latter objective as γ cross-entropy, as a generalization of cross-entropy loss. Indeed, following proposition reveal that the PSKD objectives are generalizations of cross-entropy loss.

Proposition 3. [Yu et al., 2021] When γ → 0, we have

where H(p, q) = -Ep[log q] is a cross-entropy loss.

Since the original KD uses cross-entropy loss for distillation, PSKD is a generalization of KD, which is a special case when γ → 0.

Extension to negative orders. Even though the pseudospherical scoring rule is defined on γ > 0, we extend the PSKD objectives to γ < 0 (then the corresponding scoring rule might not be strictly proper).

Note that the following holds from the convexity of - log function and Jensen's inequality:

thus we have

and it is reversed when γ < 0. One would expect Lin γ to be superior since it upper bounds Lout γ when γ > 0, and Lout γ is superior when γ < 0. This is supported empirically by comparing the performance on model compression tasks (see section 5.1.

4.2 Implementation of PSKD

Let the logits of teacher and student network by zT and zS respectively, and let pT = στ (zT ) and pS = στ (zS ) be the probability output of each teacher and student network with temperature τ > 0. Then from the homogeneity of pseudo-spherical scoring score and as pT = exp(zT /τ ), pS = exp(zS /τ ), we haveLin

and

Note that both objectives can be simply implemented by a simple log-sum-exp function, and requires no additional computation compared to the original cross-entropy.

Note that the gradient of Lout γ is given by:

Since exp(z/τ) ≈ 1 + z/τ for sufficiently large τ and by assuming that the logits have zero-mean as in [Hinton et al., 2015], we have στ (z) ≈ 1+z/τ K . Thus, it follows that

On the other hand, the gradient of Lin γ is given by:

and by using similar logic, we have

Therefore for large τ, the PSKD objective conducts logit matching, and for small τ, it focuses on having the same label for student and teacher (i.e. label matching) similar to KD [Hinton et al., 2015].

Finally, the complete objective is consists of standard classification loss and pseudo-spherical distillation loss as in eqn. (1):

where α, β > 0 are balancing weights and Lγ can be both Lin γ and Lout γ . Note that we multiply τ 2 since the gradient scales as 1/τ 2 .

5 Experiment

For the experiments, we demonstrate the effectiveness of PSKD on model compression tasks. First, we conduct ablation on PSKD objectives with respect to different values of γ. Then, we compare PSKD to cutting-edge distillation methods on CIFAR-100 and ImageNet benchmarks.

5.1 Ablation Study

Effect of γ. First, we conduct ablation studies on the performance of PSKD with different values of γ. Given a pretrained wide ResNet 40-2 [Zagoruyko and Komodakis, 2016b] teacher network, we train wide ResNet 16-2 student network on CIFAR-100. For evaluation, we report top-1 test accuracy of student network and the centered kernel alignment (CKA) similarity [Kornblith et al., 2019] of teacher and student network. The CKA similarity is a reliable measure of the similarity of two neural representations, which is invariant to orthogonal transformation and isotropic scaling. Thus, high CKA similarity indicates that the student learns similar representations to the teacher. We took the output of the penultimate layer (i.e. after average pooling) to calculate the CKA similarity.

Figure 2 demonstrates the effect of γ for both Lout γ and Lin γ . We observe that the test accuracy of student networks distilled by Lin γ increases as γ increases from -1 to 1, while the CKA similarity remains unchanged. On the other hand, the student networks distilled by Lout γ achieve their maximum around γ = -0.5 and it drops when γ becomes larger. Also, the CKA similarity significantly drops as γ increases. Remark that when distilling by Lout γ , the CKA similarity is strongly correlated with the test accuracy of student networks, yet it doesn't hold when distilling with Lin γ .

For the case when a student is trained by Lin γ , the CKA similarity doesn't change for different values of γ, showing that the value of γ on distilling representational knowledge doesn't vary. However, the test accuracy varies as γ becomes larger. We provide analysis by the following gradient. From eqn. (27), recall that the gradient of Lin can be approximated by a logit matching between zT and (γ + 1)zS . Thus when the logit of a teacher network has a high value, the larger value of γ makes student logit easier to follow the teacher logit.

From eqn. (25), the approximation on the gradient of Lout γ is a logit matching between zT and zS , which is independent to γ. But, from eqn (24), the gradient of Lout γ is expressed by the difference between στ (zT + γzS ) and στ ((γ + 1)zS ). Then for γ > 0, the softmax output becomes smoother as γ increases, thus the gradient becomes smaller and reduces the effect of distillation. Therefore, Lout γ has relatively better performance when γ < 0.

Comparison of two PSKD objectives. We further compare two objectives by conducting experiments on various teacher and student combinations. We choose γ = -0.5 and γ = 1.0 for baseline. Table 2 reports the test accuracy of student networks trained with different PSKD objectives and values of γ. Remark that Lout outperforms Lin when γ = -0.5 and Lin outperforms Lout when γ = 1.0, ascertaining the hypothesis asserted in section 4.1. While Lin with γ = 1.0 and Lout with γ = -0.5 achieve similar performances, the former achieves slightly better performance overall. Therefore, we use Lout with γ = -0.5 for the rest of the experiments.

5.2 Main Results

Setup. We followmodel compression benchmarks that were proposed in CRD [Tian et al., 2019]. For experiments on CIFAR-100, there are 13 teacher-student combinations where 7 are of similar architectures and 6 are of different architectures. Those architectures include resnet [He et al., 2016], vgg [Simonyan and Zisserman, 2014], and ShuffleNet [Ma et al., 2018; Zhang et al., 2018]. We used pre-trained teacher models provided in the official CRD repository 1 . For the experiment on ImageNet, we used the official PyTorch pretrained ResNet-34 for a teacher and ResNet-18 for a student.

Combination with other distillation methods. Recently, many distillation methods resort to transfer structural knowledge of teacher networks by using contrastive learning. It has been empirically verified that using the contrastive distillation method and KD simultaneously improves performance, showing that they are complementary to each other. We show that using PSKD and contrastive distillation together further improves the performance. To that end, we add the contrastive distillation objective Lcon to our objective:

1 https://github.com/HobbitLong/RepDistiller

where α, β, η > 0 are balancing weights. For contrastive methods, we experiment on both CRD [Tian et al., 2019] and SSKD [Xu et al., 2020].

Results on CIFAR-100 dataset. Table 1 compares the performance of various distillation methods when teacher and student share similar architectural styles. We observe that PSKD outperforms KD and other baselines except for the contrastive distillation methods. The contrastive methods such as CRD [Tian et al., 2019], SSKD [Xu et al., 2020] and WCoRD [Chen et al., 2021] exhibit slightly better performance than PSKD as they transfer structural information of teacher network. We observe that using PSKD and contrastive methods together can further boost the performance. In addition to CRD and SSKD, model compression with PSKD achieves state-of-the-art performance in 5 out of 7 tasks, showing the effectiveness of PSKD.

Table 3 compares the performance when the teachers and students are of different architectures. Remark that the distillation methods using the information extracted from intermediate layers may perform worse when the architectures of teacher and student differ, but since PSKD only uses the output of a network, it doesn't rely on architecture-specific cues. Also, when using PSKD with contrastive methods, we achieve state-of-the-art results in 5 out of 6 tasks.

Results on ImageNet. Table 4 shows the results of model compression on ImageNet. Remark that our PSKD achieves lower error rates than KD and all other distillation methods except SSKD. Surprisingly, even though PSKD does not rely on transferring structural information, it achieves better performance than CRD and WCoRD. Meanwhile, PSKD enjoys simpler implementation and faster training.

6 Conclusion

We present pseudo-spherical knowledge distillation that generalizes knowledge distillation by using pseudo-spherical scoring rule. We propose two novel loss functions for pseudospherical knowledge distillation and provide an empirical and qualitative analysis. Our method achieves state-of-the-art results on CIFAR-100 and ImageNet supervised model compression, together with structural knowledge distillation methods.

Acknowledgments

This work was supported by the Agency for Defense Development by the Korean Government ("Adversarial AI based Deep Neural Nets Attack and Defense project")