Conditional Invertible Neural Networks for Guided Image Generation

1 INTRODUCTION . Generative adversarial networks ( GANs ) produce ever larger and more realistic samples ( Karras et al. , 2017 ; Brock et al. , 2019 ) . Hence they have become the primary choice for a majority of image generation tasks . As such , their conditional variants ( cGANs ) would appear to be the natural tool for conditional image generation as well , and they have successfully been applied in many scenarios ( Ledig et al. , 2017 ; Miyato & Koyama , 2018 ) . However , a lack in diversity is especially common when the condition itself is an image , and special precautions have to be taken to avoid mode collapse and training stability continues to pose a challenge . Conditional variational autoencoders ( cVAEs ) do not suffer from the same problems . Training is generally stable , and since every data point is assigned a region in latent space , sampling yields the full variety of data seen during training . However cVAEs come with drawbacks of their own : The assumption of a Gaussian posterior on the decoder side implies an L2 reconstruction loss , which is known to cause blurriness . In addition , the partition of the latent space into diagonal Gaussians leads to either mode-mixing issues or regions of poor sample quality ( Kingma et al. , 2016 ) . There has also been some success in combining aspects of both approaches for certain tasks , such as ( Isola et al. , 2017 ; Zhu et al. , 2017b ; Park et al. , 2019 ) . We propose a third approach , by extending Invertible Neural Networks ( INNs , Dinh et al . 1Code is available , but held back for the review process due to anonymity . ( 2016 ) ; Kingma & Dhariwal ( 2018 ) ; Ardizzone et al . ( 2019 ) ) for the task of conditional image generation , by adding conditioning inputs to their core building blocks . INNs are neural networks which are by construction bijective , efficiently invertible , and have a tractable Jacobian determinant . They represent transport maps between the input distribution p ( x ) and a prescribed , easy-to-samplefrom latent distribution p ( z ) . During training , the likelihood of training samples from p ( x ) is maximized in latent space , while at inference time , z-samples can trivially be transformed back to the data domain . Previously , INNs have been used successfully for unconditional image generation , e.g . by Dinh et al . ( 2016 ) and Kingma & Dhariwal ( 2018 ) . Unconditional INN training is related to that of VAEs , but it compensates for some key disadvantages : Firstly , since reconstructions are perfect by design , no reconstruction loss is needed , and generated images do not become blurry . Secondly , each x maps to exactly one z in latent space , and there is no need for posteriors p ( z |x ) . This avoids the VAE problem of disjoint or overlapping regions in latent space . In terms of training stability and sample diversity , INNs show the same strengths as autoencoder architectures , but with superior image quality . We find that these positive aspects apply to conditional INNs ( cINNs ) as well . One limitation of INNs is that their design restricts the use of some standard components of neural networks , such as pooling and batch normalization layers . Our conditional architecture alleviates this problem , as the conditional inputs can be preprocessed by a conditioning network with a standard feed-forward architecture , which can be learned jointly with the cINN to greatly improve its generative capabilities . We demonstrate the quality of cINNs for conditional image generation and uncover emergent properties of the latent space , for the tasks of conditional MNIST generation and diverse colorization of ImageNet . Given this , we believe that the cINN architecture brings the research field of INNs and other normalizing flow models a substantial step forward . Despite the fact that INNs have received very little attention in the literature in contrast to GANs , we find that cINNs can achieve comparable quality to cGANs , with some remarkable properties absent in cGANs . This includes diverse outputs by default , compared to sparse support or mode collapse in cGANs , easier explainability and interpretability of the learned latent representation , as well as simple and intuitive manipulation and interpolation of generated or existing images . Our work makes the following contributions : • We propose a new architecture called conditional invertible neural network ( cINN ) , which combines an INN with an unconstrained feed-forward network for conditioning . It generates diverse images with high realism on par with existing approaces , while adding some noteworthy and useful properties . • We demonstrate a stable , maximum likelihood-based training procedure for jointly optimiz- ing the parameters of the INN and the conditioning network . • We take advantage of our bidirectional cINN architecture to explore and manipulate emergent properties of the latent space . We illustrate this for MNIST digit generation and image colorization . Quizsolution : Bottomrow , centerimage 2 RELATED WORK . Conditional Generative Modeling . Modern generative models learn to transform noise ( usually sampled from multivariate Gaussians ) into desired target distributions . Methods differ by the modelfamily these transformations are picked from and by the losses determining optimal solutions . Conditional generative adversarial networks ( cGANs , Mirza & Osindero ( 2014 ) ) train a pair of neural networks : a generator transforms a pair of conditioning and noise vectors to images , and a discriminator penalizes unrealistic looking images . The conditioning information is either concatenated to the noise ( Mirza & Osindero , 2014 ) , or fed into the network via conditional batch-norm layers ( Dumoulin et al. , 2017 ; Huang & Belongie , 2017 ; Park et al. , 2019 ) . Ensuring diversity of the generated images ( for fixed conditioning ) appears to be challenging in this approach . Recent BigGANs ( Brock et al. , 2019 ) successfully address this problem by using very large networks and batch sizes , but require parallel training on up to 512 TPUs . PacGANs ( Lin et al. , 2018 ) employ augmented discriminators , which evaluate entire batches of real or generated images together rather than one image at a time . CausalGANs ( Kocaoglu et al. , 2017 ) train two additional discriminator networks , called “ labeler ” and “ anti-labeler ” , with the latter explicitly penalizing the lack of diversity . Pix2pix ( Isola et al. , 2017 ) addresses the important special case when the target is conditioned on an image in a different modality , e.g . to generate satellite images from maps . In addition to the discriminator loss , it minimizes the L1 distance between generated and ground-truth targets using a paired training set , which contains corresponding images from both modalities . This leads to impressive image quality , but lack of diversity seems to be an especially hard problem in this case . In contrast , our method does not need explicit precautions to promote diversity . Bidirectional architectures augment generator networks with complementary encoder networks that learn the generator ’ s inverse and enable reconstruction losses , which exploit cycle consistency requirements . Conditional variational autoencoders ( cVAEs , Sohn et al . ( 2015 ) ) replace all distributions in a standard VAE ( Kingma & Welling , 2013 ) by the appropriate conditional distributions , and are trained to minimize the evidence lower bound ( ELBO loss ) . Since variational distributions are typically Gaussian , the reconstruction penalty is equivalent to squared loss , resulting in rather blurry generated images . This is avoided by AGE networks ( Ulyanov et al. , 2018 ) and CycleGANs ( Zhu et al. , 2017a ) , which combine standard cGAN discriminators with L1 reconstruction loss in the data domain , and bidirectional conditional GANs ( Jaiswal et al. , 2017 ) , which extend the GAN discriminator to act on the distributions in data and latent space jointly . SPADE ( Park et al. , 2019 ) , building upon pix2pix and pix2pixHD ( Wang et al. , 2018 ) , augments cGANs with additional VAE encoders to shape the latent space such that diversity is ensured . Instead of enforcing bijectivity through cycle losses , invertible neural networks are bidirectional by design , since encoder and generator are realized by forward and backward processing within a single bijective model . We focus on architectures whose forward and backward pass require the same computational effort . The coupling layer designs pioneered by NICE ( Dinh et al. , 2014 ) and RealNVP ( Dinh et al. , 2016 ) emerged as very powerful and flexible model families under this restriction . Using additive coupling layers , i-RevNets ( Jacobsen et al. , 2018b ) demonstrated that the lack of information reduction from data space to latent space does not cause overfitting . The Glow architecture ( Kingma & Dhariwal , 2018 ) combines affine coupling layers with invertible 1x1 convolutions and achieves impressive attribute manipulations ( e.g . age , hair color ) in generated faces images . This approach was recently generalized to video by Kumar et al . ( 2019 ) . Thanks to tractable Jacobian determinants , the coupling layer architecture enables maximum likelihood training ( Dinh et al. , 2014 ; 2016 ) , but experimental comparisons with other training methods are inconclusive so far . For instance , Danihelka et al . ( 2017 ) found minimization of an adversarial loss to be superior to maximum likelihood training in RealNVPs , Schirrmeister et al . ( 2018 ) trained i-RevNets in the same manner as adversarial auto-encoders , i.e . with a discriminator acting in latent rather than data space , and Flow-GANs ( Grover et al. , 2018 ) performed best using bidirectional training , a combination of maximum likelihood and adversarial loss . On the other hand , maximum likelihood training worked well within Glow ( Kingma & Dhariwal , 2018 ) , and i-ResNets ( Behrmann et al. , 2018 ) could even be trained with approximated Jacobian determinants . In this work we reinforce the view that high-quality generative models can be trained by maximum likelihood loss alone . To the best of our knowledge , we are the first to apply the coupling layer design for conditional generative models , with the exception of Ardizzone et al . ( 2019 ) , who use it to compute posteriors for ( relatively small ) inverse problems , but do not consider image generation . While Liu et al . ( 2019 ) recently proposed an effective post-hoc conditioning scheme for flow-based generative models , their method requires an additional cGAN to identify latent codes fulfilling the specified condition . Colorization . State-of-the-art regression models for colorization produce visually near-perfect images ( Iizuka et al. , 2016 ) , but do not account for the ambiguity inherent in this inverse problem . To address this , models would ideally define a conditional distribution of plausible color images for a given grayscale input , instead of just returning a single “ best ” solution . Popular existing approaches for diverse colorization predict per-pixel color histograms from a U-Net ( Zhang et al. , 2016 ) or from hypercolumns of an adapted VGG network ( Larsson et al. , 2016 ) . However , sampling from these local histograms independently can not lead to a spatially consistent colorization , requiring additional heuristic post-processing steps to avoid artefacts . In terms of generative models , both VAEs ( Deshpande et al. , 2017 ) and cGANs ( Isola et al. , 2017 ; Cao et al. , 2017 ) have been proposed for the task . However , their solutions do not reach the quality of the regression-based models , and cGANs in particular often lack diversity . To compensate , modifications and extensions to generative approaches have been developed , such as auto-regressive models ( Guadarrama et al. , 2017 ) and CRFs ( Royer et al. , 2017 ) . However , these methods are computationally very expensive and often unable to scale to realistic image sizes . Conceptually closest to our proposed method is the work of ( Ulyanov et al. , 2018 ) , where an encoder network maps color information to a latent space and a generator network learns the inverse transform , both conditioned on the grayscale image . Their experiments however are limited to a data set with only cars , and just three latent dimensions , leading to global , but no local diversity . In contrast to the above , our flow-based cINN generates diverse colorizations in one standard feedforward pass . It models the distribution of all pixels jointly , and allows for meaningful latent space manipulations .

Summary: This paper proposes conditional Invertible Neural Networks (cINN), which introduces conditioning to conventional flow-based generative models. Conditioning is injected into the model via a conditional affine coupling block, which concatenates conditioning with the input to the scaling and shifting sub-networks in the coupling block. Other small modification are proposed to improve training stability at higher learning rates, including soft clamping of scaling coefficients, and Haar wavelet downsampling, which is proposed to replace the squeeze operation (pixel shuffle) that is often used in flow-based models. The invertibility of the cINN allows for style transfer by projecting images into the latent space, and then reconstructing the image with a different conditioning. The performance of the cINN is evaluated empirically on the task of colorization, where it is shown to outperform other techniques in terms of nearness to the true colours, as well as sample diversity.