Large Scale GAN Training for High Fidelity Natural Image Synthesis

1 INTRODUCTION . The state of generative image modeling has advanced dramatically in recent years , with Generative Adversarial Networks ( GANs , Goodfellow et al . ( 2014 ) ) at the forefront of efforts to generate highfidelity , diverse images with models learned directly from data . GAN training is dynamic , and sensitive to nearly every aspect of its setup ( from optimization parameters to model architecture ) , but a torrent of research has yielded empirical and theoretical insights enabling stable training in a variety of settings . Despite this progress , the current state of the art in conditional ImageNet modeling ( Zhang et al. , 2018 ) achieves an Inception Score ( Salimans et al. , 2016 ) of 52.5 , compared to 233 for real data . In this work , we set out to close the gap in fidelity and variety between images generated by GANs and real-world images from the ImageNet dataset . We make the following three contributions towards this goal : • We demonstrate that GANs benefit dramatically from scaling , and train models with two to four times as many parameters and eight times the batch size compared to prior art . We introduce two simple , general architectural changes that improve scalability , and modify a regularization scheme to improve conditioning , demonstrably boosting performance . ∗Work done at DeepMind †Equal contribution • As a side effect of our modifications , our models become amenable to the “ truncation trick , ” a simple sampling technique that allows explicit , fine-grained control of the tradeoff between sample variety and fidelity . • We discover instabilities specific to large scale GANs , and characterize them empirically . Leveraging insights from this analysis , we demonstrate that a combination of novel and existing techniques can reduce these instabilities , but complete training stability can only be achieved at a dramatic cost to performance . Our modifications substantially improve class-conditional GANs . When trained on ImageNet at 128×128 resolution , our models ( BigGANs ) improve the state-of-the-art Inception Score ( IS ) and Fréchet Inception Distance ( FID ) from 52.52 and 18.65 to 166.5 and 7.4 respectively . We also successfully train BigGANs on ImageNet at 256×256 and 512×512 resolution , and achieve IS and FID of 232.5 and 8.1 at 256×256 and IS and FID of 241.5 and 11.5 at 512×512 . Finally , we train our models on an even larger dataset – JFT-300M – and demonstrate that our design choices transfer well from ImageNet . Code and weights for our pretrained generators are publicly available 1 . 2 BACKGROUND . A Generative Adversarial Network ( GAN ) involves Generator ( G ) and Discriminator ( D ) networks whose purpose , respectively , is to map random noise to samples and discriminate real and generated samples . Formally , the GAN objective , in its original form ( Goodfellow et al. , 2014 ) involves finding a Nash equilibrium to the following two player min-max problem : min G max D Ex∼qdata ( x ) [ logD ( x ) ] + Ez∼p ( z ) [ log ( 1−D ( G ( z ) ) ) ] , ( 1 ) where z ∈ Rdz is a latent variable drawn from distribution p ( z ) such as N ( 0 , I ) or U [ −1 , 1 ] . When applied to images , G and D are usually convolutional neural networks ( Radford et al. , 2016 ) . Without auxiliary stabilization techniques , this training procedure is notoriously brittle , requiring finely-tuned hyperparameters and architectural choices to work at all . Much recent research has accordingly focused on modifications to the vanilla GAN procedure to impart stability , drawing on a growing body of empirical and theoretical insights ( Nowozin et al. , 2016 ; Sønderby et al. , 2017 ; Fedus et al. , 2018 ) . One line of work is focused on changing the objective function ( Arjovsky et al. , 2017 ; Mao et al. , 2016 ; Lim & Ye , 2017 ; Bellemare et al. , 2017 ; Salimans et al. , 2018 ) to encourage convergence . Another line is focused on constraining D through gradient penalties ( Gulrajani et al. , 2017 ; Kodali et al. , 2017 ; Mescheder et al. , 2018 ) or normalization ( Miyato et al. , 2018 ) , both to counteract the use of unbounded loss functions and ensure D provides gradients everywhere to G. Of particular relevance to our work is Spectral Normalization ( Miyato et al. , 2018 ) , which enforces Lipschitz continuity on D by normalizing its parameters with running estimates of their first singular values , inducing backwards dynamics that adaptively regularize the top singular direction . Relatedly Odena et al . ( 2018 ) analyze the condition number of the Jacobian of G and find that performance is dependent on G ’ s conditioning . Zhang et al . ( 2018 ) find that employing Spectral Normalization in G improves stability , allowing for fewer D steps per iteration . We extend on these analyses to gain further insight into the pathology of GAN training . Other works focus on the choice of architecture , such as SA-GAN ( Zhang et al. , 2018 ) which adds the self-attention block from ( Wang et al. , 2018 ) to improve the ability of both G and D to model global structure . ProGAN ( Karras et al. , 2018 ) trains high-resolution GANs in the single-class setting by training a single model across a sequence of increasing resolutions . In conditional GANs ( Mirza & Osindero , 2014 ) class information can be fed into the model in various ways . In ( Odena et al. , 2017 ) it is provided to G by concatenating a 1-hot class vector to the noise vector , and the objective is modified to encourage conditional samples to maximize the corresponding class probability predicted by an auxiliary classifier . de Vries et al . ( 2017 ) and 1 https : //tfhub.dev/s ? q=biggan Dumoulin et al . ( 2017 ) modify the way class conditioning is passed to G by supplying it with classconditional gains and biases in BatchNorm ( Ioffe & Szegedy , 2015 ) layers . In Miyato & Koyama ( 2018 ) , D is conditioned by using the cosine similarity between its features and a set of learned class embeddings as additional evidence for distinguishing real and generated samples , effectively encouraging generation of samples whose features match a learned class prototype . Objectively evaluating implicit generative models is difficult ( Theis et al. , 2015 ) . A variety of works have proposed heuristics for measuring the sample quality of models without tractable likelihoods ( Salimans et al. , 2016 ; Heusel et al. , 2017 ; Bińkowski et al. , 2018 ; Wu et al. , 2017 ) . Of these , the Inception Score ( IS , Salimans et al . ( 2016 ) ) and Fréchet Inception Distance ( FID , Heusel et al . ( 2017 ) ) have become popular despite their notable flaws ( Barratt & Sharma , 2018 ) . We employ them as approximate measures of sample quality , and to enable comparison against previous work . 3 SCALING UP GANS . In this section , we explore methods for scaling up GAN training to reap the performance benefits of larger models and larger batches . As a baseline , we employ the SA-GAN architecture of Zhang et al . ( 2018 ) , which uses the hinge loss ( Lim & Ye , 2017 ; Tran et al. , 2017 ) GAN objective . We provide class information to G with class-conditional BatchNorm ( Dumoulin et al. , 2017 ; de Vries et al. , 2017 ) and to D with projection ( Miyato & Koyama , 2018 ) . The optimization settings follow Zhang et al . ( 2018 ) ( notably employing Spectral Norm in G ) with the modification that we halve the learning rates and take two D steps per G step . For evaluation , we employ moving averages of G ’ s weights following Karras et al . ( 2018 ) ; Mescheder et al . ( 2018 ) ; Yazc et al . ( 2018 ) , with a decay of 0.9999 . We use Orthogonal Initialization ( Saxe et al. , 2014 ) , whereas previous works used N ( 0 , 0.02I ) ( Radford et al. , 2016 ) or Xavier initialization ( Glorot & Bengio , 2010 ) . Each model is trained on 128 to 512 cores of a Google TPUv3 Pod ( Google , 2018 ) , and computes BatchNorm statistics in G across all devices , rather than per-device as is typical . We find progressive growing ( Karras et al. , 2018 ) unnecessary even for our 512×512 models . Additional details are in Appendix C. We begin by increasing the batch size for the baseline model , and immediately find tremendous benefits in doing so . Rows 1-4 of Table 1 show that simply increasing the batch size by a factor of 8 improves the state-of-the-art IS by 46 % . We conjecture that this is a result of each batch covering more modes , providing better gradients for both networks . One notable side effect of this scaling is that our models reach better final performance in fewer iterations , but become unstable and undergo complete training collapse . We discuss the causes and ramifications of this in Section 4 . For these experiments , we report scores from checkpoints saved just before collapse . We then increase the width ( number of channels ) in each layer by 50 % , approximately doubling the number of parameters in both models . This leads to a further IS improvement of 21 % , which we posit is due to the increased capacity of the model relative to the complexity of the dataset . Doubling the depth did not initially lead to improvement – we addressed this later in the BigGAN-deep model , which uses a different residual block structure . We note that class embeddings c used for the conditional BatchNorm layers in G contain a large number of weights . Instead of having a separate layer for each embedding ( Miyato et al. , 2018 ; Zhang et al. , 2018 ) , we opt to use a shared embedding , which is linearly projected to each layer ’ s gains and biases ( Perez et al. , 2018 ) . This reduces computation and memory costs , and improves training speed ( in number of iterations required to reach a given performance ) by 37 % . Next , we add direct skip connections ( skip-z ) from the noise vector z to multiple layers of G rather than just the initial layer . The intuition behind this design is to allow G to use the latent space to directly influence features at different resolutions and levels of hierarchy . In BigGAN , this is accomplished by splitting z into one chunk per resolution , and concatenating each chunk to the conditional vector c which gets projected to the BatchNorm gains and biases . In BigGAN-deep , we use an even simpler design , concatenating the entire z with the conditional vector without splitting it into chunks . Previous works ( Goodfellow et al. , 2014 ; Denton et al. , 2015 ) have considered variants of this concept ; our implementation is a minor modification of this design . Skip-z provides a modest performance improvement of around 4 % , and improves training speed by a further 18 % .

Summary: The authors present a empirical investigation of methods for scaling GANs to complex datasets, such as ImageNet, for class-conditioned image generation. They first build and describe a strong baseline based on recently proposed techniques for GANs and push the performance on large datasets with several modifications presented sequentially, to obtain strong state-of-the-art IS/FID scores, as well as impressive visual results. The authors propose a simple truncation trick to control the fidelity/variance which is interesting on its own but cannot always scale with the architecture. The authors further propose a orthogonalization-based regularization to mitigate this problem. An investigation of training collapse at large scale is also performed; the authors investigate some regularization schemes based on gathered empirical evidence. As a result, they explore and discard Spectral Normalization of the generator as a way to prevent collapse and show that a severe tradeoff between stability and quality can be controlled when using zero-centered gradient penalties in the Discriminator. In the end, no solution that can ensure quality and stability is found, except having prohibitively large amounts of data (~300M images). Models are evaluated on the ImageNet and on this internal, bigger dataset.