Training GANs with Stronger Augmentations via Contrastive Discriminator

1 INTRODUCTION . Generative adversarial networks ( GANs ) ( Goodfellow et al. , 2014 ) have become one of the most prominent approaches for generative modeling with a wide range of applications ( Ho & Ermon , 2016 ; Zhu et al. , 2017 ; Karras et al. , 2019 ; Rott Shaham et al. , 2019 ) . In general , a GAN is defined by a minimax game between two neural networks : a generator network that maps a random vector into the data domain , and a discriminator network that classifies whether a given sample is real ( from the training dataset ) or fake ( from the generator ) . Provided that both generator and discriminator attain their optima at each minimax objective alternatively , it is theoretically guaranteed that the generator implicitly converges to model the data generating distribution ( Goodfellow et al. , 2014 ) . Due to the non-convex/stationary nature of the minimax game , however , training GANs in practice is often very unstable with an extreme sensitivity to many hyperparameters ( Salimans et al. , 2016 ; Lucic et al. , 2018 ; Kurach et al. , 2019 ) . Stabilizing the GAN dynamics has been extensively studied in the literature ( Arjovsky et al. , 2017 ; Gulrajani et al. , 2017 ; Miyato et al. , 2018 ; Wei et al. , 2018 ; Jolicoeur-Martineau , 2019 ; Chen et al. , 2019 ; Schonfeld et al. , 2020 ) , and the idea of incorporating data augmentation techniques has recently gained a particular attention on this line of research : more specifically , Zhang et al . ( 2020 ) have shown that consistency regularization between discriminator outputs of clean and augmented samples could greatly stabilize GAN training , and Zhao et al . ( 2020c ) further improved this idea . The question of which augmentations are good for GANs has been investigated very recently in several works ( Zhao et al. , 2020d ; Tran et al. , 2021 ; Karras et al. , 2020a ; Zhao et al. , 2020a ) , while they unanimously conclude only a limited range of augmentations ( e.g. , flipping and spatial translation ) were actually helpful for the current form of training GANs . Meanwhile , not only for GANs , data augmentation has also been played a key role in the literature of self-supervised representation learning ( Doersch et al. , 2015 ; Gidaris et al. , 2018 ; Wu et al. , 2018 ) , especially with the recent advances in contrastive learning ( Bachman et al. , 2019 ; Oord et al. , 2018 ; Chen et al. , 2020a ; b ; Grill et al. , 2020 ) : e.g. , Chen et al . ( 2020a ) have shown that the performance gap between supervised- and unsupervised learning can be significantly closed with large-scale contrastive learning over strong data augmentations . In this case , contrastive learning aims to extract the mutual information shared across augmentations , so good augmentations for contrastive learning should keep information relevant to downstream tasks ( e.g. , classification ) , while discarding nuisances for generalization . Finding such augmentations is still challenging , yet in some sense , it is more tangible than the case of GANs , as there are some known ways to formulate the goal rigourously , e.g. , InfoMax ( Linsker , 1988 ) or InfoMin principles ( Tian et al. , 2020 ) . Contribution . In this paper , we propose Contrastive Discriminator ( ContraD ) , a new way of training discriminators of GAN that incorporates the principle of contrastive learning . Specifically , instead of directly optimizing the discriminator network for the GAN loss , ContraD uses the network mainly to extract a contrastive representation from a given set of data augmentations and ( real or generated ) samples . The actual discriminator that minimizes the GAN loss is defined independently upon the contrastive representation , which turns out that a simple 2-layer network is sufficient to work as a complete GAN . By design , ContraD can be naturally trained with augmentations used in the literature of contrastive learning , e.g. , those proposed by SimCLR ( Chen et al. , 2020a ) , which are in fact much stronger than typical practices in the context of GAN training ( Zhang et al. , 2020 ; Zhao et al. , 2020c ; a ; Karras et al. , 2020a ) . Our key observation here is that , the task of contrastive learning ( to discriminate each of independent real samples ) and that of GAN discriminator ( to discriminate fake samples from the reals ) benefit each other when jointly trained with a shared representation . Self-supervised learning , including contrastive learning , have been recently applied in GAN as an auxiliary task upon the GAN loss ( Chen et al. , 2019 ; Tran et al. , 2019 ; Lee et al. , 2021 ; Zhao et al. , 2020d ) , mainly in attempt to alleviate catastopic forgetting in discriminators ( Chen et al. , 2019 ) . For conditional GANs , Kang & Park ( 2020 ) have proposed a contrastive form of loss to efficiently incorporate a given conditional information into discriminators . Our work can be differentiated to these prior works in a sense that , to the best of our knowledge , it is the first method that successfully leverage contrastive learning alone to incorporate a wide range of data augmentations in GAN training . Indeed , for example , Zhao et al . ( 2020d ) recently reported that simply regularizing auxiliary SimCLR loss ( Chen et al. , 2020a ) improves GAN training , but could not outperform existing methods based on simple data augmentations , e.g. , bCR ( Zhao et al. , 2020c ) . 2 BACKGROUND . Generative adversarial networks . We consider a problem of learning a generative model pg from a given dataset { xi } Ni=1 , where xi ∼ pdata and xi ∈ X . To this end , generative adversarial network ( GAN ) ( Goodfellow et al. , 2014 ) considers two neural networks : ( a ) a generator network G : Z → X that maps a latent variable z ∼ p ( z ) into X , where p ( z ) is a specific prior distribution , and ( b ) a discriminator network D : X → [ 0 , 1 ] that discriminates samples from pdata and those from the implicit distribution pg derived fromG ( z ) . The primitive form of trainingG andD is the following : min G max D V ( G , D ) : = Ex∼pdata [ log ( D ( x ) ) ] + Ez∼p ( z ) [ log ( 1−D ( G ( z ) ) ) ] . ( 1 ) For a fixed G , the inner maximization objective ( 1 ) with respect to D leads to the following optimal discriminator D∗G , and consequently the outer minimization objective with respect to G becomes to minimize the Jensen-Shannon divergence between pdata and pg : D∗G : = maxD V ( G , D ) = pdata pdata+pg . Although this formulation ( 1 ) theoretically guarantees p∗g = pdata as the global optimum , the nonsaturating loss ( Goodfellow et al. , 2014 ) is more favored in practice for better optimization stability : max D L ( D ) : = V ( G , D ) , and min G L ( G ) : = −Ez [ log ( D ( G ( z ) ) ) ] . ( 2 ) Here , compared to ( 1 ) , G is now optimized to let D to classify G ( z ) as 1 , i.e. , the “ real ” . Contrastive representation learning . Consider two random variables v ( 1 ) and v ( 2 ) , which are often referred as views . Generally speaking , contrastive learning aims to extract a useful representation of v ( 1 ) and v ( 2 ) from learning a function that identifies whether a given sample is from p ( v ( 1 ) ) p ( v ( 2 ) |v ( 1 ) ) or p ( v ( 1 ) ) p ( v ( 2 ) ) , i.e. , whether two views are dependent or not . More specifically , the function estimates the mutual information I ( v ( 1 ) ; v ( 2 ) ) between the two views . To this end , Oord et al . ( 2018 ) proposed to minimize InfoNCE loss , which turns out to maximize a lower bound of I ( v ( 1 ) ; v ( 2 ) ) . Formally , for a given v ( 1 ) i ∼ p ( v ( 1 ) ) and v ( 2 ) i ∼ p ( v ( 2 ) |v ( 1 ) i ) while assuming v ( 2 ) j ∼ p ( v ( 2 ) ) for j = 1 , · · · , K , the InfoNCE loss is defined by : LNCE ( v ( 1 ) i ; v ( 2 ) , s ) : = − log exp ( s ( v ( 1 ) i , v ( 2 ) i ) ) ∑K j=1 exp ( s ( v ( 1 ) i , v ( 2 ) j ) ) , ( 3 ) where s ( · , · ) is the score function that models the log-density ratio of p ( v ( 2 ) |v ( 1 ) ) to p ( v ( 2 ) ) , possibly including some parametrized encoders for v ( 1 ) and v ( 2 ) . Many of recent unsupervised representation learning methods are based on this general framework of contrastive learning ( Wu et al. , 2018 ; Bachman et al. , 2019 ; Hénaff et al. , 2020 ; He et al. , 2020 ; Chen et al. , 2020a ) . In this paper , we focus on the one called SimCLR ( Chen et al. , 2020a ) , that adopts a wide range of independent data augmentations to define views : specifically , for a given set of data samples x = [ xi ] Ni=1 , SimCLR applys two independent augmentations , namely t1 and t2 , to the given data to obtain v ( 1 ) and v ( 2 ) , i.e. , ( v ( 1 ) , v ( 2 ) ) : = ( t1 ( x ) , t2 ( x ) ) . The actual loss of SimCLR is slightly different to InfoNCE , mainly due to practical considerations for sample efficiency : LSimCLR ( v ( 1 ) , v ( 2 ) ) : = 1 2N N∑ i=1 ( LNCE ( v ( 1 ) i ; [ v ( 2 ) ; v ( 1 ) −i ] , sSimCLR ) + LNCE ( v ( 2 ) i ; [ v ( 1 ) ; v ( 2 ) −i ] , sSimCLR ) ) , ( 4 ) where v−i : = v \ { vi } . For sSimCLR , SimCLR specifies to use ( a ) an encoder network f : X → Rde , ( b ) a small neural network called projection head h : Rde → Rdp , and ( c ) the normalized temperature-scaled cross entropy ( NT-Xent ) . Putting altogether , sSimCLR is defined by : sSimCLR ( v ( 1 ) , v ( 2 ) ; f , h ) : = h ( f ( v ( 1 ) ) ) · h ( f ( v ( 2 ) ) ) τ · ||h ( f ( v ( 1 ) ) ) ||2||h ( f ( v ( 2 ) ) ) ||2 , ( 5 ) where τ is a temperature hyperparameter . Once the training is done with respect to LSimCLR , the projection head h is discarded and f is served as the learned representation for downstream tasks .

Summary: This paper improves upon state-of-the-art GANs by incorporating recent advances of contrastive representation learning into the training of discriminator. In particular, the discriminator loss function consists of three terms: (1) the original SimCLR loss on the multi-view real data pairs; (2) the supervised contrastive loss (Khosla et al, 2020) that assigns high scores among the fake sample pairs and giving lower score among the real data pairs; (3) the usual discriminator loss in GAN training. While each of these terms alone is not entirely new, the author proposes several tricks to make the training of GANs together with the contrastive loss works. Empirically, the proposed method outperforms other GAN methods trained with auxiliary data augmentation techniques, and demonstrates good representations under the linear classifier probing setup.