A Unified Knowledge Distillation Framework for Deep Directed Graphical Models

1 INTRODUCTION . 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 . ( 2015 ) ; Gou et al . ( 2021 ) . The goal of this work is to develop a unified 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 . ( 2021 ) . 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 . ( 2020 ) . 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 specific DGMs , including generative adversarial networks ( GANs ) Aguinaldo z ! z '' '' ! '' '' # $ ! $ '' z ! z '' '' ! '' '' # z ! z '' '' ! '' '' # $ # , ! $ # , '' $ % , ! $ & $ % , '' '' ! '' '' # $ ! $ '' ( a ) ( b ) ( c ) ( d ) Figure 1 : Illustration of DGM in four different forms . Diamonds are deterministic variables and circles are random variables . ( a ) Original form ; ( b ) Auxiliary form ; ( c ) Our semi-auxiliary form ; ( d ) Compact semi-auxiliary form . et al . ( 2019 ) ; Li et al . ( 2020 ) , auto-regressive models in natural language processing ( NLP ) Lin et al . ( 2020a ) , and VQ-VAE Roy et al . ( 2018 ) . They fail to generalize to the general deep DGMs , especially to those with multiple latent variables or complex dependence structures , as illustrated in Fig . 1 . 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 . 2 . 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 unified knowledge distillation framework for deep DGMs . Specifically , we first 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 specific 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 . ( 2019 ) , 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 unified 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 . 2 PRELIMINARIES . 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 θ . Specifically , 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 defined by the DAG of a DGM . Then the conditional distribution pθ ( y , z|x ) has its factorized form as follows , pθ ( y , z|x ) = ∏ j pθ ( yj |Pa ( yj ) , x ) ∏ i pθ ( zi|Pa ( zi ) , x ) , ( 1 ) where yj denotes the jth target variable in y , zi denotes the ith latent variable in z . Without loss of generality , we assume for two variables yi and yj , if yj is an ancestor of yi , 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 . ( 2015 ) . Given an empirical distribution of the training data pdata ( 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 Lkd = Epdata ( x ) [ d ( pφ ( y|x ) , pθ ( y|x ) ) ] , ( 2 ) 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 Lkd = Epdata ( x ) ∑ j Epφ ( y < j |x ) [ KL ( pφ ( yj |y < j , x ) ‖ pθ ( yj |y < j , x ) ] ) . ( 3 ) A tractable estimation to Lkd 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 specific 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 ) = ∫ p ( y , z|x ) dz . However , the integration ∫ p ( y , z|x ) dz 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 zi in a DGM , p ( zi|Pa ( zi ) , x ) , we convert it to a deterministic variable by adding an auxiliary variable i to its dependence , as shown in Figure 1 ( c ) . Here i is a root node of the DGM with an independent marginal distribution of p ( i ) . By choosing appropriate p ( i ) and deterministic transformation g ( · ) Kingma ( 2013 ) , we can have zi = g ( Pa ( zi ) , x , i ) , where zi is determined by its parent variables Pa ( zi ) , input variables x , and the corresponding auxiliary variable i. i serves as the source of stochasticity of zi . 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 .

Summary: This paper proposes to use the reparameterization trick to convert the latent variables in DGMs to deterministic variables in the context of KD. It then proposes a surrogate distillation loss and latent distillation loss and evaluates the performance of the proposed method in three applications. Experimental results confirm the effectiveness of the proposed model.