A U NIFIED K NOWLEDGE D ISTILLATION F RAMEWORK FOR D EEP D IRECTED G RAPHICAL M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
Knowledge distillation (KD) is a technique that transfers the knowledge from a large teacher network to a small student network.
It has been widely applied to many different tasks, such as model compression and federated learning.
However, the existing KD methods fail to generalize to general deep directed graphical models (DGMs) with arbitrary layers of random variables.
We refer by deep DGMs to DGMs whose conditional distributions are parameterized by deep neural networks .
In this work, we propose a novel uniﬁed knowledge distillation framework for deep DGMs on various applications.
Speciﬁcally, we leverage the reparameterization trick to hide the intermediate latent variables, resulting in a compact
DGM.
Then we develop a surrogate distillation loss to reduce error accumulation through multiple layers of random variables.
Moreover, we present the connections between our method and some existing knowledge distillation approaches.
The proposed framework is evaluated on three applications: deep generative models compression, discriminative deep DGMs compression and VAE continual learning.
The results show that our distillation method outperforms the baselines in datafree compression of deep generative models, including variational autoencoder (VAE), variational recurrent neural networks (VRNN) and Helmholtz Machine (HM).
Moreover, our method achieves good performance for discriminative deep
DGMs compression.
Finally, we also demonstrate that it signiﬁcantly improves the continual learning performance of VAE.
I NTRODUCTION
Knowledge distillation (KD) aims at transferring the knowledge of a large teacher model to a small student model, which tries to mimic the behavior of the teacher model to attain a competitive or superior performance Hinton et al.
Gou et al.
The goal of this work is to develop a uniﬁed knowledge distillation (KD) framework for deep directed graphical models (DGMs).
Applications of the proposed framework include
: (i) data-free deep generative models compression, (ii) discriminative deep DGMs compression, and (iii) continual learning with VAEs
Deep directed graphical models (DGMs) refer to DGMs whose conditional distributions are parameterized by deep neural networks (DNNs), which is in contrast to the regular DGMs with tabular conditional probability.
One good example is variational autoencoders (VAEs), whose posterior probability of latent variables is parameterized by DNNs.
A general deep DGM may have complex structures, consisting of arbitrary number of input variables, target variables, and latent variables.
Deep DGMs have been widely used in various applications, such as image generation Vahdat &
Kautz (2020), text generation Bowman et al.
(2015), and video prediction Wu et al.
This work is motivated by the growing popularity of recent over-parameterized deep DGMs with millions of parameters to improve their accuracy in various tasks.
However, the large models are very computationally expensive.
As a result, it is not practical to deploy them on resource-constrained edge devices, such as mobile phones and IoT systems Li et al.
One possible solution to this problem is KD, which enables a smaller student model to approximate the performance of a large teacher.
Recently, KD has been widely applied to many different tasks, such as model compression Hinton et al.
(2015), continual learning Li & Hoiem (2017);
Zhai et al.
(2019), and federated learning Lin
et al.
(2020b)
; Li & Wang (2019).
To our knowledge, the existing KD methods, however, are only applicable to some speciﬁc DGMs, including generative adversarial networks (GANs) Aguinaldo z !
z " " ! " " # $ ! $ " z !
z " " ! "
" # z !
z " " ! " " # $ #,!
$ #," $ %,!
$ & $ %," " ! "
" # $ ! $ "
[Figure] [Figure]
et al.
; Li et al.
(2020), auto-regressive models in natural language processing (NLP)
Lin et al.
(2020a), and VQ-VAE Roy et al.
They fail to generalize to the general deep DGMs, especially to those with multiple latent variables or complex dependence structures , as illustrated in Fig.
Generalizing knowledge distillation to deep DGMs poses two major challenges.
First, distillation by marginalizing all latent variables is generally intractable.
Secondly, distilling each layer locally and independently may suffer from error accumulation, as shown Fig.
We can observe that the accumulated error (i.e., KL divergence) between the teacher and student grows linearly for local distillation.
To address these challenges, we propose a novel uniﬁed knowledge distillation framework for deep DGMs.
Speciﬁcally, we ﬁrst adopt the reparameterization trick Kingma (2013); Kingma et al.
(2014) to convert a DGM into a compact semi-auxiliary form .
By semi-auxiliary form , we mean the latent variables, z , in both the student and teacher models are converted to deterministic variables with auxiliary variables, while the input variables and target variables remain unchanged, as shown in Fig.
1 (c).
Note that different from the classical reparameterization for model training that requires continuous latent variables, ours can be applied to both continuous and discrete variables.
Then a surrogate distillation loss is derived as a new objective of KD.
To mitigate gradient vanishing, we further incorporate a latent distillation loss that penalizes the dissimilarity of latent variables between the teacher and student into our objective.
Moreover, we illustrate the connections between our approach and some existing KD methods for speciﬁc DGMs.
We demonstrate that our method is a proper generalization of these existing methods.
We apply our distillation method to three different applications: deep generative models compression Rezende et al.
(2014); Chung et al.
(2015), VAE continual learning Rao et al.
and discriminative deep DGMs compression Alemi et al. (2016).
For deep generative models compression, the student model distilled by our method in a data-free manner outperforms that trains from scratch and the other baselines.
In addition to generative modeling, our method is able to compress discriminative deep DGMs with high accuracy.
Finally, we also illustrate that it can better mitigate the catastrophic forgetting issue than the generative replay approach in continual learning.
In summary, our contributions include: 1) a new uniﬁed KD framework is proposed for general deep
DGMs based on reparameterization trick, 2) we derive a novel distillation loss that combines the latent distillation loss and surrogate distillation loss to improve the performance of KD, 3) evaluation results on multiple benchmark datasets show that our approach can not only achieve high accuracy for deep DGMs compression but also improve the continual learning performance of VAEs.
P RELIMINARIES
Directed Graphical Models (DGMs) DGMs such as Bayesian Networks (Darwiche, 2009) are an expressive class of probabilistic graphical models (Koller & Friedman, 2009), in which the joint distribution is factorized into the product of many conditional distributions according to a directed acyclic graph (DAG) that captures variable conditional dependencies.
In this work, we primarily study knowledge distillation for deep DGMs, especially for those with complicated dependency structures
For deep DGMs, we are interested in modeling the conditional distribution p θ ( y , z | x ) for target variables y and latent variables z given input variables x , parameterized by θ .
Speciﬁcally, when there is no input variables, i.e., x = ∅ , we actually model the joint distribution p θ ( y , z ) .
Let Pa ( · ) denote the parent random variables of a certain variable which is deﬁned by the DAG of a DGM.
Then the conditional distribution p θ ( y , z | x ) has its factorized form as follows,
[Equation]
where y j denotes the j th target variable in y , z i denotes the i th latent variable in z .
Without loss of generality, we assume for two variables y i and y j , if y j is an ancestor of y i , then it holds that i > j .
Knowledge Distillation (KD) KD aims to transfer the knowledge of a large teacher model to a smaller student model.
One commonly used vanilla distillation method is to encourage the student to mimic the output of the teacher model Hinton et al.
Given an empirical distribution of the training data p data ( x ) , its distillation loss is an expected dissimilarity measure between the output of the student and that of the teacher.
A general form of distillation loss is given by
[Equation]
where p φ ( y | x ) and p θ ( y | x ) denote the output conditional distribution of the teacher and student models, respectively.
φ and θ denote their corresponding parameters.
d ( · , · ) is a dissimilarity measure between two probability distributions.
Kullback-Leibler (KL) divergence (Hinton et al., 2015; Chen et al., 2017), for example, is one of some typical choices.
The above Eq.
(2) and its extended version have been applied to different DGMs, such as vanilla neural networks Chen et al.
(2017), GANs Aguinaldo et al.
(2019), and fully-visible auto-regressive models (e.g.
, Transformer) Kim & Rush (2016); Lin et al.
(2020a).
Take fully-visible auto-regressive models as an example.
When we set d ( · , · ) to KL divergence, Eq.
(2) can be factorized as
[Equation]
A tractable estimation to L kd above can be obtained by Monte Carlo method.
It can also be viewed as distilling each conditional distribution in a local and independent manner, as shown in Fig.
3 (d).
However, to our best knowledge, the existing KD approaches are only designed for some speciﬁc
DGMs.
They fail to be applied to a general DGM with multiple latent variables or complicated dependence structures.
Hence, the question is, how can we distill the knowledge from the teacher to the student given a general DGM structure?
Intuitively, there are two naive methods: (i) marginalized distillation : it marginalizes all latent variables to get p ( y | x ) =
(cid:82) p ( y , z | x )d z .
However, the integration (cid:82) p ( y , z | x ) d z is generally intractable and thus the loss as in Eq.
(2) is also intractable.
(ii) local distillation : it treats latent variables z equally as target variables and then distills each conditional distributions for both z and y locally and independently, as shown in Fig. 3(d).
The local distillation may suffer from error accumulation through multiple layers if each conditional distribution in the student slightly deviates from the teacher.
Fig.
2 shows a toy example that the accumulated error of local distillation, i.e., the KL-divergence between teacher and student, increases linearly as the number of layers of latent variables rises.
For detailed discussions, please refer to Appendix A.
Reparameterization Trick Reparameterization trick Kingma (2013); Kingma et al.
(2014), also called the auxiliary form of a DGM, is originally proposed to backpropagate through a random node that is not differentiable during training.
Its basic idea is introduced as follows.
Given a conditional distribution of random variable z i in a DGM, p ( z i | Pa ( z i ) , x ) , we convert it to a deterministic variable by adding an auxiliary variable (cid:15) i to its dependence, as shown in Figure 1 (c).
Here (cid:15) i is a root node of the DGM with an independent marginal distribution of p ( (cid:15) i ) .
By choosing appropriate p ( (cid:15) i ) and deterministic transformation g ( · ) Kingma (2013), we can have z i
= g ( Pa ( z i ) , x , (cid:15) i ) , where z i is determined by its parent variables Pa ( z i ) , input variables x , and the corresponding auxiliary variable (cid:15) i .
(cid:15) i serves as the source of stochasticity of z i .
[Figure] In this work, note that we do not primarily use reparameterization trick for model training.
Rather, we leverage it to convert the latent variables z in DGMs to deterministic variables so that we can effectively distill knowledge from a compact form of DGM.
Note that different from the classical reparameterization for model training that requires continuous latent variables, ours can be applied to a wider range of variables z , including both continuous and discrete variables.
Besides, the transformation function g ( · ) Kingma & Welling (2013) in our framework can be either differentiable or non-differentiable.
Hence, our method can be applied to much more DGMs than the classical one.
Below, we will elaborate this general idea in more detail.
M ODELING
In this section, we ﬁrst introduce the semi-auxiliary form of DGMs using reparameterization trick.
Then we propose a new surrogate loss function and latent distillation loss for our KD method.
3. S EMI - AUXILIARY F ORM
As discussed in Section 2, the two naive methods, marginalized distillation and local distillation, do not work well due to intractable distillation loss or error accumulation.
In order to address these issues, we propose a novel idea that converts DGM to its semi-auxiliary form based on reparameterization trick.
Speciﬁcally, we convert all the latent variables, z , in both the teacher and student to deterministic variables with auxiliary variables, while keeping target variables y and input variables x unchanged.
This is because our ultimate goal is to encourage student to mimic the output (target variable) of the teacher based on input variables.
Hence, we can omit the deterministic (latent) variables in a DGM, yielding a compact semi-auxiliary form that only consists of target variables, input variables and auxiliary variables.
In this way, each target variable has a tractable and direct dependence on input variables or prior target variables, as shown in Fig.
1 (c).
Fig.
1 illustrates a toy example of converting DGM to its semi-auxiliary form.
Fig.
1 (a) is the original form of a DGM, and its corresponding auxiliary form is shown in Fig.
1 (b).
In this paper, we only assign auxiliary variables to latent variables z , as shown in Fig.
1 (c).
We can observe from it that latent variables z 1 and z 2 are deterministic when their ancestors are given.
We thus can omit the deterministic variables, leading to a compact semi-auxiliary form, as shown in Fig.
1 (d).
3. S URROGATE D ISTILLATION L OSS
After obtaining the semi-auxiliary form of a DGM in Fig.
1 (d), our distillation loss becomes the dissimilarity between p φ ( y | (cid:15) , x ) and p θ ( y | (cid:15) , x ) .
We call it surrogate distillation loss .
Our goal is to minimize the surrogate distillation loss w.r.t.
student’s parameters θ below.
[Equation]
where (cid:15) denotes a set of auxiliary variables (cid:15) .
The expectation of dissimilarity is taken over both empirical data distribution p data ( x ) and auxiliary variable distribution p φ ( (cid:15) ) .
The expectation can be estimated using Monte Carlo method.
Note that p φ ( (cid:15) ) is generally chosen to be simple and ﬁxed distribution with no parameters, such as unit Gaussian or standard uniform distribution.
Thus, it implies that teacher’s p φ ( (cid:15) ) is equivalent to student’s p θ ( (cid:15) ) , so there is no need to distill p φ ( (cid:15) ) to p θ ( (cid:15) ) .
An illustration is given in Fig.
3(e).
Proposition 3.1.
The surrogate distillation loss as deﬁned in Eq.
(4) is an upper bound of the distillation loss as deﬁned in Eq.
(2) when the dissimilarity measure is chosen to be KL divergence.
We provide the detailed proof in Appendix B.
Next, we discuss the advantages of our method over the two naive methods mentioned above.
Firstly, the proposed method bypasses the intractable computation in marginalized distillation.
While marginalized distillation measures p ( y | x ) which is intractable in general, we instead measure p ( y | (cid:15) , x ) which is easily tractable by function composition with no need of integral or sum operation.
Here p ( y | (cid:15) , x ) is tractable because p ( y | (cid:15) , x ) = (cid:81) i p
( y i | (cid:15) ≤ i , y <
i , x ) .
Secondly, our method can make the DGMs shallower than local distillation, mitigating error accumulation through multiple layers of latent variables.
As illustrated in Fig. 3(d) and 3(e), we can see our method only constraints the target variables y in the student while local distillation constrains both latent variables z and target variables y .
As a result, local distillation may suffer from error accumulation issue.
Fig.
2 shows a toy example of comparing the accumulated error, i.e., KL divergence between the teacher and student, for our method and local distillation.
In this illustrative experiment, we let student mimic the output distribution of the teacher with L -layers latent variables for L ∈ 1 , . . .
Each layer of the teacher follows the Gaussian distribution p ( z i +
| z i ) =
N ( µ ( z i ) , 0 .
01 I ) , where µ ( z i ) is z 1 .
i for z
i ≥ 0 and − ( − z i ) 1 . 1 for z
i < 0 .
p ( z 1 ) is a uniform distribution U [ − 1 , 1] .
The student is parameterized by neural networks with proper residual structure He et al.
Then density ratio estimation Nguyen et al.
; Sugiyama et al.
(2012) is used to measure the KL divergence between the teacher and student.
We can observe from Fig.
2 that the accumulated error (KL divergence) grows linearly w.r.t the number of layers for local distillation because each layer of the student deviates from the teacher to an extent.
In contrast, the accumulated error (KL divergence) of our method increases slowly.
Furthermore, we show that our method can still mitigate error accumulation issue for a DGM whose latent variables are discrete in Appendix I.
3. L ATENT D ISTILLATION L OSS
The surrogate distillation loss in Eq.(4) can sufﬁciently achieve a satisfactory performance for knowledge distillation.
Nevertheless, there are still two limitations of the proposed method for some special DGMs.
Firstly, it fails to back-propagate the gradients when there exist discrete latent variables.
Secondly, it might suffer from gradient vanishing when the network structure of latent variables is very deep and complex.
To deal with these issues, we propose to penalize the dissimilarity of latent variables z in the teacher and student model.
The resulting latent distillation loss for latent variables z i is given by
[Equation]
where (cid:15) ≤ i is a set of all the ancestral auxiliary variables of z i .
z i is deterministic when (cid:15) ≤ i and x are given.
In order to better penalize the dissimilarity of latent variables z i , one good choice is to (only) convert the current z i back to its original form by removing (cid:15) i from its dependence.
Then we have the following latent distillation loss
[Equation]
The proposed latent distillation loss can beneﬁt our optimization process.
When latent variables are continuous, the latent distillation loss may provide shallower and supplementary supervisory signals to hasten the convergence.
When latent variables are discrete, it can deal with the back-propagation cutting off problem.
3. F INAL T ARGET -F REE L OSS
Combining the above surrogate distillation loss with latent distillation loss, our ﬁnal distillation loss is given by
[Equation]
where λ is a hyper-parameter that controls the importance of latent distillation loss.
Similar to Hinton et al.
(2015), our loss in Eq.
(7) is a target-free distillation loss, which means target data is not required for calculating it.
Eq.
(7) can also be applied to DGMs with no input variables (i.e. x = ∅ ).
In this case, Eq.
(7) can be computed in a completely data-free manner.
We summarize our KD method for a general DGM in Algorithm 1 in Appendix C.
3. C ONNECTIONS TO O THER D ISTILLATION M ETHODS
We present the connections between our method and some existing knowledge distillation methods.
Vanilla KD and Sequence-Level KD.
When there is no latent variable in a DGM, our distillation method in Eq.
(4) is naturally reduced to Eq.
(2) by removing the dependence on auxiliary variables (cid:15) , which is a typical vanilla KD Hinton et al. (2015).
Also, when the DGM is a fully visible autoregressive model, by removing dependence on (cid:15) and letting d ( · , · ) be KL divergence, Eq.
(4) can be reduced to Eq.
(3), which is the Monte Carlo approximation of intractable sequence-level distillation loss Kim & Rush (2016).
Please note that further simpliﬁcation has been made in Kim & Rush (2016) to enhance the practicability.
Feature Based KD.
Feature based knowledge distillation uses the intermediary representations of a teacher network to supervise a student network Romero et al.
When the teacher and student share the same size of latent features, the feature distillation loss can be written as
[Equation]
where f φ ( x ) and f θ ( x ) are intermediary deterministic features of the teacher and student, respectively.
r f ( · , · ) is a distance between two vectorized feature maps.
A vanilla neural network with multiple intermediate features can be viewed as a DGM with multiple deterministic latent variables.
Deterministic variables can be viewed as following the degenerate distributions.
In general, for p ≥ 1 , p-Wasserstein distance (inf E [ r ( a 1 , a 2 ) p ]) 1 p between two degenerate distributions located at a 1 and a 2 is equivalent to r ( a 1 , a 2 ) .
Thus, by viewing the intermediate features as latent variables following degenerate distributions, and choosing d ( · , · ) in our latent distillation loss L z,i as Wasserstein distance (inf E [ r f ( z φ , z θ ) p ]) 1 p , our latent distillation loss is reduced to feature distillation loss in Eq.
GAN Distillation.
GAN distillation in Li et al.
; Aguinaldo et al.
(2019) also incorporates the idea of feature distillation into their model, which is given by
[Equation]
The ﬁrst term above is the output distillation loss of the generator and the second one is the intermediary feature distillation loss.
Similar to feature distillation loss, by viewing the intermediate features and generator output as latent variables following degenerate distributions, and choosing d ( · , · ) in L sd and L z,i as p -Wasserstein distance (inf E [ r o ( y φ , y θ ) p ]) 1 p and (inf E [ r f ( z φ , z θ ) p ]) 1 p respectively, our ﬁnal distillation loss in Eq.
(7) can be reduced to GAN distillation loss as well.
3. A PPLICATIONS
We evaluate the performance of our method in three applications below.
• Data-Free Deep Generative Models Compression.
We ﬁrst apply our method to compress deep generative models in a data-free manner.
We conduct experiments on three representative models: VAE, VRNN, and Helmholtz Machine (HM) Chung et al.
; Saul et al.
Recent studies show that VAEs generate better with over-parametrization Vahdat & Kautz (2020).
However, it is challenging to deploy them to edge devices with limited computing resources.
In this work, we apply our method to compress a large 5-layer hierarchical VAE model Sønderby et al.
(2016) to smaller models, as illustrated in Fig.
6(a) in Appendix D.
Then it is applied to compress a sequence generative model with many latent variables, VRNN.
Finally, we adopt our method to compress the HM model with discrete variables.
The classical HM only has one target variable and each layer is parametrized by a linear transformation.
In order to demonstrate the applicability of our method, we extend it to a 5-layer deep HM with two targets, as shown in Fig.
6(d).
• VAE Continual Learning.
In addition to model compression, we evaluate the performance of our method on VAE continual learning.
Our goal is to model a new distribution while retaining the ability of modeling a learned old distribution without access to the old dataset.
Prior work on VAE continual learningYe & Bors (2020); Ramapuram et al.
(2020) mainly resort to generative replay strategy, in which a set of samples generated from VAE model learned on old datasets p (cid:48)
θ ( y ) are reused when learning a new dataset.
In fact, our experiments will show that generative replay is inferior to our distillation method for VAE continual learning because of the blurry nature of VAE generation.
• Discriminative Deep DGMs Compression.
Finally, we apply the proposed method to compress a 2-layer discriminative DGM in Alemi et al.
(2016), as shown in Fig. 6(e).
Incorporating our distillation loss into the original training loss, we have the following loss function.
[Equation]
E XPERIMENTS
In this section, we evaluate the performance of our method on three applications mentioned in Section 3.6.
We carry out extensive experiments on seven benchmark datasets: Old Faithful Geyser H ¨ ardle et al.
(1991), IAM online handwriting Liwicki & Bunke (2005), SVHN Netzer et al.
(2011), CIFAR10 Krizhevsky et al.
(2009), CelebA Liu et al.
MNIST LeCun et al.
(1998), and
Fashion-MNIST Sohn et al.
The detailed data splitting for training and testing is described in
Appendix E.
In addition, we present the detailed model conﬁgurations and hyperparameter settings in Appendix F. Source code will be publicly available upon publication.
4. E VALUATION ON D ATA - FREE D EEP G ENERATIVE M ODELS C OMPRESSION
We ﬁrst apply our distillation method to data-free compression of deep generative models, including
VAE, RVNN, and HM.
We compare the performance of the student model using our method and two baselines: training from scratch and local distillation Achille et al.
Frechet inception distance (FID) Heusel et al.
(2017) is adopted to evaluate their generative performance.
Data Free Hierarchical VAE Compression
We compress a large 5-layer hierarchical VAE Sønderby et al.
to a smaller model.
Experiments are conducted on three benchmark datasets: SVHN, CIFAR10, and CelebA. Fig.
4 illustrates the comparison of FID for different methods.
We can see that our method consistently outperforms the baselines for different student model sizes.
Importantly, our data-free distillation method outperforms student trained from scratch.
This is because directly optimizing a capacity-limited student VAE may not learn a decent model that can generate high-ﬁdelity samples.
Conversely, our method can help student VAE learn better performance from the teacher.
To demonstrate the similarity of the teacher and student, we further compare the FID between samples generated from the student and teacher for different methods, as shown in Fig.
7 in Appendix G.
We observe that our distillation method helps the student better approximate the teacher’s performance.
We also demonstrate that our method is not very sensitive to the hyper-parameter λ .
Please refer to
Appendix G.3 for more details.
Data-free VRNN Compression.
We then adopt our distillation method to compress a deep sequence generative model, VRNN.
Fig 5 shows the generated strokes by our method and the baselines.
It can be observed that our method can generate more readable and clearer strokes than the baselines.
The student VRNN distilled by our method in Fig 5 (c) has comparable generative performance to that of the teacher in Fig 5 (b).
Hence, we can conclude that the proposed method is able to achieve good performance for data-free VRNN compression.
Data-Free HM Compression.
In addition, our method is applied to compress Helmholtz Machine (HM) with discrete latent variables .
The experiment results, as shown in Fig.
9 in Appendix H, show that it can still achieve good performance as the teacher for HM compression.
For more details, please refer to the experimental results in Appendix H. [Figure] [Figure]
4. E VALUATION ON VAE C ONTINUAL L EARNING
We also evaluate the performance of the proposed method on VAE continual learning using CelebA dataset.
For each experiment, one of forty ground-truth attributes is selected.
Based on this attribute, we divide the whole dataset into two parts for continual learning.
Speciﬁcally, we ﬁrst learn to model one part of images, then we learn to model the other part while trying to retain the ability of modeling the ﬁrst part.
We repeat this experiment for four different attributes.
The single input variable x in this model is a binary variable used to indicate which part the generated images belong to.
We compare our method with local distillation and other two generative replay approaches: CURL Rao et al.
and LGM Ramapuram et al.
Table 1 illustrates the comparison of FID for different methods.
We can observe from Columns Old that the increased FID of our method is much smaller than that of baselines, which means our method suffers from the least catastrophic forgetting.
Besides, we can see from Columns New that our method achieves comparable FID to LGM and
CURL on new distribution, because these methods put only necessary constraints on the model to prevent catastrophic forgetting.
However, local distillation does not work well on new distributions because of its redundant constraints, leaving little free space for learning new distributions [Table] [Table]
[Table]
E VALUATION ON D ISCRIMINATIVE DGM S C OMPRESSION
Finally, we apply the proposed method to compress a discriminative 2-layer DGM with the same
MLP parameterization as Alemi et al.
The training loss in Alemi et al.
(2016) is denoted by
L VIB .
We compare the performance of different models using our method (our-VIB) in Eq.
(10) and the other ﬁve baselines: (i) combining L VIB with local distillation loss (Local-VIB), (ii) combining
L VIB with Monte Carlo based marginalized distillation loss (MC-VIB), (iii) model trained with L VIB only (VIB), (iv) model trained with vanilla cross entropy loss only (Vanilla), and (v) model trained using vanilla cross entropy loss with dropout (Dropout).
The experimental results are illustrated in Table 2.
We can see that our method beats the baselines in terms of classiﬁcation accuracy and negative log likelihood (NLL) on Fashion MNIST and SVHN datasets.
Besides, our method has lower ECE than all the baselines except for Dropout method.
[Table]
R ELATED W ORK
In the past few years, there is a line of work on knowledge distillation (KD) for different DGMs.
Most of existing approaches focus on vanilla DNNs or multi-target DNNs Hinton et al.
; Chen et al.
(2017); M ¨ uller et al.
; Romero et al.
; Zagoruyko & Komodakis (2016).
For instance,
Chen et al. Chen et al.
(2017) proposed a weighted cross entropy loss for multi-class object detection models using KD.
Zagoruyko et al.
M ¨ uller et al.
(2019) developed an attention based mechanism to transfer knowledge from the teacher CNN to the student for image recognition tasks.
These models can be viewed as DGMs with only one stochastic layer, consisting of one input variable and one or multiple conditionally independent target variables.
Some researchers also study KD for fully visible auto-regressive networks Kim & Rush (2016); Jiao et al.
; Oord et al.
; Huang et al.
Among them, sequence-level knowledge distillation (SeqKD) Kim & Rush (2016) is a promising strategy that supervises student with the teacher’s sequence distribution over the space of all possible sequences.
For example, some researchers Lin et al.
(2020a)
adopted imitation-based knowledge distillation (ImitKD) method to compress autoregressive models for language generation tasks.
Tan et al.
(2019) proposed a multilingual distillation framework for multilingual Neural Machine Translation (NMT).
These auto-regressive models are speciﬁc DGMs with multiple stochastic layers, but with no latent variables.
In addition, recent studies apply KD to compress deep generative models with one latent variable.
To reduce the number of parameters used in GANs, researchers Aguinaldo et al.
; Li et al.
(2020) devised new knowledge distillation methods for compressing GANs.
Besides, Lee et al.
(2020) distilled the learned representation from VAE models to GAN for high-ﬁdelity synthesis.
Roy et al.
(2018) leveraged VQ-VAE with KD to develop a non-autoregressive machine translation model.
However, these works did not apply reparameterization trick to deal with the latent variables.
In summary, the above existing KD methods are mainly focused on speciﬁc DGMs, but fail to generalize to the general deep DGMs, especially to those with multiple layers of random variables or complex dependence structures.
Different from prior work, we developed a novel uniﬁed KD framework for a general deep DGM using reparameterization trick.
C ONCLUSION
This paper proposed a new uniﬁed KD framework for deep directed graphical models (DGMs).
Speciﬁcally, we ﬁrst adopted reparametrization trick to convert latent variables into deterministic variables with auxiliary variables, resulting in a compact semi-auxiliary form of DGM.
Then a novel objective that combines the surrogate distillation loss and latent distillation loss was proposed to improve the performance of KD.
We further illustrated that our framework is a proper generalization of some existing KD methods.
We evaluated the performance of our method on different tasks.
The results showed that it can better compress VAE models in a data-free manner than the baselines.
It can also achieve high accuracy for discriminative DGMs compression and DGMs with discrete variables.
R EPRODUCIBILITY S TATEMENT
To replicate our experimental results in Section we have uploaded the source code on GitHub ( https://github.com/papersourcecode/ICLR2022-SourceCode.git ).
In our experiments, we adopt seven benchmark datasets: Old Faithful Geyser, IAM online handwriting,
SVHN, CIFAR10, CelebA, MNIST and Fashion-MNIST.
The detailed data splitting for training and testing is described in Appendix E.
In addition, we present the detailed model conﬁgurations and hyperparameter settings in Appendix F. For Proposition 3.
we provide the detailed proof in
Appendix B.
A T WO N AIVE M ETHODS
We present two naive methods for distilling knowledge from a teacher to a student for a general DGM as follows.
A. M ARGINALIZED D ISTILLATION
To minimize the desired distillation loss in Eq.
(2), one candidate solution is to marginalize over all latent variables, which, however, is generally intractable in DGMs with latent variables.
This is because in general there is no analytical solution for the likelihood p ( y | x ) =
(cid:82) p ( y , z | x )d z , much less for d ( p φ ( y | x ) , p θ ( y | x )) .
Thus, we try to use Monte Carlo method to estimate the conditional distribution p ( y | x ) as below
[Equation]
where { z ( n ) } Nn =1 is a set of samples drawn from p ( z | x ) .
However, there are two main issues of this method.
First, when y consists of more than one random variables (e.g., in deep generative modeling) or it is a continuous variable (e.g, auto-regressive models), p ( y | x ) is often not a simple distribution.
As a result, Eq.
(2) generally does not have a closed-form solution.
Besides, it is not often scalable to the scenario where z contains many latent variables or each latent variable is a high-dimensional vector, since the variance of Monte Carlo estimation may be increased.
The two limitations discussed above suggest that marginalized distillation is only applicable to singlelabel classiﬁcation tasks, where latent variables in the deep DGMs must be very low dimensional.
A. L OCAL D ISTILLATION
Another naive method is called local distillation, which applies distillation in a layer-wise manner in deep DGMs, as shown in Fig. 3(d).
In local distribution, each of student’s local component is encouraged to mimic its corresponding teacher’s local component in an independent manner.
With this method, we can factorize the joint distribution of a deep DGM into a product of a series of conditional distributions, as shown in Eq.
(1), so that the distillation loss of each conditional distribution can be calculated separately.
Formally, the distillation loss for target variable y j is given by
[Equation]
Similarly, the distillation loss for latent variable z
i is
[Equation]
For the above equations, all the expectations can be estimated using Monte Carlo method, and Monte
Carlo samples can be drawn using ancestral sampling from the teacher model.
Finally, we sum up all the local distillation losses and then jointly minimize the overall distillation loss as follows.
[Equation]
Let d ( · , · ) be KL divergence, then Eq.
(14) can be simpliﬁed as
[Equation] [Equation]
The above equation is the KL divergence between the joint probability on all the latent variables and target variables of the teacher and that of the student model.
It can be easily proved by factorizing p ( y , z | x ) .
For the above Eq.
(15), if there are no latent variables in a DGM, and if we choose d ( · , · ) to be
KL divergence, the local distillation loss can be simply reduced to the distillation loss in Eq.
(2) or can be expanded to Eq.
However, when there exist latent variables, local distillation will add extra constraints on latent variables, which is redundant and dampens the distillation performance in practice.
This is because we are only interested in minimizing the dissimilarity between target variables rather than latent variables.
Next, we provide an insight of why local distillation does not work well in deep DGMs.
Take the DGM in Fig. 3(d) as an illustrative example.
After distillation, each conditional distribution in the student model may still slightly deviate from the teacher model because of the inferior capacity of the student model or the randomness of optimization process.
We call this deviation from student to teacher imitation error .
Imitation errors will accumulate through multiple layers.
For instance, in the imitation process of variable y 1 , its local distillation loss in Eq.
(12) is L y, 1 = E p data ( x ) p φ ( z 1 | x )
[ d ( p φ ( y 1 | z 1 ) , p θ ( y 1 | z 1 ))] .
Even if this local distillation loss for y 1 is very well optimized, when the student’s latent distribution p θ ( z 1 | x ) deviates far from the teacher’s distribution, p φ ( z 1 | x ) , there will be no guarantee that p θ ( y 1 | x ) is a good approximation for p φ ( y 1 | x ) because p ( y 1 | x ) =
E p ( z 1 | x ) p ( y 1 | z 1 ) .
p θ ( y 1 | z 1 ) is only well constrained in the domain of p φ ( z 1 | x ) .
As a result, imitation errors may accumulate through layers of latent variables in a deep DGM.
We also conduct a toy experiment to show the error accumulation issue in local distillation.
In our experiment, we let student mimic the output distribution of the teacher with L -layers latent variables for L ∈ 1 , . . .
Each layer of the teacher follows the Gaussian distribution p ( z i +
| z i ) =
N ( µ ( z i ) , 0 .
01 I ) , where µ ( z i ) is z 1 .
i for z
i ≥ 0 and − ( − z i ) 1 . 1 for z
i < 0 .
p ( z 1 ) is a uniform distribution U [ − 1 , 1] .
The student is parameterized by neural networks with proper residual structure
He et al.
Then density ratio estimation Nguyen et al.
; Sugiyama et al.
(2012) is used to measure the KL divergence between the teacher and student.
We can observe from Fig.
2 that the accumulated error (KL divergence) grows linearly w.r.t the number of layers for local distillation because each layer of the student deviates from the teacher to an extent.
In contrast, the accumulated error (KL divergence) of our method increases slowly.
B A BOUT U PPER B OUND OF V ANILLA D ISTILLATION L OSS
We present the detailed proof of Proposition 3.1 below.
Proof.
Since the distributions of the auxiliary variables for both student and teacher are the same, i.e., p φ ( (cid:15) )
=
p θ ( (cid:15) ) , the surrogate distillation loss as in Eq.
(4) can be rewritten as follows.
[Equation]
Notice that p φ ( y | (cid:15) , x ) p φ ( (cid:15) )
=
p φ ( y , (cid:15) | x ) =
p φ ( (cid:15) | y , x ) p φ ( y | x ) .
Similarly, it also holds that p θ ( y | (cid:15) , x ) p θ ( (cid:15) )
=
p θ ( (cid:15) | y , x ) p θ ( y | x ) .
Thus, we can further rewrite the surrogate distillation loss
L sd as below.
[Equation]
Since the ﬁrst term is also non-negative [Equation] it further holds that
[Equation]
which ﬁnishes our proof
[Equation]
C A LGORITHM S UMMARY
We summarize our knowledge distillation method for a general DGM in Algorithm 1.
Please note that we use v i to denote a non-root random variable in DGMs, which is either a latent variable z i or a target variable y i .
v φ,i and v θ,i denote a sampling point of node v i in the teacher DGM and the student, respectively.
It can be observed from the distillation process that we only need to traverse the graph once to compute both the surrogate distillation loss and the latent distillation losses.
Line 9 is our latent distillation loss.
In Line 10-12, we describe the latent variables of the teacher and student model that are shared indirectly via sharing the auxiliary variables (cid:15) .
Line 14 computes the surrogate distillation loss.
Line 15-16 describe that the target variables of teacher and student model are shared directly.
Line 19 presents the back propagation and gradient descent.
Require: teacher DGM G φ , student DGM G θ , learning rate η , hyperparameter λ , empirical distribution of input dataset p data ( x ) (only required when x (cid:54) = ∅ )
Ensure: distilled student DGM G θ 1: while θ not converged do 2: sample mini-batched x from p data ( x ) //
(only required when x (cid:54) = ∅ ) 3: set L our =4: while G not fully traversed do 5: grab variable v i where Pa ( v i ) have all been sampled 6: set p φ,i = p φ ( v i | Pa ( v i ) , x ) 7: set p θ,i = p θ ( v i | Pa ( v i ) , x ) 8: if v i is latent variable then 9:
L our ← L our + λd ( p φ,i , p θ,i ) 10: sample (cid:15)
i from p φ ( (cid:15) i ) 11: sample v φ,i from p φ,i according to (cid:15) i // set v φ,i =
g φ ( Pa ( v i ) , x , (cid:15) i ) 12: sample v θ,i from p θ,i according to (cid:15) i // set v θ,i =
g θ
( Pa ( v i ) , x , (cid:15) i ) 13: else { v i is target variable } 14:
L our ← L our + d ( p φ,i , p θ,i ) 15: sample v φ,i from p φ,i 16: v θ,i ← v φ,i 17: end if 18: end while 19: θ ← θ − η ∂ L our ∂θ 20: end while
D I LLUSTRATION OF G RAPHICAL M ODELS FOR D IFFERENT A PPLICATIONS
Fig.
6 shows different forms of deep DGMs used in our experiments.
[Figure]
E D ATASETS
To evaluate the performance of the proposed distillation method, we conduct experiments on the following benchmark datasets.
• Old Faithful Geyser H¨ardle et al.
(1991): this is 2-dimensional Geyser data.
We use all the 272 samples for training.
• IAM Online Handwriting Liwicki & Bunke (2005): it has 13040 samples and each sample is a sequence of ( x, y ) coordinate together with binary indicator of pen up/down.
We use all the examples for training.
• SVHN Netzer et al.
This dataset is about house numbers in street views withcategorical label of 10 classes.
We use 73257 images for training and 26032 images for testing respectively.
• CIFAR
Krizhevsky et al.
This dataset consists of 60000 natural world images, with 1 categorical label of 10 classes.
The data is split into 50000 and 10000 images for training and testing.
• CelebA Liu et al.
(2015): It contains 182 , 732 RGB 128 × 128 × 3 images of celebrity faces.
We use 162770 images for training and 19962 images for testing.
• MNIST LeCun et al.
This dataset has 70000 images about grey-scale digits.
We use 60000 images for training and 10000 images for testing.
• Fashion-MNIST Sohn et al.
(2015): This dataset is about Grey-scale dressing images.
We use 60000 images for training and 10000 images for testing in our experiment.
F M ODEL C ONFIGURATIONS AND HYPERPARAMETER SETTINGS
We summarize the detailed model conﬁgurations and hyperparameter settings for the proposed method in different tasks below.
For all the tasks, we choose d ( · , · ) in the surrogate distillation loss L sd to be KL divergence.
When latent variables are discrete, d ( · , · ) in the latent distillation loss L z,i is chosen to be KL divergence and λ is set to 1.
When latent variables are continuous, we choose d ( · , · ) in L z,i to be squared 2-Wasserstein distance because it has analytical solution on two Gaussian distributions.
In addition, we normalize L z,i w.r.t.
the size of vector z i .
We empirically found that normalized latent distillation loss works slightly better than unnormalized one.
In all the following experiments, λ is chosen after a simple grid search.
The main training conﬁgurations for different tasks are summarized in the Table 3 below.
[Table]
F. D ATA - FREE VAE C OMPRESSION
In our experiments
[Table]
we distill a VAE with the same DGM and parameterization as the Ladder VAE
Sønderby et al.
The DGM structure of its generative model (decoder) is shown in Fig. 6(a).
Slightly different from Sønderby et al.
, we use Convolutional Neural Network (CNN) to parameterize both the bottom-up and top-down network in VAE, instead of simple MLP.
In addition, we adopt spatial latent variables (latent variables are spatially arranged, as the feature map in CNN) in the neural networks.
Each CNN is composed of consecutive residual blocks He et al.
We ﬁx the size of the teacher model to 5.4M parameters, with almost symmetric structure on the bottom-up network and top-down network.
To demonstrate the robustness of our method, we vary the student model size with different compression times.
In the experiments, we set the depth of latent variables to 32 for both the teacher and student VAE.
The hidden dimension of residual blocks varies from 4 to 64 in the student model, and is set tofor the teacher VAE.
Mini-batch sizes are set to 256 for SVHN and CIFAR10 , and 128 for CelebA.
When VAE is trained from scratch, we adopt linear β annealing, which is called warm up in Sønderby et al.
(2016) for ﬁrst half of the iterations to stabilize model learning.
For CelebA dataset, we resize the images to the shape 64 × 64 × 3 .
Model training is stopped after 57k, 46k, 63k iterations for
SVHN, CIFAR10, and CelebA, respectively.
F. VAE C ONTINUAL L EARNING
In our experiment, we still use the same 5-layer VAE as the above VAE compression experiment.
Following Rao et al.
(2019), we add an additional discrete input variable to the graph for the sake of indicating where the data comes from.
As a result, our VAE model is turned to a 6-layer Conditional
VAE, as illustrated in Fig.
6(d).
For VAE continual learning, we have the following distillation loss.
[Equation]
(16) where c is a constant independent of θ .
p old ( y ) and p new ( y ) are the empirical distribution of old and new dataset, respectively.
For the above loss, since p old ( y ) is not accessible, we use the old learned VAE’s distribution p (cid:48) θ ( y ) as a replacement.
Accordingly, we replace KL [ p old ( y ) (cid:107) p θ ( y )] with KL [ p (cid:48) θ ( y ) (cid:107) p θ ( y )] .
This results in our desired distillation loss, transferring knowledge from teacher p (cid:48) θ ( y ) to student p θ ( y ) , which, however, is intractable as discussed earlier.
To tackle this issue, we use our distillation loss L our instead.
As a result, the total loss for continual learning is rewritten as
[Equation] [Equation]
In this experiment, we follow the model conﬁgurations as the above VAE compression, except that we resize images to 32 × 32 × 3 and do not use running statistics for batch normalization layers during evaluation.
F. D ATA - FREE VRNN C OMPRESSION
In this experiment, we build and distill a VRNN with the same DGM and parameterization as in
Fraccaro et al.
(2016), as illustrated in Fig 6(c).
VRNN consists of as many latent variables and target variable as its sequence length.
We ﬁx the number of hidden layer in all densely connected networks to 1, and ﬁx the RNN to be 1-layer LSTM as in Fraccaro et al.
The dimension of latent variables z is set to 16.
For the teacher model, the hidden dimension of LSTM is set to 1200.
Feature size of target variables, latent variables, and hidden dimension of all densely connect networks are set to 512.
For the student model, the hidden dimension of LSTM is set to 600.
Feature dimension of target/latent variables, and hidden dimension of all densely connect networks are set to 128.
F. D ATA - FREE HM C OMPRESSION
Helmholtz Machine (HM) Neal (1992); Saul et al. (1996) is a classical generative modeling technologyGregor et al.
(2014); Bornschein & Bengio (2014); V´ertes & Sahani (2018); Bornschein et al.
(2016) that consists of: 1) Sigmoid Belief Network (SBN) as generative graphical model; 2) another
SBN as amortized inference network; 3) Wake-Sleep learning strategy.
The classical HM has only one target variable and each layer is parameterized by a linear transformation.
In our experiment, in order to demonstrate the applicability of our distillation method, we extend it to a 5-layer DGM with two target variables, as shown in Fig.
6(d).
The inference network is also modiﬁed to match the true posterior dependency structure.
Following the similar idea in Krishnan et al.
and
Bornschein & Bengio (2014), each conditional distribution in both generative and inference network are parameterized by 2-layer MLP instead of linear transformation, which turns our HM to deep HM .
In our experiment, we ﬁx latent variables to two-dimensional vectors with binary elements.
The hidden dimension of MLP is set to 8 and 2 for teacher and student model respectively.
F. D ISCRIMINATIVE D EEP DGM C OMPRESSION
In this experiment, we distill a Discriminative Deep DGM with the same 3-layer MLP parameterization as Alemi et al.
We ﬁx the hyperparameter β in original training loss to be 1e-3.
Latent variable is set to be 256 dimensional vector.
The hidden size of MLP is set to 1024 and 256 for the teacher and the student model, respectively.
Models in our experiment are trained for 46 k iterations for MNIST, Fashion-MNIST dataset, and are trained for 57 k iterations for SVHN dataset.
Images in
SVHN dataset are ﬁrst turned to grayscale.
G E XPERIMENTS ON VAE C OMPRESSION G. S IMILARITY BETWEEN T EACHER AND S TUDENT
Fig.
7 compares the FID between samples generated from the student and teacher for different methods.
We can observe that our method beats the baselines: local distillation and trained from scratch 0[ [Equation]
    ) , ' RXUGLVWLOODWLRQORFDOGLVWLOODWLRQWUDLQVFUDWFK (a) CelebA 0
