Whitening and Coloring Batch Transform for GANs

1 INTRODUCTION . Generative Adversarial Networks ( GANs ) ( Goodfellow et al . ( 2014 ) ) have drawn much attention in the last years due to their proven ability to generate sufficiently realistic short videos ( Vondrick et al . ( 2016 ) ) or still images ( Gulrajani et al . ( 2017 ) ) , possibly conditioned on some input information such as : a class label ( Odena et al . ( 2017 ) ) , an image ( Isola et al . ( 2017 ) ; Tang et al . ( 2019 ) ) , a textual description ( Reed et al . ( 2016 ) ) , some structured data ( Tang et al . ( 2018 ) ; Siarohin et al . ( 2018a ) ) or some perceptual image attribute ( Siarohin et al . ( 2018b ) ) ) . When the generation process depends on some input data , the framework is commonly called conditional GAN ( cGAN ) . In this paper we deal with both conditional and unconditional GANs and we propose to replace Batch Normalization ( Ioffe & Szegedy ( 2015 ) ) ( BN ) with a Whitening and Coloring ( WC ) transform ( Hossain ( 2016 ) ) in order to jointly speed-up training and increase the network representation capacity . BN , proposed in ( Ioffe & Szegedy ( 2015 ) ) for discriminative tasks and then adopted in many discriminative and generative networks , is based on feature standardization of all the network layers . The original motivation behind the introduction of BN is to alleviate the internal Covariate Shift problem caused by the continuous modification of each layer ’ s input representation during training ( Ioffe & Szegedy ( 2015 ) ) . However , recent works ( Santurkar et al . ( 2018 ) ; Kohler et al . ( 2018 ) ) showed that a major motivation of the empirical success of BN relies also on a higher training stability , the latter being due to a better-conditioning of the Jacobian of the corresponding loss when BN is used compared to training without BN . Basically , feature normalization makes the loss landscape smoother , hence it stabilizes and speeds-up training . Based on this intuition , BN has very recently been extended in ( Huang et al . ( 2018 ) ) to full-feature decorrelation using ZCA whitening ( Kessy et al . ( 2017 ) ) in a discriminative scenario . On the other hand , GAN training is known to be particularly unstable , because GAN optimization aims at finding a Nash equilibrium between two players , a problem which is more difficult and less stable than the common discriminative-network optimization and that frequently leads to nonconvergence . Specifically , Odena et al . ( 2018 ) show the relation between stability of GAN training and conditioning of the Jacobian . This motivates our proposal to extend BN to full-feature whitening in a GAN scenario ( more details on this in Sec . 6 ) . Differently from Huang et al . ( 2018 ) , our feature whitening is based on the Cholesky decomposition ( Dereniowski & Marek ( 2004 ) ) , and we empirically show that our method is much faster and stable and achieves better results in GAN training with respect to ZCA whitening . Importantly , while in ( Huang et al . ( 2018 ) ) feature whitening is followed by a per-dimension scaling and shifting of each feature , as in the original BN ( Ioffe & Szegedy ( 2015 ) ) , we propose to use a coloring transform ( Hossain ( 2016 ) ) to keep the network representation capacity unchanged with respect to a non-whitened layer . Our coloring projects each whitened feature vector in a new distribution using learnable multivariate filters and we show that this transformation is critical when using whitening for GANs . The second contribution of this paper is based on the aforementioned coloring transform to represent class-specific information in a cGAN scenario . Formally , given a set of class labels Y = { y1 , ... , yn } , we want to generate an image I = G ( z , y ) , where the generator G is input with a noise vector z and a class label y . For instance , if y = cat , we want a foreground object in I depicting a cat , if y = dog , then I should represent a dog , etc . The categorical conditional information y is commonly represented in G using conditional batch normalization ( cBN ) ( Dumoulin et al . ( 2016b ) ; Gulrajani et al . ( 2017 ) ; Miyato & Koyama ( 2018 ) ) , where a class-specific pair of scaling-shifting parameters is learned for each class y . In our proposed conditional Whitening and Coloring ( cWC ) the class-specific scaling-shifting parameters are replaced with a combination of class-agnostic and class-specific coloring filters . The intuitive idea behind the proposed cWC is that multivariate transformation filters are more informative than the scalar parameters proposed in ( Dumoulin et al . ( 2016b ) ) and , thus , they can more accurately represent class-specific information . Finally , in order to have a set of class-specific coloring filters ( { Γy } ) which grows sub-linearly with respect to the number of classes n , we represent each Γy using a common , restricted dictionary of class-independent filters together with class-filter soft-assignment weights . The use of a classindependent dictionary of filters , together with a class-agnostic branch in the architecture of our cGAN generators ( see Sec . 4 ) makes it possible to share parameters over all the classes and to reduce the final number of weights which need to be learned . In summary , our contributions are the following : • We propose a whitening-coloring transform for batch normalization to improve stability of GAN training which is based on the Cholesky decomposition and a re-projection of the whitened features in a learned distribution . • In a cGAN framework we propose : ( 1 ) A class-specific coloring transform which generalizes cBN ( Dumoulin et al . ( 2016b ) ) and ( 2 ) A soft-assignment of classes to filters to reduce the network complexity . • Using different datasets and basic GAN-cGAN frameworks , we show that our WC and cWC consistently improve the original frameworks ’ results and training speed . Our code is publicly available https : //github.com/AliaksandrSiarohin/wc-gan . 2 RELATED WORK . Batch Normalization . BN ( Ioffe & Szegedy ( 2015 ) ) is based on a batch-based ( per-dimension ) standardization of all the network ’ s layers . Formally , given a batch of d-dimensional samples B = { x1 , ... , xm } , input to a given network layer , BN transforms each sample xi ∈ Rd in B using : BN ( xi , k ) = γk xi , k − µB , k√ σ2B , k + + βk , ( 1 ) where : k ( 1 ≤ k ≤ d ) indicates the k-th dimension of the data , µB , k and σB , k are , respectively , the mean and the standard deviation computed with respect to the k-th dimension of the samples in B and is a constant used to prevent numerical instability . Finally , γk and βk are learnable scaling and shifting parameters , which are introduced in ( Ioffe & Szegedy ( 2015 ) ) in order to prevent losing the original representation capacity of the network . More in detail , γk and βk are used to guarantee that the BN transform can be inverted , if necessary , in order to represent the identity transform ( Ioffe & Szegedy ( 2015 ) ) . BN is commonly used in both discriminative and generative networks and has been extended in many directions . For instance , Ba et al . ( 2016 ) propose Layer Normalization , where normalization is performed using all the layer ’ s activation values of a single sample . Luo ( 2017 ) proposes to whiten the activation values using a learnable whitening matrix . Differently from ( Luo ( 2017 ) ) , our whitening matrix only depends on the specific batch of data . Similarly to our work , batch-dependent full-feature decorrelation is performed also by Huang et al . ( 2018 ) , who use a ZCA-whitening and show promising results in a discriminative scenario . However , the whitening part of our method is based on the Cholesky decomposition and we empirically show that , in a GAN scenario where training is known to be highly unstable , our method produces better results and it is much faster to compute . Moreover , while in ( Huang et al . ( 2018 ) ) only scaling and shifting parameters are used after feature decorrelation , we propose to use a coloring transform . Analogously to the role played by γk and βk in Eq . 1 , coloring can potentially invert the whitening transform if that were the optimal thing to do , in this way preserving the original representation capacity of the network . We empirically show ( see Sec . 5.1.1 ) that whitening without coloring performs poorly . Last but not least , our conditional coloring can be used to represent rich class-specific information in a cGAN framework ( see below ) . Representing conditional information in cGANs . cGANs have been widely used to produce images conditioned on some input . For instance , Isola et al . ( 2017 ) transform a given image in a second image represented in another “ channel ” . Siarohin et al . ( 2018a ) extend this framework to deal with structured conditional information ( the human body pose ) . In ( Reed et al . ( 2016 ) ) a textual description of an image is used to condition the image generation . When the conditioning data y is a categorical , unstructured variable , as in the case investigated in this paper , information about y can be represented in different ways . cBN ( Dumoulin et al . ( 2016b ) ) uses y to select class-specific scaling and shifting parameters in the generator ( more details in Sec . 4 ) . Another common approach is to represent y using a one-hot vector , which is concatenated with the input layer of both the discriminator and the generator ( Mirza & Osindero ( 2014 ) ; Perarnau et al . ( 2016 ) ) or with the first convolutional layer of the discriminator ( Perarnau et al . ( 2016 ) ) . Miyato & Koyama ( 2018 ) split the last discriminator layer in two branches . The first branch estimates the class-agnostic probability that the input image is real , while the second branch represents class-conditional information . This solution is orthogonal to our proposed cWC and we show in Sec . 5.2 that the two approaches can be combined for boosting the results obtained in ( Miyato & Koyama ( 2018 ) ) . 3 THE WHITENING AND COLORING TRANSFORM . We start describing the details of our Whitening and Coloring transform in the unconditional case , which we refer to as WC in the rest of the paper . We always assume an image domain and convolutional generator/discriminator networks . However , most of the solutions we propose can potentially be applied to non-visual domains and non-convolutional networks . Let F ∈ Rh×w×d be the tensor representing the activation values of the convolutional feature maps for a given image and layer , with d channels and h × w spatial locations . Similarly to BN , we consider the d-dimensional activations of each location in the h×w convolutional grid as a separate instance xi ∈ Rd inB . This is done also to have a larger cardinality batch which alleviates instability issues when computing the batch-related statistics . In fact , a mini-batch of m′ images corresponds to m = |B| = m′ × h× w ( Ioffe & Szegedy ( 2015 ) ) . Using a vectorial notation , our proposed WC-based batch normalization is given by WC ( xi ) = Coloring ( x̂i ) , where x̂i = Whitening ( xi ) : Coloring ( x̂i ) = Γx̂i + β ( 2 ) Whitening ( xi ) = WB ( xi − µB ) . ( 3 ) In Eq . 3 , the vector µB is the mean of the elements in B ( being µB , k in Eq . 1 its k-th component ) , while the matrix WB is such that : W > BWB = Σ −1 B , where ΣB is the covariance matrix computed using B . This transformation performs the full whitening of xi and the resulting set of vectors B̂ = { x̂1 , ... , x̂m } lies in a spherical distribution ( i.e. , with a covariance matrix equal to the identity matrix ) . Note that , similarly to µB , k and σB , k in Eq . 1 , µB and WB are completely data-dependent , being computed using only B and without any learnable parameter involved . On the other hand , Eq . 2 performs the coloring transform ( Hossain ( 2016 ) ) , which projects the elements in B̂ onto a multivariate Gaussian distribution with an arbitrary covariance matrix ( see Sec . 2 for the motivation behind this projection ) . Eq . 2 is based on the learnable parameters β = ( β1 , ... , βk , ... , βd ) > and Γ = ( ... ck ... ) , where ck is a vector of parameters ( a coloring filter ) which generalizes the scalar parameter γk in Eq . 1 . Coloring is a linear operation and can be simply implemented using 1×1×d convolutional filters ck and the corresponding biases βk . We compute WB using the Cholesky decomposition , which has the important properties to be efficient and well conditioned . We provide the details in Sec . A , where we also analyze the computational costs , while in Sec . E we compare our approach with the ZCA-based whitening proposed by Huang et al . ( 2018 ) . WC is plugged before each convolutional layer of our generators ( G ) . In most of our experiments , we use a ResNet-like architecture ( He et al . ( 2016 ) ) forG with 3 blocks ( Gulrajani et al . ( 2017 ) ) and a final convolutional layer ( i.e . 3× 2 + 1 = 7 convolutional layers , corresponding to 7 WC layers ) . Fig . 1 ( a ) shows a simplified block scheme . We do not use WC nor we introduce any novelty in the discriminators ( more details in Sec . 5 ) .

Summary: This paper proposed Whitening and Coloring (WC) transform to replace batch normalization (BN) in generators for GAN. WC generalize BN by normalizing features with decorrelating (whitening) matrix, and then denormalizing (coloring) features by learnable weights. The main advantage of WC is that it exploits the full correlation matrix of features, while BN only considers the diagonal. WC is differentiable and is only 1.32x slower than BN. The authors also apply conditional WC, which learn the parameters of coloring conditioned on labels, to conditional image generation.  Experimental results show WC achieves better inception score and FI distance comparing to BN on CIFAR-10, CIFAR-100, STL-10 and Tiny Imagenet. Furthermore, the conditional image generation results by WC are better than all previous methods.