Information-Theoretic GAN Compression with Variational Energy-based Model
Abstract
We propose an information-theoretic knowledge distillation approach for the compression of generative adversarial networks, which aims to maximize the mutual information between teacher and student networks via a variational optimization based on an energy-based model.
Because the direct computation of the mutual information in continuous domains is intractable, our approach alternatively optimizes the student network by maximizing the variational lower bound of the mutual information.
To achieve a tight lower bound, we introduce an energy-based model relying on a deep neural network to represent a ﬂexible variational distribution that deals with high-dimensional images and consider spatial dependencies between pixels, effectively.
Since the proposed method is a generic optimization algorithm, it can be conveniently incorporated into arbitrary generative adversarial networks and even dense prediction networks, e . g ., image enhancement models.
We demonstrate that the proposed algorithm achieves outstanding performance in model compression of generative adversarial networks consistently when combined with several existing models.
Introduction
Generative adversarial networks (GANs) [1] have accomplished impressive achievements on various computer vision tasks including image synthesis [2–5], image-to-image translation [6–9], video prediction [10–12], and many others.
However, despite the outstanding performance, their applicability to resource-hungry systems, e . g ., edge or mobile devices, is limited due to high computational costs.
To deal with this issue, GAN compression, a task reducing the size of generator networks, has recently drawn the attention of many machine learning researchers.
The straightforward way to compress generators is to directly employ traditional network compression techniques such as low-rank approximation [13–15], network quantization [16–19], pruning [20– 25], and knowledge distillation [26–30].
However, the compression of generative models typically involves a higher-dimensional output space and requires matching the output distributions between the original and compressed models instead of the output instances given by the two models.
Hence, as reported in [31–33], the synthesized images obtained from the compressed generators via the channel pruning [34, 23, 24] are often suboptimal—producing images with artifacts and suffering from image quality degradation, which implies that a naïve application of the methods designed for discriminative models with low dimensional output spaces is not effective for GAN compressions.
Existing approaches for compressing GANs [31, 35, 32, 36–38, 33, 39] often utilize knowledge distillation and manage to achieve competitive accuracy of student models.
Speciﬁcally, these methods transfer the knowledge of a teacher network to a student model via intermediate representations, ﬁnal outputs, and embedded features [40, 41].
To this end, the existing algorithms employ (cid:96) 1 or (cid:96)norms to compute the difference between the representations of the two networks, but they 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 student.
For the compression of generative models, the use of the mutual information between two random variables makes more sense than the direct pixel-wise comparison between two outputs, which is the standard way to transfer knowledge from a teacher to a student in discriminative models.
However, because the computation of the mutual information is generally intractable in the continuous domain, we employ a variational approach [42], where the lower bound of the mutual information is alternatively optimized using a variational distribution, an approximation for the conditional distribution of the teacher output given the student output.
We further adopt an energy-based model for the ﬂexible representation of the variational distribution and tighten the lower bound for effective knowledge transfer.
Note that the proposed energy-based model is a convenient way to deal with high-dimensional data e . g ., images, and an effective tool to handle spatial dependency between pixels within an image.
Combined with the existing approaches, our framework achieves substantial performance improvement on GAN compression tasks.
The main contributions of our work are summarized below:
• We propose a novel energy-based framework of GAN compression for knowledge distillation, which maximizes the mutual information between a teacher and a student.
• The proposed energy-based approach is effective for representing a ﬂexible variational distribution and can be incorporated into many GAN compression methods thanks to its generality.
• Experimental results verify that the proposed method improves performance consistently 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 are described in Section 3 and 4, and the experimental results are presented in Section 5.
Finally, we conclude this paper in Section 6.
Related Work
2.1 Generative Adversarial Networks (GANs) 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 by synthesizing realistic images from random noise samples while the discriminator attempts to distinguish the real data from the fake ones.
Unlike the vanilla GANs, conditional GANs [43, 6, 7] control the image generation conditioned on additional inputs such as a class label [43] or an image to be translated into another image [6, 7] ( e .
g ., Horse → Zebra).
The conditional GANs are not conﬁned to the conditional image generation tasks, but applied to various areas such as image superresolution [44], text-to-image synthesis [45], and image inpainting [46].
Since the computational costs of the conditional GANs are typically high in terms of FLOPs, the proposed GAN compression algorithm is mainly tested on those networks but we also verify the effectiveness of our approach on the unconditional case.
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 [47], output distributions [26], intermediate activations [48], or attention maps [27].
Speciﬁcally, Hinton et al .
[26] learn the student representations by minimizing the relative entropy from its normalized softmax outputs to the teacher’s while [47, 48] minimize the mean squared error between the representations of the teacher and student networks.
Both Variational Information Distillation (VID) [49] and Contrastive Representation Distillation (CRD)
[50] employ mutual information between teacher and student to estimate the similarity of the representations from the two networks.
The former relies on a variational approach [42] for the tractable computation of mutual information between the two networks.
On the other hand, the latter derives an InfoNCE-like loss function from the lower bound of the mutual information to learn the representations of a student.
However, VID maximizes the variational lower bound using a strong assumption about the variational distribution, a fully factorized Gaussian distribution, and CRD is difﬁcult to be employed in generative models due to its contrastive learning objective, which is prone to be trivial in image generation tasks without semantic attribute annotations, e .
g ., class labels.
Our algorithm is related to VID and CRD since they all attempt to maximize the mutual information between the teacher and student networks, but has clear advantage over them in terms of the ﬂexibility of the variational distribution and the applicability to generative models.
2.3 GAN Compression
GAN compression algorithms typically aim to reduce the size of the generator without considering the discriminator since it is not used for inference.
Co-Evolution [31] utilizes an evolutionary algorithm [51] to reduce the channel size in the generator of CycleGAN [7] based on a ﬁtness formulated based on FLOPs and the mean squared reconstruction error between the original uncompressed and pruned generators.
GAN-Compression [35] employs a once-for-all network training scheme [52] combined with the reconstruction loss as well as the intermediate feature distillation loss [48], and selects the best performing subnetwork given a target budget.
Compression and Teaching (CAT)
[38] designs an efﬁcient building block for generator, and prunes the network using a magnitude-based channel pruning method [24] while maximizing the similarity between the two networks based on centered kernel alignment [53].
Content-Aware GAN Compression (CAGC)
[33] uses the feature distillation and reconstruction loss based only on the object parts of a generated image while additionally considering LPIPS [41] as a perceptual distance metric for synthesized images.
On the contrary, Online Multi-Granularity Distillation (OMGD) [39] iteratively trains teacher networks with original conditional GAN objectives while optimizing a student network with the total variation loss [54] for spatial smoothness and the distillation loss for making the student mimic the ﬁnal outputs and the intermediate feature maps of the teachers.
Generator-discriminator Cooperative
Compression (GCC) [55] also adopts the online distillation scheme while selecting channels in the original discriminator to match the reduced capacity of the student generator at each iteration.
To validate the effectiveness of our framework, we present how the proposed algorithm can improve performance via the combinations with existing methods such as GCC, OMGD, CAGC, CAT, and
GAN-Compression.
Variational GAN Compression with Energy-based Model
Our goal is to compress generators in GANs through knowledge distillation.
Let S and T be random variables 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 to the student, we maximize the 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 indicate realizations of the teacher and student output distributions, respectively.
A practical challenge in this objective is that the computation of the mutual information for multivariate continuous random variables is typically intractable.
To tackle this issue, we adopt a variational method [42], where we derive the lower bound of the mutual information, ˜
I
( T ; S ) , using a variational distribution, q ( t | s ) , which is given by
[Equation]
To maximize the mutual information, we increase the lower bound by minimizing the KL-divergence between the true and variational distributions, which is deﬁned as
[Equation]
Variational Information Distillation (VID) [49] optimizes the variational lower bound by employing the fully factorized Gaussian distribution as the variational distribution q ( t | s ) .
However, as noted in [56], the statistics of the natural images usually exhibit non-Gaussian properties.
Furthermore, optimizing the lower bound using the fully factorized distribution encourages the student to generate blurry outputs [57, 46, 58].
To the contrary, we employ an energy-based model as a ﬂ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 parametrized with θ and Z θ ( s ) is the partition function.
The variational distribution given by the energy-based model is effective for the student model to maximize the variational lower bound.
In addition, differentiated from VID [49], our approach employs the variational distribution that considers the spatial dependency among pixels in t and s .
Overall, we optimize the energy-based model, the student network, and the teacher network, which is achieved by the following iterative procedure:
Fix { φ s , φ t } and optimize θ .
Under the ﬁxed parameters { φ s , φ t } , minimize the objective in (3) with respect to θ . 2.
Fix { θ, φ t } and optimize φ s .
Under the ﬁxed parameters { θ, φ t } maximize the objective in (2) while minimizing the standard knowledge distillation losses with respect to φ s . 3.
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 model should be learned ofﬂine in advance.
These three steps are repeated until convergence.
Training Procedures
This section describes the optimization procedure of our approach for knowledge distillation, which iteratively maximizes the mutual information between a teacher and a student using the energy-based model.
4.1 Optimization of Energy-based Model
The ﬁrst step of the iterative procedure is the optimization of 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 θ , the energy-based model minimizes the following objective:
[Equation]
To optimize the energy-based model using the standard stochastic gradient descent method, we derive the gradient of the objective in (5), which is given by
[Equation]
where the second expectation of the right-hand side requires sampling from q θ ( t | s ) via a Markov
Chain Monte Carlo (MCMC) method, for example, Gibbs sampling.
4.1.2 Improved MCMC sampling
The main drawback of the MCMC methods is the high computational cost, so we rely on the Langevin dynamics, which leads to the following recursive sample generation:
[Equation]
[Figure]
where w k is a sample drawn from N (0 , λI ) and λ is a ﬁxed step size.
Theoretically, ˜ t k becomes identical to a sample drawn from q θ ( t | s ) as k → ∞ and λ → 0 [59].
Following Short-Run
MCMC [60], we also run K -step MCMC starting from an initial distribution, q 0 , to approximate samples from q θ ( t | s ) , where K is the number of MCMC steps.
4.1.3 Initial distribution
We initialize ˜ t 0 as the output of the student network for the acceleration of the MCMC optimization via the Langevin sampling process.
This initialization is a good approximation to the model distribution; intuitively, since the student model is supposed to mimic the samples drawn from the data distribution, the optimization of q θ ( t | s ) towards the data distribution p ( t | s ) from the current student output would be more efﬁcient.
As a result, the proposed initialization strategy often results in improved performance of student models compared to the existing methods [61, 60, 62–65] even with a small number of MCMC steps.
On the contrary, we can naturally adopt the data distributions of teacher networks, which are static in Contrastive Divergence (CD) [62] and dynamic in Persistent Contrastive
Divergence (PCD) [63, 64], for setting q 0 while Short-Run MCMC [60] can initialize the variational distribution using the uniform noise distribution.
Since the data distribution is prone to induce bias in the CD estimate [65] and the random uniform distribution is far from the data distribution, the existing methods require many steps of MCMC and consequently increase training time [60].
In addition, PCD updates all samples in the memory buffer, which leads to signiﬁcantly increased memory overhead, especially in generation tasks.
4.1.4 Inference
In contrast to the previous approaches for energy-based models [60, 65, 64], we do not require any
MCMC sampling for inference but generate images via a single forward process of the student generator.
Therefore, the proposed algorithm incurs no extra computational cost at inference time.
4.2 Optimization of Student Generator 4.2.1 Loss 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 objective is equivalent to maximize the expected log likelihood as shown in the above equation.
Since the joint distribution p ( t , s ) depends on φ s , we employ a reparameterization trick using a function G s .
Then, the gradient of the objective function in (8) with respect to φ s is derived as follows:
[Equation]
[Table]
Note [Table] we do not need to perform MCMC additionally to obtain the samples from q θ ( t | G sφ s ( x )) because they are already available from the training procedure for the energy-based model.
Figureillustrates the overview of our framework for maximizing the mutual information.
4.2.2 Algorithm-speciﬁc loss
The proposed variational lower bound maximization framework using the energy-based model is incorporated into model compression algorithms based on knowledge distillation such as OMGD [39],
GCC [55], GAN-Compression [35], CAT [38], and CAGC [33].
The brief descriptions about these methods are presented in Section 2.3, and we discuss the details of the algorithm-speciﬁc losses in the supplementary document.
4.2.3 Combined loss
The full objective function for the student network is given by
[Equation]
[Equation]
where L algo is the algorithm-speciﬁc loss and ˜
I ( T ; S ) is the variational lower bound deﬁned in (8) while λ MI is a hyperparameter balancing the two terms.
4.3 Optimization of Teacher Generator
The teacher generator is trained by the standard method of the corresponding GAN algorithm, e .
g .,
Pix2Pix [6] and CycleGAN [7].
Optionally, the teacher model learns to minimize the (cid:96) 1 distance between the generated image and the ground-truth when 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)distance may enforce the instance mapping between two domains.
To sum up, the teacher network is optimized by the adversarial loss, denoted by L GAN , and the reconstruction loss, L rec , as follows:
[Equation]
(11) where λ rec is a hyperparameter to balance the two loss terms.
Note that, for unconditional GANs, there is 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 the variational energybased model, referred to as VEM, we conduct extensive experiments on the standard datasets using various architectures.
Since VEM is a model-agnostic algorithm component, we incorporate the module into existing GAN compression approaches and the combined models are expressed in the form of “[Algorithm Name] + VEM” throughout this section.
[Table] [Table]
5.1 Datasets and Models
Image-to-image translation We adopt Pix2Pix [6] on the Edges → Shoes [67] and Cityscapes [68] 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 data while CycleGAN is mainly proposed for unpaired image-to-image translation.
Following the original models [6, 7], we employ generators based on U-Net [69] and ResNet [70] for Pix2Pix and CycleGAN, respectively.
For all the datasets, images are resized to 256 × 256 before feeding them into the models.
Image generation To evaluate the performance of our approach for unconditional GANs
[Table]
we run experiments with StyleGAN2 [66] on Flickr-Faces-HQ (FFHQ) [71] and Self-Attention GAN (SAGAN) [4] on CelebA [72].
As a preprocessing, we resize input images to 64 × 64 for CelebA and 256 × 256 for FFHQ.
5.2 Evaluation Metrics
Following the protocol of the previous approaches, we report the Fréchet Inception Distance (FID) [73] scores on Edges → Shoes, Horse → Zebra, Summer →
Winter, CelebA, and FFHQ while the mean Intersection over Union (mIoU) is selected for Cityscapes.
The FID score is widely used for assessing the quality of the generated images by measuring the distance between the distributions of embedding features extracted from generated and real images, which are given by the
Inception V3 model [74] pretrained on ImageNet.
In the case of the Cityscapes dataset, we adopt
DRN-D-105 [75] as a model to perform semantic segmentation on the translated images and derive mIoU scores for the evaluation of image-to-image translation networks.
5.3 Implementation Details
The proposed algorithm is implemented in PyTorch [76] based on the publicly available code 1 .
Additionally, we incorporate the proposed mutual information maximization framework into existing approaches by using the ofﬁcial codes of GAN-Compression 2 [35], CAT 3 [38], CAGC 4 [33], and
GCC 5 [55] to validate its beneﬁt and generality.
For fair comparisons, we do modify neither experimental settings nor network structures of the compared algorithms.
Speciﬁcally, we set the batch size to 4 for Pix2Pix, 1 for CycleGAN, 64 for SAGAN, and 16 for StyleGAN2 while we adopt the Adam optimizer with an initial learning rate of 0.0002, which decays to zero linearly.
When we train our energy-based model, we employ the Adam optimizer with a learning rate of 0.0001.
Note that we only run 10 steps of Langevin dynamics to draw samples from the energy-based model to save training time.
We set the standard deviation of random noise to 0.005 and use a step size of 100 for each gradient step of Langevin dynamics for all datasets except for Cityscapes and
FFHQ.
For these two datasets, we use a step size of 50 instead.
Details about the hyperparameter settings are provided in the supplementary document.
In the case of paired datasets such as Edges →
Shoes and Cityscapes, we maximize the mutual information between 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 the Pix2Pix models on Edges → Shoes and Cityscapes, where the proposed algorithm, denoted by OMGD + VEM, is compared with existing approaches including the vanila OMGD [39].
As shown in the table, OMGD + VEM achieves the best performance in terms of FID and MAC among all compared algorithms on Edges → Shoes.
On Cityscapes,
OMGD
+ VEM presents the best mIoU by large margins even with the highest compression rate.
Note that we report the results from OMGD based on our reproductions using the ofﬁcial code, which are denoted by asterisks (*) in all tables.
CycleGAN We test our compression algorithm for CycleGAN on Horse → Zebra and Summer →
Winter.
As shown in Table 2, OMGD
+ VEM outperforms all the compared methods in both datasets, especially by large margins in Horse → Zebra.
Unconditional GANs We compress SAGAN and StyleGAN2 using GCC [55] and CAGC [33], respectively.
Table 3 illustrates that the proposed method is also effective to improve performance of the unconditional GAN compression techniques.
In particular, our approach also works well with a large-scale generator such as StyleGAN2.
5.5 Ablation Study about Mutual Information Maximization Techniques
To maximize the mutual information between the teacher and student networks, one can use a fullyfactorized Gaussian distribution for the variational distribution as in VID [49] or employ a contrastive loss similar with infoNCE [77] following the strategy in CRD [50].
To validate the effectiveness of our energy-based models, VEM, in comparison with VID and CRD, we perform an ablation study, where the three add-ons are tested through their integrations into the baselines such as OMGD [39],
GAN-Compression [35], CAT [38], GCC [55], and CAGC [33].
Table 4 demonstrates that VEM consistently outperforms VID and CRD in all the tested GAN compression algorithms.
This is because, compared to VID, we derive a tighter lower bound of the mutual information theoretically and consider spatial dependencies among pixels thanks to the more ﬂexible variational distribution based on the energy function.
On the other hand, the application of CRD to generative models may not be beneﬁcial since instance-level contrastive losses are less effective due to large gaps between positive and negative pairs.
[Table]
[Table]
[Table]
5.6 Effect of λ MI for Mutual Information
We analyze the effect of the balancing parameter λ MI using the Pix2Pix algorithm, which is tested on the Cityscapes dataset, when VEM is combined with OMGD [39].
As shown in Table 5, VEM is not sensitive to the hyperparameter and improves the performance of OMGD
[Table]
5.7 Qualitative Results
Figure 2 visualizes generated images given by the original uncompressed backbones of Pix2Pix and
CycleGAN, where we also present the images synthesized by the compressed counterparts using
OMGD and OMGD + VEM.
Pix2Pix is tested on the Edges → Shoes and Cityscapes datasets, and
CycleGAN is evaluated on the Horse → Zebra and Summer → Winter datasets.
The results imply that our approach is effective to enhance generation quality.
For example, OMGD often produces unusual artifacts in the background with the Cityscapes dataset, which are signiﬁcantly reduced in
OMGD + VEM.
Figure 3 depicts generation results given by unconditional GANs on CelebA and
FFHQ.
Compared with the uncompressed versions, the proposed method preserves or often improves image quality even with high compression rates.
Also, our approach synthesizes images with better visual quality when combined with the previous GAN compression methods such as GCC and CAGC. [Figure] [Figure]
[Figure]
Conclusion
We presented a novel GAN compression method via knowledge distillation based on mutual information, which maximizes the lower bound of the mutual information by incorporating an energy-based variational distribution.
Our energy-based model is straightforward to be incorporated into various existing model compression algorithms for generative models.
It provides a more ﬂexible variational distribution, which leads to theoretically tighter lower bounds and facilitates the consideration of the dependency between pixels.
The proposed algorithm demonstrates outstanding performance in various GAN compression scenarios on multiple datasets compared to the state-of-the-art approaches.
Acknowledgments This work was partly supported by SAIT, Samsung Electronics Co., Ltd, the
Bio & Medical Technology Development Program of the National Research Foundation (NRF) funded by the Korea government (MSIT)
[No. 2021M3A9E4080782], and Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT)
[No. 2021-0-01343, No. 2022-0-00959].
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 ?? .
(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 to them?
[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 experimental results (either in the supplemental material or as a URL)?
[Yes] We attach the code and instructions in the supplementary material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] See Section 5.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[No] We did not report them.
(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 amount spent on participant compensation?
[N/A]