[ 0[ 0[ 0[ VWXGHQWPRGHOVL]H      ) , ' RXUGLVWLOODWLRQORFDOGLVWLOODWLRQWUDLQVFUDWFK (
0
[ 0[ 0[ 0[ VWXGHQWPRGHOVL]H     ) , ' RXUGLVWLOODWLRQORFDOGLVWLOODWLRQWUDLQVFUDWFK (c) SVHN
G. Q UALITATIVE R ESULTS ON C ELEB A
We also conduct experiments to compare the generated samples by different knowledge distillation (KD) methods on the CelebA dataset, as illustrated in Fig.
It can be observed from Fig. 8(c) that images generated by our method are high-ﬁdelity, as good as ground truth and those generated by the teacher, which means our method can imitate teacher’s performance very well.
However, a limited-size VAE model is incapable of learning such a complicated distribution from scratch on its own. Hence, the images generated by the student VAE trained from scratch are low quality, as illustrated in Fig. 8(d). G. H YPERPARAMETER S TABILITY IN THE L OSS We further demonstrate the hyper-parameter stability of our method.
Fig.
9 shows the relationship between λ and FID score.
It can be observed that as λ varies from to 0.0 to 1 e 3 , the FID score of our distillation method does not vary too much.
It means that our method is stable and robust to the hyper-parameter λ .
H E VALUATION ON D ATA - FREE H ELMHOLTZ M ACHINE C OMPRESSION
We apply the proposed approach to compress a 5-layer HM trained with Wake-Sleep algorithm to a smaller HM in a data-free manner.
We still compare it with the baselines: training from scratch and [Figure]
H H H H H  ODPEGD  
 

  ) , '
[Figure]
local distillation.
When a model is trained from scratch instead of distillation, we add an inference model (encoder) with symmetric structure and almost the same number of parameters as its generative model.
The experimental results are shown in Fig.
We can observe that our distillation method still outperforms the student model trained from scratch, just as VAE and VRNN compression.
In addition, it shows that optimizing a capacity-limited HM may not learn a powerful model that can generate high-ﬁdelity samples.
In contrast, our method enables a small student HM to attain good performance as a larger teacher model.
[Figure]
In our experimental results, we ﬁnd that there is tiny difference between our method and local distillation, partly because this 2-dimensional modeling task is too easy for highly ﬂexible HMs to learn.
Both our method and local distillation can converge after a few iterations. Hence, they have comparable performance to each other. I S CALABILITY ON L ARGE N UMBER OF D ISCRETE V ARIABLES Our proposed method can also be scalable to a DGM with a large number of discrete latent variables.
We take a DGM with L (e.g, 20) layers of discrete latent variables as an example to demonstrate the good performance of our approach, as shown in Fig.
In our experiment, we let student mimic the output distribution of the teacher with L -layers latent variables for L ∈ 1 , . . .
Each layer of the teacher follows the binomial distribution p ( z i +1 | z i ) = B (1000 , p ( z i )) , where p ( z i ) = ( z i 1000 ) 1 . 2 ; p ( z 1 ) is a uniform distribution U [0 , 1] .
The student is parameterized by neural networks with a residual structure He et al.
Then density ratio estimation Nguyen et al.
; Sugiyama et al.
(2012) is used to measure the KL divergence between the teacher and the student.
We can observe from Fig.
11 that even if all the latent variables are discrete, our method still has the capability of mitigating the imitation error occurred in the shallow layers and achieves a very good imitation performance.
For the local distillation method, however, its imitation errors are accumulated as the number of layers increases.