Data-Efficient GAN Training Beyond (Just) Augmentations: A Lottery Ticket Perspective

Training generative adversarial networks ( GANs ) with limited real image data generally results in deteriorated performance and collapsed models . To conquer this challenge , we are inspired by the latest observation , that one can discover independently trainable and highly sparse subnetworks ( a.k.a. , lottery tickets ) from GANs . Treating this as an inductive prior , we suggest a brand-new angle towards data-efficient GAN training : by first identifying the lottery ticket from the original GAN using the small training set of real images ; and then focusing on training that sparse subnetwork by re-using the same set . We find our coordinated framework to offer orthogonal gains to existing real image data augmentation methods , and we additionally present a new feature-level augmentation that can be applied together with them . Comprehensive experiments endorse the effectiveness of our proposed framework , across various GAN architectures ( SNGAN , BigGAN , and StyleGANV2 ) and diverse datasets ( CIFAR-10 , CIFAR-100 , Tiny-ImageNet , ImageNet , and multiple few-shot generation datasets ) . Codes are available at : https : //github . com/VITA-Group/Ultra-Data-Efficient-GAN-Training . 1 Introduction The quantity , diversity , and high quality of natural images available in the general domain have played an essential role in the achieved breakthroughs of Generative Adversarial Networks ( GANs ) [ 2–11 ] over the past few years . However , it could become challenging or even infeasible for specific application domains to collect a sufficiently large-scale dataset , due to various constraints on the imaging expense , subject type , image quality , privacy , copyright status , and more . That prohibits GANs ’ broader applications in these domains , e.g. , for generating synthetic training data [ 12 ] . Examples of such domains include medical images , images from scientific experiments , images of rare species , or photos of a specific person or landmark . Eliminating the need of immense datasets for GAN training is highly demanded for those scenarios . Naively training GAN with scarce samples leads to overconfident discriminators that overfit the small training data [ 13–15 , 1 ] ; it usually ends up with training divergence and drastic performance degradation ( evidenced later in Figure 2 ) . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . This paper addresses the above issue from a brand new perspective by decomposing the challenging GAN training in limited data regimes into two sequential sub-problems : ( i ) finding independent trainable subnetworks ( i.e. , lottery tickets in GANs ) [ 16 , 17 ] ; then ( ii ) training the located subnetworks , which we show is more data-efficient by itself , and can further benefit from aggressive augmentations ( both the input data and feature levels ) . Either sub-problem becomes much less data-hungry to train , and the two sub-problems re-use the same small training set of real images . Although this paper focuses on tackling the data-efficient training of GANs , such a coordinated framework might potentially be generalized to training other deep models with higher data efficiency too . Our key enabling technique is to leverage the lottery ticket hypothesis ( LTH ) [ 18 ] . LTH shows the feasibility to locate highly sparse subnetworks ( called “ winning tickets ” ) that are capable of training in isolation to match or even outperform the performance of original unpruned models . Recently , [ 17 , 16 ] revealed the existence of winning ticket in GANs ( called “ GAN tickets ” ) . However , none of the existing works discuss the influence of training data size on locating and training those tickets . Our work takes one step further , and shows that one can identify the same high-quality GAN tickets even in the data-scarce regime . The found GAN tickets also serve as a sparse structural prior to solve the second sub-problem with less data , while maintaining an unimpaired trainability blessed by the LTH assumption [ 18 ] . Figure 1 ( the outer circle ’ s blue dots ) evidences that we can identify sparse GAN tickets that achieve superior performance than full GANs in the data-scarce scenarios . The new lottery ticket angle complements the existing augmentation techniques [ 19–21 ] , and we further show that they can be organically combined to boost performance further1 . When we train the identified lottery ticket , we demonstrate its training can benefit as well from the latest data-level augmentation strategies , ADA [ 15 ] and DiffAug [ 1 ] . Furthermore , we introduce a novel feature-level augmentation that can be applied in parallel to data-level . It injects adversarial perturbations into GANs ’ intermediate features to implicitly regularize both discriminator and generator . Combining the new feature-level and existing data-level augmentations in training GAN tickets leads to more stabilized training dynamics , and establishes new state-of-the-arts for data-efficient GAN training . Extensive experiments are conducted on a variety of the latest GAN architectures and datasets , which consistently validate the effectiveness of our proposal . For example , our BigGAN tickets at 36.00 % and 67.24 % sparsity levels reach an ( FID , IS ) of ( 23.14 , 52.98 ) and ( 70.91 , 7.03 ) , on Tiny-ImageNet 64× 64 and ImageNet 128× 128 , with 10 % and 25 % training data , respectively . On CIFAR-10 and CIFAR-100 , for SNGAN and BigGAN tickets at 67.24 % ∼ 86.58 % sparsity , our results with only 10 % training data can even surpass their dense counterparts . Impressively , our method can generate high-quality images on par with other GAN transfer learning approaches , by training on as few as 100 real samples and without using any pre-training . 2 Related Work . GANs and Data-Efficient GAN Training . GANs [ 22 ] have gained popularity in diverse computer vision scenarios . To stabilize GAN training and improve the visual fidelity and diversity of generated images , extensive studies have been conducted , such as sophisticated network architectures [ 23 , 8 , 24– 26 ] , improved training recipes [ 27 , 4 , 28 , 29 ] , and more stable objectives [ 30–35 ] . [ 36 , 37 ] utilize semi- and self-supervised learning to pursue label efficiency in GAN training . Recently , how to train GANs without sufficient real images in the target domain sparkles new interests . There have been efforts on adapting a pre-trained GAN generator , including BSA [ 38 ] , AdaFM [ 39 ] , Elastic Weight Consolidation [ 40 ] , and Few-Shot GAN [ 41–43 ] . However , those methods assume a large , related source domain as pre-training , based on which they further alleviate target domain data limitation by only tuning small subsets of weights . They are hence in a completely different track from our “ stand-alone ” data-efficient training goal where no pre-training is leveraged . [ 44 , 45 ] select core-sets of training data to speed up GAN training . A few recent attempts [ 1 , 15 ] leverage differentiable or adaptive data augmentations to significantly improve GAN training in limited data regimes . Lately , [ 46 ] investigates a regularization approach , on constraining the distance between the current prediction of the real image and a moving average variable that tracks the historical predictions of the generated image , that complements the data augmentation methods . Lottery Ticket Hypothesis and GAN Tickets . [ 18 ] claims the existence of independently trainable sparse subnetworks that can match or even surpass the performance of dense networks . [ 47 , 48 ] scale 1We also tried to add augmentations in the lottery ticket finding stage , but did not observe visible impact . up LTH by rewinding [ 49 , 50 ] . Follow-up researches evidence LTH across broad fields , including visual recognition [ 18 , 47 , 51–60 ] , natural language processing [ 48 , 61 , 50 , 62–64 ] , graph neural network [ 65 ] , and reinforcement learning [ 61 ] . Recently , LTH has been extended to GANs by [ 16 , 17 ] , who validated the existence of winning tickets in the min-max game beyond minimization . Compared with the aforementioned work , our work is the first to study LTH in the data-scarce regime ( for GANs , and in general ) . Besides finding highly compact yet same capable subnetworks , our work reveals LTH ’ s power in saving training data - an appealing perspective never being examined before . Adversarial Training and Augmentations . Deep neural networks suffer from severe performance degradation [ 66 , 67 ] when facing adversarial inputs [ 67–69 ] . To address this notorious vulnerability , various defense mechanisms [ 70–79 ] have been proposed . Among others , adversarial training-based approaches achieve superior adversarial robustness [ 67–69 ] , although at the price of sacrificing benign generalization [ 80 , 70–75 ] . Several recent works investigate enhancing model ( benign ) generalization ability with adversarial training [ 81–86 ] . They adopt adversarially perturbed input images , embeddings , or intermediate features , into model training to ameliorate performance on the clean test sets . Specifically , the damaging effects of adversarial training could be controlled by extra batch normalization [ 81 ] or so . Different from those minimization problems that previous work has focused on , the two-player GAN optimization is more challenging . Generally , the adversarial competition between two players poses impediments to exploit extra adversarial information during training GANs . 3 Methodology . 3.1 Revisiting GANs and the Overfitting Challenge . Generative adversarial networks ( GANs ) are dedicated to modeling the target distribution with the two-player game formulation of a generator G and a discriminator D. Specifically , the generator G takes a random sampled latent vector z ( e.g. , from a Gaussian distribution ) as input and outputs the fake sample G ( z ) . The discriminator D aims to distinguish generated fake samples G ( z ) from real samples x . Alternative optimizations for the discriminator ’ s loss LD and the generator ’ s loss LG are adopted in the standard GAN training , which can be depicted as follows : LD : = Ex∼pdata ( x ) [ fD ( −D ( x ) ) ] + Ez∼p ( z ) [ fD ( D ( G ( z ) ) ) ] LG : = Ez∼p ( z ) [ fG ( −D ( G ( z ) ) ) ] , where loss functions fD ( x ) , fG ( x ) have multiple choices , e.g. , the non-saturating loss [ 3 ] with fD ( x ) = fG ( x ) = log ( 1 + e x ) , and the hinge loss [ 23 ] with fD ( x ) = max ( 0 , 1 +x ) and fG ( x ) = x. pdata ( x ) and p ( z ) represent the data distribution of real samples and latent vectors . LD is maximized to update D ’ s parameters φ ( i.e. , D ( · ) : = D ( · , φ ) ) , and LG is minimized to update G ’ s parameters θ ( i.e. , G ( · ) = G ( · , θ ) ) . Algorithm 1 Data-Efficient Iterative Magnitude Pruning Procedures 1 : Input : Initial two masks mg = 1‖θ‖0 and md = 1 ‖φ‖0 ; Initialization weights θ0 and φ0 2 : Output : { G ( · , θ0 mg ) , D ( · , φ0 md ) } 3 : repeat 4 : Training { G ( · , θ0 mg ) , D ( · , φ0 md ) } for t epochs with limited training data 5 : Pruning ρ = 20 % of remaining weights in both G and D 6 : Updating the binary masks mg and md accordingly 7 : Rewinding weights of G , D to θ0 and φ0 8 : until masks reach the desired sparsity level Training Failures of GANs under Limited Data . [ 15 , 1 ] observe that GANs ’ performance has severely deteriorated when only limited training data is available . The discriminator tends to memorize and heavily overfit the small training samples , leaving a gap between the real sample ’ s and the generated sample ’ s distribution . As shown in Figure 2 , with only 10 % CIFAR-10 data available for training SNGAN [ 23 ] , the training and validation accuracies of the discriminator D quickly saturate to nearly 100 % ( ideally close to 50 % ) , which indicates D to become over-confident in distinguishing real and generated samples . It demonstrates that D simply memorizes the training data , and such overfitting leads to training collapses and deteriorated quality of generated images . To address this dilemma , we suggest a new data-efficient GAN training workflow , decomposed into two stages : ( i ) finding winning tickets in GANs via Algorithm 1 ; then ( ii ) training the found GAN tickets , potentially with both data- and feature-level augmentations , via Algorithm 2 . Blessed by LTH , the located GAN ticket shows improved generalization ability , and is further enhanced by augmentations that prevent D from becoming too confident .

Summary: This paper proposed to extend Lottery Ticket Hypothesis (LTH) for data-efficient GAN training. Specifically, the authors have found that sparse subnetworks lead to better sample qualities of Image GANs. In addition, they have proposed AdvAug that combines the ideas of feature augmentation and adversarial training.