Revisiting Discriminator in GAN Compression: A Generator-discriminator Cooperative Compression Scheme

1 Introduction . Generative Adversarial Networks ( GANs ) have been widely popularized in diverse image synthesis tasks , such as image generation [ 16 , 50 ] , image translation [ 64 , 30 ] and super resolution [ 35 , 59 ] . However , the ultra-high computational and memory overhead of GANs hinder its deployment on resource-constrained edge devices . To alleviate this issue , a series of traditional model compression algorithms , i.e. , network pruning [ 18 , 36 ] , weight quantization [ 39 , 26 ] , low-rank decomposition [ 10 , 19 ] , knowledge distillation [ 22 , 52 ] , and efficient architecture design [ 29 , 55 ] have been applied to reduce calculational consumptions of GANs . Previous work [ 56 , 12 , 61 ] attempted to directly employ network pruning methods to compress the generator but obtain unimpressive results , as shown in Figure 1 ( a ) . The similar phenomenon also occurred in SAGAN shown in Appendix A . A potential reason is that these methods failed to take into account the generator and the discriminator must follow the Nash equilibrium state to avoid the mode ∗This work was done when Shaojie Li was a intern at Bytedance Inc. †Corresponding Author : rrji @ xmu.edu.cn 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . collapse in the adversarial learning . In other words , these methods simply compress the generator while maintaining the original capability of the discriminator , which resulted in breaking the balance in adversarial learning . As shown in Figure 1 ( b ) and 1 ( c ) , when compressing the generator , the loss of the discriminator gradually tends to zero , such a situation indicates that the capacity of the discriminator significantly surpasses that of the lightweight generator . Furthermore , the capacity imbalance between the generator and discriminator leads to the mode collapse problem [ 47 ] . In order to retain the Nash equilibrium between the lightweight generator and discriminator , we must revisit the role of the discriminator in the procedure of GAN compression . A straightforward method is to reduce the capacity of the discriminator to re-maintain the Nash equilibrium between it and the lightweight generator . Therefore , we train a lightweight generator against the discriminators with different channel numbers . The results are shown in Figure 2 , we can obtain the following observations : i ) The channel numbers in the discriminator significantly influence the performance of the lightweight generator ; ii ) The discriminator with the same channel compression ratio as the generator may not get the best result . iii ) The optimal value of channel numbers for discriminators is task-specific . In short , it is arduous to choose an appropriate channel number for the discriminator in different tasks . To solve the above issues , we design a generator-discriminator cooperative compression scheme , term GCC in this paper . GCC consists of four sub-networks , i.e. , the original uncompressed generator and discriminator , a pruned generator , and a selective activation discriminator . GCC selectively activates partial convolutional channels of the selective activation discriminator , which guarantees the whole optimization process to maintains in the Nash equilibrium stage . In addition , the compression of the generator will damage its generating ability , and even affect the diversity of generated images , resulting in the occurrence of mode collapse . GCC employs an online collaborative distillation that combines the intermediate features of both the original uncompressed generator and discriminator and then distills them to the pruned generator . The introduction of online collaborative distillation boosts the lightweight generator performance for effective compression . The contributions of GCC can be summarized as follows : • To the best of our knowledge , we offer the first attempt to revisit the role of discriminator in GAN compression . We propose a novel selective activation discriminator to automatically select and activate the convolutional channels according to a local capacity constraint and a global coordination constraint . The dynamic equilibrium relationship between the original models provides a guide to choose the activated channels in the discriminator . • A novel online collaborative distillation is designed to simultaneously employ intermediate features of teacher generator and discriminator to guide the optimization process of the lightweight generator step by step . These auxiliary intermediate features provide more complementary information into the lightweight generator for generating high-fidelity images . • Extensive experiments on various GAN-based generation tasks ( i.e. , image generation , image-to-image translation , and super-resolution ) demonstrate the effectiveness and generalization of GCC . GCC contributes to reducing 80 % computational costs while maintains comparable performance in image translation . 2 Related Work . Generative Adversarial Networks . Generative Adversarial Network , i.e. , GANs [ 16 ] has attracted intense attention and has been extensively studied for a long time . For example , DCGAN [ 50 ] greatly enhanced the capability of GANs by introducing the convolutional layer . The recent works [ 17 , 33 , 63 , 4 ] employed a series of advanced network architecture to further improve the fidelity of the generated images . In addition , several novel loss functions [ 54 , 17 , 2 , 45 ] were designed to stabilize the adversarial training stage of GANs . A popular research direction of GANs is image-to-image translation . The image-to-image translation is designed to transfer images from the source domain to the target domain while retaining the content representation of the image . Image-to-image translation is widely used in computer vision and image processing , such as semantic image synthesis [ 30 , 48 ] , style transfer [ 64 , 1 ] , super resolution [ 35 , 59 ] , image painting [ 49 , 65 ] and so on . Mode Collapse . Mode collapse is the main catastrophic problem during the adversarial training stage of GANs , which leads to the diversity of generated images are monotonous , only fitting part of the real data distribution . To address this issue , recent efforts were focus on new objective functions [ 47 , 17 , 34 ] , architecture modifications [ 5 , 15 ] and mini-batch discrimination [ 54 , 33 ] . Actually , the capacity gap between the generator and the discriminator is a more essential reason for mode collapse . In this paper , we provide a different perspective to eliminate the capacity gap , we select the convolutional channel to be activated in the discriminator according to a local capacity constraint and a global coordination constraint . GAN Compression . To reduce huge resource consumption during the inference process of GANs , great efforts have emerged in the field of GAN compression in the recent two years . The searchbased methods [ 56 , 37 , 12 , 23 , 40 , 58 ] exploited neural architecture search and genetic algorithm to obtain a lightweight generator . However , they fell into the trap of manual pre-defined search space and huge search costs . The prune-based methods [ 38 , 31 , 57 , 61 ] directly pruned a lightweight generator architecture from the original generator architecture . However , these works failed to take discriminator pruning into account , which would seriously destroy the Nash equilibrium between the generator and the discriminator . The discriminator-free methods [ 51 , 12 ] directly use the large GAN model as a teacher to distill the lightweight generator without discriminator , which also achieved good performance . Slimmable GAN [ 23 ] correspondingly shrank the network width of the discriminator with the generator . However , as shown in Fig . 2 , there is no linear between the channel numbers of the optimal discriminator and generator . Anycost GAN [ 40 ] proposed a generator-conditioned discriminator , which generates the discriminator architecture via passing the generator architecture into the fully connected layer . However , the fully connected layer is only trained through the original GAN loss , so it is difficult to accurately predict the optimal discriminator architecture for the lightweight generator . In this work , we design a generator-discriminator cooperative compression ( GCC ) scheme , which employ a selective activation discriminator and novel online collaborative distillation to boost the performance of the lightweight generator . 3 Methodology . This section , we illustrate the proposed GCC framework in Figure 3 . We first briefly introduce the composition and loss function of GAN in Section 3.1 . Then , we present how to obtain the lightweight architecture from the original generator in Section 3.2 . The selective activation discriminator that automatically chooses activated convolutional channels is illustrated in Section 3.3 . In Section 3.4 , we propose a novel online collaborative distillation scheme to further improve the performance of the lightweight generator . 3.1 Preliminaries . Generative Adversarial Networks ( GANs ) contain two fundamental components , generator G and discriminator D. Among them , G maps the input z into a fake image for cheating the discriminator , while the discriminator distinguishes the generator outputs from the real images x . The generator and discriminator are alternately optimized via the adversarial term to achieve their respective optimization objectives . The optimization objectives are defined as follows : LG = Ez∼p ( z ) [ fG ( −D ( G ( z ) ) ) ] ( 1 ) LD = Ex∼pdata [ fD ( −D ( x ) ) ] ︸ ︷︷ ︸ LDreal +Ez∼p ( z ) [ fD ( D ( G ( z ) ) ) ] ︸ ︷︷ ︸ LDfake ( 2 ) where fG and fD are the loss functions of the generator and discriminator , respectively . Theoretically , the generator and discriminator maintain the Nash equilibrium when the generator and discriminator have a balanced capacity , and the discriminator is cheated . In this paper , we observe from experiments that the Nash equilibrium maintained in the GAN compression stage to avoid model collapse and obtain an impressive performance . 3.2 Generator Compression . In order to avoid additional computational overhead brought by the neural architecture search , we directly employ the conventional network pruning methods to compress the generator . Previous work [ 43 ] revealed that the core of network pruning searches for a suitable sub-architecture from the original network architecture instead of parameter inheritance . Therefore , we leverage the one-shot pruning method [ 36 , 42 ] to obtain the appropriate lightweight generator architecture from the original generator architecture . In our experiments , we adopt the slimming [ 42 ] to prune the generator when the network contains the Batch Normalization ( BN ) layers . We use L1-norm pruning [ 36 ] on the contrary ( without BN layers ) . The detailed pruning process is described in Appendix B . Finally , we only keep the architecture of the pruned generator and train it from scratch with a selective activation discriminator proposed in Section 3.3 . 3.3 Selective Activation Discriminator . As shown in Figure 2 , the convolutional channel number in the discriminator have an essential impact on the performance of the lightweight generator . Therefore , we design a selective activation discriminator to automatically choose and activate convolutional channels via learnable retention factors . We assign a learnable retention factor to each convolution kernel of the original discriminator , each retention factor represents the probability that the corresponding convolution kernel participates in the inference process . When the retention probability is larger than the given threshold , the corresponding convolution kernel is retained , otherwise it is suppressed . We denote retention factors set as A = { αi } LDi=1 where αi = { αi1 , αi2 , ... , αini } ∈ Rni is the retention factors of the i-th convolution layer and ni is the number of convolution kernel in the i-th layer . Given the threshold τ , the retention factor αij determines whether the j-th convolution kernel of the i-th layer activate or not : Iij = { 1 if αij ∈ [ τ , 1 ] 0 if αij ∈ [ 0 , τ ) . ( 3 ) Then activated convolution kernel ’ s output O′ij is calculated by : O′ij = Iij ∗Oij , ( 4 ) where Oij is the original output feature map . Since the binary function is not differentiable for αij , we employ STE [ 27 ] to estimate the gradient of αij as : ∂Larch ∂αij = ∂Larch ∂Iij , ( 5 ) where Larch is a loss function , and the retention factor αij is manually clipped to the range of [ 0 , 1 ] . The learning of retention factor is considered as the architecture search of the discriminator . The selective activation discriminator is optimized according to a local capacity constraint and a global coordination constraint . On the one hand , the incremental loss gap between the generator and the discriminator will break the Nash equilibrium in the adversarial training scheme , which leads to the discriminator effortlessly distinguishing the authenticity of the input images . The pruned generator and the selective activation discriminator should pursuit capacity balance to maintain the Nash equilibrium . Therefore , a local capacity constraint Llocal is designed as follows : Llocal = |LSG − LSDfake | , ( 6 ) where the superscript S denotes the student ( compressed model ) . On the other hand , the loss gap between the original model ’ s combination internally reflects the relative capacity relationship between them . Therefore , we design a global coordination constraint to adjust the relative ability of the student generator and discriminator to be consistent with that of the teacher combination . This constraint helps to pull the loss gap between teacher and student models into close proximity . In this way , the uncompressed model is regarded as teacher model to guide the learning of retention factors to automatically adjust the discriminator capacity to match the lightweight generator . The global coordination constraint Lglobal is calculated by : Lglobal = |Llocal − |LTG − LTDfake || ( 7 ) The superscript T represents the teacher ( uncompressed ) model . Due to the instability of adversarial training , we use Exponential Moving Average ( EMA ) to stabilize the loss corresponds to the teacher models . In short , the optimization objective of the retention factor is defined as follows : Larch = LSD + Lglobal . ( 8 ) where LSD guarantees the basic identification function of the discriminator . Furthermore , we use the following bilevel optimization [ 7 , 11 ] to optimize α : i ) Update discriminator weights to minimize LD for t times . ii ) Update retention factor α to minimize Larch . iii ) Repeat step i ) and step ii ) until the end of training . In the optimization process of the discriminator weights , we freeze the retention factors and vice versa .

Summary: This work revisits the role of discriminator in GAN compression and finds that the capacity mismatch when compressing the generator with a fixed discriminator will lead to model collapse. To this end, this work proposes a generator-discriminator co-operative compression framework to automatically adapt the size of discriminator to maintain the Nash equilibrium with the lightweight generator. The extensive evaluation results validate the superiority of the proposed framework over SOTA GAN compression methods.