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 a new knowledge distillation framework for directed graphical models based on the reparameterization trick. The new distillation framework overcomes the intractable marginals in marginalized distillation, and error accumulation in local distillation. Empirically, the new distillation framework surpasses the baselines on deep generative model compression, VAE continual learning, and discriminative DGMs compression.