Information-Theoretic Generative Model Compression with Variational Energy-based Model
Anonymous Author(s)
AfﬁliationAddress email
Abstract
We propose an information-theoretic knowledge distillation approach for the compression of generative adversarial networks, which aims to maximize the mutualinformation between the teacher and student networks via a variational optimization based on an energy-based model.
Because the direct computation of themutual information in continuous domains is intractable, our approach alternativelyoptimizes the student network by maximizing the variational lower bound on themutual information.
To achieve a tight lower bound, we introduce an energy-basedmodel relying on a deep neural network to represent a ﬂexible variational distribution that deals with high-dimensional images and consider spatial dependenciesbetween pixels, effectively.
Since the proposed method is a generic optimizationalgorithm, it can be conveniently incorporated into arbitrary generative adversarialnetworks and even dense prediction networks, e.g., image enhancement models.
We demonstrate that the proposed algorithm achieves outstanding performance inmodel compression of generative adversarial networks consistently when combinedwith several existing models.
Introduction
Generative adversarial networks (GANs) [1] have accomplished impressive achievements on variouscomputer vision tasks including image synthesis [2, 3, 4, 5], image-to-image translation [6, 7, 8, 9],video prediction [10, 11, 12] and many others.
However, despite the outstanding performance, theirapplicability to resource-hungry systems, e g ., , edge or mobile devices is limited due to their hugecomputational and energy costs.
To deal with the issues, GAN compression, a task reducing the sizesof generator networks, has recently drawn attention to many machine learning researchers.
The straightforward way to compress generators is to directly employ traditional network compressiontechniques such as low-rank approximation [13, 14, 15], network quantization [16, 17, 18, 19],pruning [20, 21, 22, 23, 24] and knowledge distillation [25, 26, 27, 28, 29].
However, the compressionof generative models typically deals with a higher-dimensional output space and requires to match theoutput distributions between the original and compressed models instead of the output instance pairgiven by the two models.
Hence, as reported in [30, 31, 32], the synthesized images obtained from thecompressed generators via the channel pruning [33, 22, 23] are often suboptimal—producing imageswith artifacts and suffering from image quality degradation, which implies that naïve application ofthe methods designed for discriminative models is not effective for GAN compressions.
To deal with the issue, existing approaches for compressing GANs [30, 34, 31, 35, 36, 37, 32, 38]utilize knowledge distillation to achieve competitive accuracy of a student model.
Speciﬁcally,these methods transfer the knowledge of a teacher network to a student model via intermediaterepresentations, ﬁnal outputs, and embedded features [39, 40].
For the purpose, the existing algorithms employ (cid:96) 1 or (cid:96) 2 norms to compute the distance between the representations of the two networks, butthey are not necessarily effective for minimizing the distance between the two distributions.
We propose to compress GAN models by maximizing the mutual information for knowledge distillation between the outputs of a teacher and a students.
For the compression of generative models,the use of the mutual information between two random variables makes more sense than the directcomparison between two outputs, which is the standard way to transfer knowledge from a teacher toa student in discriminative models.
However, because the computation of the mutual informationin continuous domain is not tractable in general, we employ a variational approach [41], where thelower bound of the mutual information is alternatively optimized using a variational distribution, anapproximation for the conditional distribution of the teacher output given the student output.
Wefurther adopt an energy-based model for the ﬂexible representation of the variational distribution andtighten the lower bound for the effective knowledge transfer.
Note that the proposed energy-basedmodel is convenient to deal with high-dimensional data (images) and effective to handle spatialdependency between pixels within an image.
Combined with the existing approaches, our frameworkachieves substantial performance improvement on GAN compression.
The main contributions of ourwork are summarized below:
• We propose a novel energy-based framework of GAN compression via knowledge distillation, which maximizes mutual information between a teacher and a student network.•
The proposed energy-based approach is effective to represent a ﬂexible variational distribu- tion and can be incorporated into many GAN compression methods due to its generality.
• Experimental results show that the proposed method can improve the performance when combined with the existing methods.
The rest of the paper is organized as follows.
Section 2 discusses the related work to GAN compression.
The details of our approach is described in Section 3 and 4, and the experimental results arepresented in Section 5.
Finally, we conclude this paper in Section 6.
Related Work
2.1 Generative Adversarial Networks
Generative Adversarial Networks (GANs) [1] consist of two components, a generator and a discriminator, and train the models using an adversarial loss; the generator aims to fool the discriminator bysynthesizing realistic images from random noise samples while the discriminator wants to distinguishthe real data from the fake ones.
Unlike the vanilla GANs, conditional GANs [42, 6, 7] can controlimage generation conditioned on inputs such as a class label [42] or an image translated into anotherimage [6, 7] ( e . g ., Horse → Zebra).
The conditional GANs are not conﬁned to the conditional imagegeneration tasks, but applied to various areas such as super-resolution [43], text-to-image [44], andimage inpainting [45].
Since the conditional GAN models are typically large in terms of FLOPs, theproposed GAN compression algorithm is mainly tested on them but we also verify the effectivenessof our approach on the unconditional counterpart.
2.2 Knowledge Distillation
Knowledge distillation techniques aim to transfer knowledge from a powerful network ( i .
e ., teacher)to a compact network ( i .
e ., student) by matching logits [46], output distributions [25], intermediateactivations [47], or attention maps [26].
Speciﬁcally, Hinton et al .
[25] learn the representation ofstudent by minimizing the relative entropy from its normalized softmax output to that of teacher while[46, 47] minimize the mean squared error between the representations of the teacher student networks.
Both Variational Information Distillation (VID) [48] and Contrastive Representation Distillation(CRD) [49] employ the mutual information between teacher and student to estimate the similarityof the representations from the two networks.
The former relies on a variational approach [41] forthe tractable computation of mutual information between the two networks.
On the other hand, thelatter derives an InfoNCE-like loss function from the lower bound of the mutual information to learnthe representation of a student.
However, VID maximizes the variational lower bound using a strongassumption about the variational distribution, a fully factorized Gaussian distribution, and CRD is to be trivial in image generation tasks without semantic attribute annotations, e .
g ., class labels.
Ouralgorithm is related to VID and CRD since they all attempt to maximize the mutual informationbetween teacher and student networks, but has clear advantage over them in terms of the ﬂexibility ofthe variational distribution and the applicability to generative models.
2.3 GAN Compression
GAN compression methods typically aim to reduce the size of the generator without considering thesize of the discriminator since it is not used for inference.
Co-Evolution [30] utilizes an evolutionaryalgorithm [50] to prune the channel size of CycleGAN based on a ﬁtness formulated based on FLOPsand the reconstruction error between the original and pruned generators.
GAN-Compression [34]employs a once-for-all network training scheme [51] combined with the reconstruction loss as wellas the intermediate feature distillation loss [47], and selects the best performing subnetwork giventhe target budget.
Compression and Teaching (CAT)
[37] designs an efﬁcient building block in thegenerator, and prunes the network using the magnitude-based pruning method [23] while maximizingthe similarity between the two networks based on centered kernel alignment [52].
Content-Aware GAN Compression (CAGC)
[32] uses the feature distillation and reconstruction loss based only on theobject part of a generated image while additionally considering LPIPS [40] as a perceptual distancefor the synthesized ones.
On the contrary, Online Multi-Granularity Distillation (OMGD) [38]iteratively trains a teacher network with original conditional GAN objectives and a student networkusing the total variation loss for spatial smoothness and a distillation loss that makes the studentmimic the ﬁnal outputs and channel-wise attention weights for the intermediate feature maps of theteacher networks.
Generator-discriminator Cooperative Compression (GCC) [53] also adopts theonline distillation scheme while selecting channels for the original discriminator to match the reducedcapacity of the student generator at each iteration.
To validate the effectiveness of our framework, wepresent how the proposed algorithm can improve performance via the combinations with existingmethods such as OMGD, GCC, GAN-Compression, CAT, and CAGC.
Variational GAN Compression with Energy-Based Model
Our goal is to compress generators in GANs through knowledge distillation.
Let S and T be randomvariables of the outputs from a student, G s ( · ; φ s ) :
X → Y and a teacher, G t ( · ; φ t ) :
X → Y ,respectively, where φ s and φ t are their model parameters.
For knowledge transfer from the teacher tothe student, we aim to maximize mutual information between T and S , I ( T ; S ) , which is given by
[Equation]
where D KL ( ·||· ) denotes the Kullback-Leibler (KL) divergence and H ( · ) means the entropy function.
Note that all the expectations are taken over the joint data distribution, p ( t , s ) , where t and s indicateteacher and student output distributions, respectively.
The practical challenge in this objective is that the computation of the mutual information formultivariate random variables is typically intractable.
To tackle the issue, a variational method [41]derives the lower bound on the mutual information, ˜ I
( T ; S ) , using a variational distribution, q ( t | s ) ,which is given by
[Equation]
To maximize the mutual information, we aim to increase the variational lower-bound by minimizingthe KL-divergence between the true and variational distributions, which is deﬁned as
[Equation]
Variational Information Distillation (VID) [48] optimizes the variational lower bound by employingthe fully factorized Gaussian distribution as the variational distribution q ( t | s ) .
However, as notedin [54], the statistics of the natural images usually exhibit non-Gaussian property.
Furthermore,optimizing the variational lower bound using the fully factorized distribution encourages the student to generate the blurry outputs [55, 45, 56].
On the other hand, we employ an energy-based model asa ﬂexible distribution for the variational distribution, which is given by
[Equation]
where E θ ( t , s ) :
Y × Y → R is the energy function deﬁned by a deep neural network with parameterθ and Z θ ( s ) is the partition function.
The variational distribution given by the energy-based modelis effective for the student model to maximize the variational lower bound.
In addition, contrary to VID [48], our approach employs the variational distribution that considers the spatial dependencyamong pixels in t and s .
Overall, we optimize the energy-based model, the student network, and the teacher network, whichare achieved by the following iterative procedure:
Fix { φ s , φ t } and optimize θ .
Under ﬁxed parameters { φ s , φ t } , minimize the objective (3) with respect to θ .
Fix { θ, φ t } and optimize φ s .
Under ﬁxed parameters { θ, φ t } , maximize the objective (2) while minimizing standard knowledge distillation losses with respect to φ s .
Optimize φ t .
Minimize the standard conditional GAN loss [6, 7] with respect to φ t using a training dataset.
This step is only required for online knowledge distillation.
Otherwise, the teacher modelshould be learned ofﬂine in advance.
These steps are iterated until convergence.
Training
This section discusses our optimization procedure based on knowledge distillation, which iterativelymaximizes the mutual information between a teacher and a student using the energy-based model.
4.1 Optimization for Energy-based Model
We now discuss how to optimize the energy-based model for minimizing E s [ D KL ( p ( t | s ) ||
q θ ( t | s ))]with respect to θ given the ﬁxed model parameters { φ s , φ t } .
4.1.1 Objective
Since the conditional data distribution p ( t | s ) is constant with respect to the model parameter θ , theenergy-based model minimizes the following objective:
[Equation]
To optimize the energy-based model using the standard stochastic gradient descent method, we derivethe gradient of the objective in (5), which is given by
[Equation]
where the second expectation requires sampling from q θ ( t | s ) via a Markov Chain Monte Carlo(MCMC) method such as Gibbs sampling.
4.1.2 Improved Sampling for MCMC
The main drawback of the MCMC methods is high computational cost, so we rely on Langevindynamics, which leads to the following recursive sample generation:
[Equation]
(7) where w k is a sample drawn from N (0 , λI ) and λ is a ﬁxed step size.
Theoretically, ˜ t k becomesidentical to a sample drawn from q θ ( t | s ) as k → ∞ and λ → 0 [57].
Following Short-Run
MCMC [58], we also run K -step MCMC starting from an initial distribution, q 0 , to approximatesamples from q θ ( t | s ) , where K is the number of MCMC steps. [Figure]
4.1.3 Initial Distribution
We initialize ˜ t 0 as the output of the student network for the acceleration of the MCMC optimizationvia the Langevin sampling process.
This initialization can be a good approximation to the modeldistribution. Intuitively, since the student model is supposed to mimic the samples drawn from thedata distribution, the optimization of q θ ( t | s ) towards the data distribution p ( t | s ) from the currentstudent output would be more efﬁcient.
As a result, the proposed initialization strategy often resultsin improved performance of student models compared to existing methods [58, 59, 60, 61, 62]even with a small number of MCMC steps.
Note that Contrastive Divergence (CD) [59] uses thedata distribution for the initialization of q 0 while Persistent Contrastive Divergence (PCD) [60, 61]utilizes the samples collected in the previous epoch and [58] initializes the uniform noise distribution.
Since the data distribution is prone to induce bias in the CD estimate [62] and the random uniformdistribution is far from the data distribution, the existing methods require many steps for MCMC andincrease training time [58].
In addition, PCD updates all samples in the memory buffer, which leadsto signiﬁcantly increased memory overhead especially for generation tasks.
4.1.4 Inference
In contrast to the previous approaches for energy-based models [58, 62, 61], we do not require any
MCMC sampling for inference but generate images via a single forward process of the studentgenerator.
Therefore, the proposed algorithm incur no extra computational cost at inference time.
4.2 Optimization for Student Generator
4.2.1 Objective for Mutual Information
The student generator learns to maximize the variational lower bound ˜
I ( T ; S ) in (2) with respect toφ s , where the objective function is formally given by
[Equation]
Note that the entropy, H ( T ) , is independent of the parameters in the student network and the objectiveis equivalent to maximize the expected log likelihood as shown in the above equation.
Since thejoint distribution p ( t , s ) depends on φ s , we employ a reparameterization trick using the function G s .
Then, the gradient of the objective function in (8) with respect to φ s is derived as follows:
[Equation]
Note that we do not need to run MCMC additionally to obtain the samples from q θ ( t | G sφ s ( x )) becausethey are already available from the training procedure for the energy-based model.
Figure 1 illustratesthe overview of our framework for maximizing the mutual information.
[Table] 4.2.2 Objectives Specialized for Compression Algorithms
The proposed variational lower bound maximization framework based on the energy-based modelis incorporated into distillation methods such as OMGD [38], GCC [53], GAN-Compression [34],
CAT [37], and CAGC [32].
The brief descriptions about these methods are presented in Section 2.3,and we will discuss the details of the algorithm-speciﬁc losses in the supplementary document.
4.2.3 Full Objective
The full objective function for the student network, which is given by
[Equation]
(10) where L algo is an algorithm-speciﬁc loss, λ MI is a hyperparameter, and ˜ I
( T ; S ) is the variationallower bound deﬁned in (8).
4.3 Optimization for Teacher Generator
The teacher generator is trained using the standard methods of backbone algorithms.
Optionally, theteacher model learns to minimize the (cid:96) 1 distance between the generated image and the ground-truthwhen the ground-truth is available.
In the case of the unsupervised image-to-image translation tasks,the cycle consistency loss [7] based on the (cid:96) 1 distance enforces the instance mapping between twodomains.
To sum up, the teacher network is optimized by the adversarial loss, denoted by the GANloss, L GAN , and the reconstruction loss, L 1 , as follows:
[Equation]
(11) where λ 1 is a hyperparameter to balance the two loss terms.
Note that, for unconditional GANs, thereis nothing to reconstruct and we train the network based only on the GAN loss.
Experiments
To validate the effectiveness of the proposed compression algorithm based on variational energy-basedmodel, denoted by VEM, we conduct extensive experiments on the standard datasets using variousarchitectures.
5.1 Datasets and Models
Image-to-image translation We adopt Pix2Pix [6] on the Edge → Shoes [64] and Cityscapes [65]datasets while taking CycleGAN [7] on Horse → Zebra [7] and Summer → Winter [7].
Note that
Pix2Pix is the image-to-image translation network for paired datasets while CycleGAN is designed forunpaired image-to-image translation.
Following the original models [6, 7], we employ a U-Net-basedgenerator [66] for Pix2Pix and a ResNet-based generator [67] for CycleGAN.
For all the datasets,images are resized to 256 × 256 before feeding them into the models. [Table] [Table]
Image generation To evaluate the performance of our approach for unconditional GANs [Table] weemploy StyleGAN2 [63] on the Flickr-Faces-HQ (FFHQ) dataset [68] and Self-attention GAN(SAGAN) [4] on the CelebA dataset [69].
For image preprocessing, we resize images to 64 × 64 for CelebA and 256 × 256 for FFHQ datasets.
5.2 Evaluation Metrics
Following the protocol of the previous approaches, we report the FID scores on the Edge → Shoes, Horse → Zebra, Summer → Winter, CelebA [Table]
and FFHQ datasets while mIoU is selected for the Cityscapes dataset.
The FID score is widely utilized for assessing the quality of the generated imagesby measuring the distance between the distributions of embedding features extracted from generatedimages and real images, which are given by the Inception V3 model [70] pretrained on ImageNet.
Inthe case of the Cityscapes dataset, we adopt DRN-D-105 [71] as a segmentation model to derive themIoU scores for the generated images.
5.3 Implementation Details
The proposed algorithm is implemented in PyTorch [72] based on the publicly available code.
Weincorporate the proposed mutual information maximization framework into existing approachesincluding GAN-Compression, CAT, CAGC, and GCC to validate its beneﬁt and generality.
For faircomparisons, we do not modify any experimental settings and network structures of the comparedalgorithms.
Speciﬁcally, we set the batch size to 4 for Pix2Pix, 1 for CycleGAN, 64 for SAGAN, and16 for StyleGAN2 while we adopt the Adam optimizer with an initial learning rate of 0.0002, whichis linearly decayed to zero.
When we train our models integrated with the energy-based model, we employ the Adam optimizerwith a learning rate of 0.0001.
Note that we only run 10 steps of Langevin dynamics to draw samples [Table] [Figure] from the energy-based model for saving training time.
We set the standard deviation of random noiseto 0.005 and use a step size of 100 for each gradient step of Langevin dynamics for all datasets exceptfor the Cityscapes and FFHQ datasets.
For the two datasets, we use a step size of 50 instead.
Detailsabout the hyperparameter setting will be provided in the supplementary document.
In the case ofthe paired datasets such as Edges → Shoes and Cityscapes, we maximize the mutual informationbetween the outputs given by the student model and the real outputs instead of the teacher outputs.
5.4 Results
Pix2Pix Table 1 presents the results from Pix2Pix on the Edges → Shoes and Cityscapes datasetsby the proposed algorithm using the U-Net backbone in comparison to the original network andrecent methods including OMGD and OMGD + VID.
As shown in the table, OMGD + VEM achievesthe best performance in terms of FID and MAC among all compared algorithms on Edges → Shoes.
On Cityscapes, OMGD + VEM presents the highest mIoU by large margins even with the highestcompression rate.
Note that we report the results from OMGD based on our reproduction using theofﬁcial code, which are denoted by asterisks (*) in all tables.
CycleGAN We test our compression algorithm using CycleGAN on Horse → Zebra and Summer→ Winter.
OMGD
+ VEM outperforms all the compared methods on the both datasets, especially bylarge margins in Horse
[Figure] Unconditional GANs We compress SAGAN and StyleGAN2 using GCC [53] and CAGC [32],respectively.
Table 3 illustrates the proposed method is also effective to improve performance of theunconditional GAN compression techniques.
In particular, our approach works well with a large-scalegenerator such as StyleGAN2.
5.5 Ablation Study about Mutual Information Maximization Techniques
For maximizing mutual information with the teacher and student networks, one can use a fullyfactorized Gaussian distribution for the variational distribution followed by VID [48] or CRD [49]employing InfoNCE-like contrastive loss.
To validate the effectiveness of our energy-based modelscompared with VID and CRD, we perform an ablation study to combine with the state-of-the-art
GAN compression approaches such as GAN-Compression [34], CAT [37], and OMGD [38].
Table 4 demonstrates that VEM outperforms VID and CRD signiﬁcantly in all the tested algorithms.
This is because, compared to VID, we derive a tight lower bound of the mutual information theoretically and consider spatial dependencies among pixels thanks to the more ﬂexible variationaldistribution based on the energy function.
Also, note that the application of CRD to generative modelsis not effective since the instance-level contrastive loss is prone to be trivial due to large gaps betweenpositive and negative pairs.
5.6 Qualitative Results
Figure 2 visualizes generated images given by the original uncompressed backbones and the compressed counterparts given by OMGD and OMGD + VEM (ours) with U-Net backbones for Pix2Pixand CycleGAN on the Horse → Zebra, Summer
→ Winter, Edges → Shoes, and Cityscapes datasets.
According to the presented images, our approach is effective to enhance the generation quality.
Forexample, OMGD fails to produce the ﬁne details such as building in the background in the Cityscapesdataset while OMGD + VEM is successful.
Figure 3 depicts the generation quality by unconditional GANs on the CelebA and FFHQ datasets.
Compared with the uncompressed versions, the proposedmethod preserves or often improves the generation quality even with large compression ratios.
Also,our approach generates visually improved images compared to the previous GAN compressionmethods such as GCC and CAGC when combined with them.
Conclusion
We presented a novel GAN compression approach via knowledge distillation based on the mutualinformation, which maximizes the variational lower bound by incorporating an energy-based variational distribution.
Our energy-based model is straightforward to be incorporated into variousexisting model compression algorithms for generative models.
It provides a more ﬂexible variationaldistribution, which leads to theoretically tighter lower bounds and facilitates the consideration ofthe dependency between pixels.
The proposed algorithm demonstrates outstanding performance invarious GAN compression scenarios on multiple datasets compared to the state-of-the-art algorithms.
Checklist
For all authors..
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] See the abstract and Section 1.
(b) Did you describe the limitations of your work?
[Yes] See Section A.4 in the the supplemental material.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes] we have read the guidelines and ensured that our paper conforms to them.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[N/A] (b) Did you include complete proofs of all theoretical results?
[N/A]3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?
[Yes] We will attachthe code and instructions in the supplemental material.
(b) Did you specify all the training details (e g., data splits, hyperparameters, how theywere chosen)?
[Yes] See Section 5.
(c) Did you report error bars (e g., with respect to the random seed after running experi- ments multiple times)?
[No] We ran one time.
(d) Did you include the total amount of compute and the type of resources used (e g., type of GPUs, internal cluster, or cloud provider)?
[No] We did not include them.
If you are using existing assets (e
.g. , code data models or curating/releasing new assets.. (a) If your work uses existing assets, did you cite the creators?
[Yes] See Section 5.
(b) Did you mention the license of the assets?
[No] We did not mention the license.
(c) Did you include any new assets either in the supplemental material or as a URL?
[No] No we did not introduce any new assets.
(d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[Yes] See Section 5.
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[No] No we did not.
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]