Learning Disentangled Representations for Image Translation

1 INTRODUCTION . Learning disentangled representations from a set of observations is a fundamental problem in machine learning . Such representations can facilitate generalization to downstream discriminative and generative tasks as well as improving interpretability ( Hsu et al. , 2017 ) , reasoning ( van Steenkiste et al. , 2019 ) and fairness ( Creager et al. , 2019 ) . Recent advances have contributed to various tasks such as novel image synthesis ( Zhu et al. , 2018 ) and person re-identification ( Eom & Ham , 2019 ) . Image translation is an extensively researched task that benefits from disentanglement . Its goal is to generate an analogous image in a target domain ( e.g . cats ) given an input image in a source domain ( e.g . dogs ) . Although this task is generally poorly specified , it is often satisfied under the assumption that images in different domains share common attributes ( e.g . head pose ) which can be transferred during translation - we name those content . In many cases , the class ( domain ) and common attributes do not uniquely specify the target image e.g . there are many dog breeds with the same head pose . This multi-modal translation motivates the specification of the particular classspecific attributes that we wish the target image to have - we name those style . The ability to transfer the content of a source image to a target class and style has been the objective of several methods e.g . MUNIT ( Huang et al. , 2018 ) , FUNIT ( Liu et al. , 2019 ) and StartGAN-v2 ( Choi et al. , 2020 ) . Unfortunately , we show that despite their visually pleasing results , the translated images still retain many class-specific attributes of the original image resulting in limited translation quality . For example , when translating dogs to wild animals , current methods are prone to transfer facial shapes which are unique to dogs and should not be transferred precisely to wild animals . As demonstrated in Fig . 1 , our model transfers the semantic head pose more reliably . In this work , we analyze the class-supervised setting and present a principled objective for disentangling image class and attributes . We explain why LORD , introduced by ( Gabbay & Hoshen , 2020 ) , can not be applied for multi-modal translation . We then show that introducing an additional style representation overcomes this issue and propose a practical method for high-fidelity image translation by learning disentangled representations . Our method achieves this in two stages ; i ) Disentanglement : disentangled representation learning in a non-adversarial framework , leveraging latent optimization and well-motivated content and style bottlenecks . ii ) Synthesis : the disentangled representations learned in the previous stage are used to ” supervise ” a synthesis network that generalizes to unseen images and classes . As synthesis network training is well-conditioned , we can effectively incorporate a GAN loss resulting in a high-fidelity image translation model . Our approach illustrates that adversarial optimization , which is typically used for domain translation , is not necessary for disentanglement , and its main utility lies in generating perceptually pleasing images . Our model learns disentangled representations and achieves better translation quality and output diversity than current methods . Our contributions are : i ) Introducing a non-adversarial disentanglement method that enables multimodal solutions . ii ) Learning statistically disentangled representations . iii ) Extending domain translation methods to cases with many ( 10k ) domains . iv ) State-of-the-art results in image-translation . 1.1 RELATED WORK . Image Translation Translating the content of images across different domains has attracted much attention . In the unsupervised setting , CycleGAN ( Zhu et al. , 2017 ) introduces a cycle consistency loss to encourage translated images preserves the domain-invariant attributes ( e.g . pose ) of the source image . MUNIT ( Huang et al. , 2018 ) recognized that a given content image could be transferred to many different styles ( e.g . colors and textures ) in a target domain and extends UNIT ( Huang & Belongie , 2017 ) to learn multi-modal mappings by learning style representations . DRIT ( Lee et al. , 2018 ) tackles the same setting using an adversarial constraint at the representation level . MSGAN ( Mao et al. , 2019 ) added a regularization term to prevent mode collapse . StarGAN-v2 ( Choi et al. , 2020 ) and DMIT ( Yu et al. , 2019 ) extend previous frameworks to translation across more than two domains . FUNIT ( Liu et al. , 2019 ) further allows translation to novel domains . Class-Supervised Disentanglement In this parallel line of work , the goal is to anchor the semantics of all the images within each class into a separate class representation while modeling all the remaining class-independent attributes by a content representation . Several methods encourage disentanglement by adversarial constraints ( Denton et al. , 2017 ; Szabó et al. , 2018 ; Mathieu et al. , 2016 ) while other rely on cycle consistency ( Harsh Jha et al. , 2018 ) or group accumulation ( Bouchacourt et al. , 2018 ) . LORD ( Gabbay & Hoshen , 2020 ) takes a non-adversarial approach and trains a generative model while directly optimizing over class and content codes . Most works in this area demonstrate domain translation results on simple datasets but not in the multi-modal ( many-tomany ) settings . Moreover , their focus is to achieve disentanglement at the representation level rather than designing architectures for high-resolution image translation resulting in weak performance on competitive benchmarks . In this work , we draw inspiration from LORD in relying on the inductive bias conferred by latent optimization to learn a disentangled content representation . In contrast to LORD , we tackle the multi-modal image translation setting by modeling style . Moreover , we add an adversarial term in the synthesis network that increases the image quality and resolution . 2 BACKGROUND : REPRESENTATION LEARNING IN IMAGE-TRANSLATION . Image translation takes as input a set of N images and corresponding class labels ( x1 , y1 ) , ( x2 , y2 ) , ... , ( xN , yN ) . Let us assume that an image xi is fully specified by its class yi and attributes ai . As a motivational example , let us consider the images xi to be of animals , and the class label yi denotes the species . Attributes ai , may include attributes aci common to all classes such as head angle , and class-specific attributes asi such as dog or cat breed . An unknown function G∗ maps yi , aci and a s i to the image xi : xi = G ∗ ( yi , a s i , a c i ) ( 1 ) The goal of image translation is to replace the common attributes aci of target xi by a c j of source xj : xij = G ∗ ( yi , a s i , a c j ) ( 2 ) The main challenge is that during training , only the class yi of each image is given , but the attributes aci and a s i are unknown . We define three representations corresponding to each of the physical properties above - the class embedding eyi represents the class yi , the content code ci represents the common attributes aci and the style code si represents the class-specific attributes a s i . The computational task is to learn the representations eyi , si , ci such that they faithfully encode their corresponding physical properties of image xi and are sufficient to reconstruct the image using a generator G : xi = G ( eyi , si , ci ) ( 3 ) This constraint however is insufficient to specify the representations uniquely . A popular constraint , utilizes the fact that common attributes aci are independent of the class yi . In the motivating animals example , we assume that head pose acts similarly on cats or dogs . It therefore requires independence between the content code ci and class yi . Adversarial methods : Adversarial methods implement this constraint by training a domain confusion discriminator on the translated image xij = G ( eyi , si , cj ) . The discriminator attempts to distinguish translated images from original images of the target class . Unfortunately , we empirically show that state-of-the-art methods do not learn content codes ci that are disentangled from the class yi . The consequence is that ci transfers not only common attributes but also class-dependent attributes . We hypothesize that this failure is due to the challenging adversarial optimization . In order to constrain the style si , current methods ( e.g . StarGAN-v2 , FUNIT ) rely on localitypreserving architectures that bias the content code ci to represent the structure , while the style operates in a global manner per channel and is applied at a relatively deep level within the generator ( typically from 16×16 spatial resolution ) . Unfortunately , such style is unable to model class-specific attributes asi of spatial nature e.g . the spatial facial shape of different dog breeds . Non-adversarial methods : LORD is a non-adversarial approach for disentangled representation learning . It assumes that all attributes are common to all classes I ( ac ; y ) = 0 and H ( as ) = 0 - where I denotes mutual information and H denotes entropy . It therefore learns only class and content representations eyi and ci but not style codes ( si = 0 ) . In order to learn disentangled class and content representations it uses a simple but effective constraint of minimal information in the content codes ci . After learned disentangled representations eyi and ci , LORD learns a synthesis network which generalizes to unseen images and classes . When the assumptions of LORD are satisfied ( H ( as ) = 0 ) , the content codes ci are indeed shown to recover the information of the common attributes aci . However , when class-specific attributes exist , LORD does not achieve its goals . The reason is that the content code ci , contains the information of both the common aci and class specific a s i attributes ( I ( a s ; c ) > 0 ) . This has two issues : i ) as the content code now contains class-specific information ( e.g . cat breed ) , it is no longer disentangled from the class ( we can predict if the content is of a cat ) ; ii ) when transferring content codes across classes ( e.g . cats to dogs ) , we are missing the required class specific information in the target class ( i.e . given the class code of edog and content of a cat ccat , the generator is not provided with the information of what dog breed to generate ) . Due to its effective non-adversarial approach , LORD learns disentangled representations and achieves state-of-the-art results on uni-modal synthetic and low-res datasets . However , due to i ) ignoring class-specific attributes as highlighted above ; ii ) using a non-adversarial synthesis network , it does not produce high fidelity images and can not scale to high-resolution . 3 OVERLORD : PRINCIPLED DISENTANGLEMENT IN IMAGE TRANSLATION . Based on the above analysis , we present our disentangled image-translation method , OverLORD . 3.1 LEARNING A DISENTANGLEMENT MODEL . In order to learn disentangled representations for class , content and style , we propose a nonadversarial framework with several principles . Reconstruction : As the image should be fully specified by the representations eyi , si , ci , we employ a reconstruction term : Lrec = ` ( G ( eyi , si , ci ) , xi ) - where ` is a measure of similarity between images ( we use the VGG perceptual loss ) . Content bottleneck : In order to constrain the amount of information in each content variable ci , we parameterize it as a noisy channel consisting of a vector c′i and an additive Normal noise z ∼ N ( 0 , I ) , ci = c′i + z . Using the Shannon-Hartley theorem , the information capacity of the channel is a function of the signal-to-noise ratio . For the noisy channel ci , the SNR ratio is given by ‖c′i‖2 . We therefore define the content bottleneck loss as Lcb = ∑ i ‖c′i‖2 . Style : We train a style encoderEs : X −→ S to infer the class-specific attributes as . To encourage invariance to the common attributes ac , the input image first undergoes a random transformation . This creates a transformed version xtransi of the input image that shares the same class specific attributes as , but exhibits random common attributes ac ( as style should be invariant to these attributes ) . The nature of the transformation is setting dependent e.g . if common attributes include spatial locations , then crop and rotation transformations will remove spatial attributes from the transformed image and make the style code si invariant to them . si = Es ( x trans i ) ( 4 ) Our complete objective can be summarized as follows : min c′i , eyi , Es , G Ldisent = ∑ i ` ( G ( eyi , Es ( x trans i ) , c ′ i + z ) , xi ) + λcb‖c′i‖2 z ∼ N ( 0 , I ) ( 5 ) There are several fundamental differences from the cVAE objective : i ) the noise variance is a fixed hyper-parameter and is not learned ; ii ) there is no content encoder , instead the latent content code for each variable is learned directly , in an unamortized way often referred to as latent optimization ; iii ) the input to the style encoder is a transformed version of the image . Gabbay & Hoshen ( 2020 ) discovered that latent optimization improves disentanglement significantly over encoder-based methods . The difference lies in the initialization as latent optimization starts with no class-content correlation ( code initialization is IID ) while a content encoder starts with near perfect correlation ( class can be predicted from the output of a randomly initialized content encoder ) . We further elaborate on latent optimization and its inductive bias in Appendix A.3 .

Summary: The paper presents a principled approach to style transfer by disentangling class-specific attributes from common (eq. class-independent) attributes. In order to do so, the paper leverages the formulation of a recently proposed disentangling approach called "LORD". The proposed approach is called OverLORD, and includes two main augmentations to LORD. The first is the introduction of a style encoder to learn a latent code for class-specific attributes, and the second is to introduce an adversarial learning in the second stage for high-quality style-transferred generation of images. The results are shown on three datasets: AFHQ (dog, cat, wildlife), CelebA (human faces), and CelebA-HQ (hi-res human faces), and are compelling in both qualitative and quantitative comparisons.