Semantic Segmentation With Generative Models: Semi ...

1 Semantic Segmentation with Generative Models: Semi-Supervised Learning and Strong Out-of-Domain GeneralizationDaiqing Li1*Junlin Yang1,3 Karsten Kreis1 Antonio Torralba4 Sanja Fidler1,2,51 NVIDIA2 University of Toronto3 Yale University4 MIT5 Vector InstituteAbstractTraining deep networks with limited labeled data whileachieving a strong generalization ability is key in the questto reduce human annotation efforts. This is the goal ofsemi-supervised learning, which exploits more widely avail-able unlabeled data to complement small labeled data this paper, we propose a novel framework for discrim-inative pixel-level tasks using a Generative model of bothimages and labels. Concretely, we learn a Generative ad-versarial network that captures the joint image -label dis-tribution and is trained efficiently using a large set of un-labeled images supplemented with only few labeled build our architecture on top of StyleGAN2 [45], aug-mented with a label synthesis branch. image labeling attest time is achieved by first embedding the target imageinto the joint latent space via an encoder network and test-time optimization, and then generating the label from the in-ferred embedding.

2 We evaluate our approach in two impor-tant domains: medical image Segmentation and part-basedface Segmentation . We demonstrate strong in-domain per-formance compared to several baselines, and are the firstto showcase extreme out-of-domain generalization, such astransferring from CT to MRI in medical imaging, and pho-tographs of real faces to paintings, sculptures, and evencartoons and animal faces. Project Page: IntroductionDeep learning is now powering the majority of com-puter vision applications ranging from autonomous driv-ing [93,73] and medical imaging [78,38] to image edit-ing [69,15,98,88,70,74]. However, deep networks areextremely data hungry, typically requiring training on large-scale datasets to achieve high accuracy. Even when largedatasets are available, generalizing the network s perfor-mance to out-of-distribution data, for example, on imagescaptured by a different sensor, presents challenges, sincedeep networks tend to overfit to artificial statistics in the*Correspondence Test ImagesTrainingxrayCTMRIF igure 1:Out-of-domain model trained on realfaces generalizes to paintings, sculptures, cartoons and even outputs plau-sible segmentations for animal faces.

3 When trained on chest x-rays, it gen-eralizes to multiple hospitals, and even hallucinates lungs under clothedpeople. Our model also generalizes well from CT to MRI medical data. Labeling large datasets, particularly for densepixel-level tasks such as Semantic Segmentation , is alreadyvery time consuming. Re-doing the annotation effort eachtime the sensor changes is especially undesirable. This isparticularly true in the medical domain, where pixel-levelannotations are expensive to obtain (require highly-skilledexperts), and where imaging sensors vary across sites. Inthis paper, we aim to significantly reduce the number oftraining data required for attaining successful performance,while achieving strong out-of-domain learning (SSL) facilitates learning withsmall labeled data sets by augmenting the training set withlarge amounts of unlabeled data. The literature on SSL isvast and some classical SSL techniques include pseudo-labeling [50,2,8,83], consistency regularization [80,49,87,23,83], and various data augmentation techniques [7,6,91] (also see ).

4 State-of-the-art SSL performanceis currently achieved by contrastive learning, which aims totrain powerful image feature extractors using unsupervisedcontrastive losses on image transformations [13,29,61,32].Once the feature extractors are trained, a smaller amount oflabels is needed, since the features already implicitly en-code Semantic information. While SSL approaches havebeen more widely explored for classification, recent meth-ods also tackle pixel-wise tasks [37,62,40,47,22,68].Although SSL techniques allow to train models with lit-tle labeled data, they usually do not explicitly model the dis-8300tribution of the input data itself and therefore can still easilyoverfit to the training data, hampering their generalizationcapabilities. This is especially critical in Semantic segmen-tation, where annotations are expensive and hence the avail-able amount of labeled data can be particularly address this, we propose a fully Generative approachbased on a Generative adversarial network (GAN) that mod-els thejointimage-label distribution and synthesizes bothimages and their Semantic Segmentation masks.

5 We buildon top of the StyleGAN2 [45] architecture and augment itwith a label generation branch. Our model is trained on alarge unlabeled image collection and a small labeled sub-set using only adversarial objectives. Test-time predictionis framed as first optimizing for the latent code that recon-structs the input image , and then synthesizing the label byapplying the generator on the inferred showcase our method in the medical domain andon human faces. It achieves competitive or better in-domain performance even when compared to heavily engi-neered state-of-the-art approaches, and shows significantlyhigher generalization ability on out-of-domain tests. Wealso demonstrate the ability to generalize to domains thatare drastically different from the training domain, such asgoing from CT to MRI volumes, and natural photographsof faces to sculptures, paintings and cartoons, and even an-imal faces (see Figure1).In summary, we make the following contributions: (i)We propose a novel Generative model for Semantic segmen-tation that builds on the state-of-the-art StyleGAN2 andnaturally allows semi-supervised training.

6 To the best ofour knowledge, we are the first work that tackles semanticsegmentation with a purely Generative method that directlymodels the joint image -label distribution. (ii) We exten-sively validate our model in the medical domain and on faceimages. In the semi-supervised setting, we demonstrate re-sults equal to or better than available competitive baselines.(iii) We show strong generalization capabilities and outper-form our baselines on out-of-domain Segmentation tasks bya large margin. (iv) We qualitatively demonstrate reason-able performance even on extreme out-of-domain Related WorkOur paper touches upon various topics, including medi-cal image analysis, Semantic Segmentation , semi-supervisedlearning, Generative modeling and neural network Learning and Semantic Segmenta-tion: in the medical domain, semi-supervised Semantic seg-mentation has been tackled via pseudo-labeling [2], adver-sarial training [64,53], and transformation-consistency [53]in a mean-teacher framework [87].

7 In computer vision, [59]is the first work using an adversarial objective to train a seg-mentation network. Later this idea was extended to semi-supervised setups via self-taught losses and discriminatorfeature matching [37,62]. Recently, [47] proposed an ap-proach using a flaw detector to approximate pixel-wise pre-diction confidence. Further relevant approaches to semi-supervised Segmentation have been developed in weakly-supervised setups [35,51,94].For simpler classification tasks, a plethora of SSL meth-ods have been developed, based on pseudo-labeling [50,8,72], self-supervision [8], entropy-minimization [27],consistency-regularization [80,49,87,23], adversarialtraining [63], data augmentation [91], and combinationsthereof [7,83,6]. However, current state-of-the-art semi-supervised methods are based on self-supervised learningwith contrastive objectives [13,29,61,32]. These ap-proaches use unlabeled data in an often task-agnostic man-ner to learn general feature representations that can be fine-tuned using a smaller amount of labeled data.

8 Relatedideas have been applied to semi-supervised Semantic seg-mentation [40] and tailored data augmentation strategieshave been explored [22,68]. Furthermore, many worksemploy carefully designed pretext tasks to learn useful rep-resentations from unlabeled images [18,66,24,26]. Ourmethod is related to these works in the sense that our taskfor learning strong features is image generation itself, in-stead of an auxiliary pretext above works train discriminative models of theformp(y|x), in contrast to our fully Generative , Generative approaches to SSL have been pro-posed before. [48] leverages variational autoencoders and[67,81] use GANs in which the discriminator distinguishesbetween different classes. A related approach to semi-supervised Semantic Segmentation uses Generative modelsto augment the training data with additional synthesizeddata [84,41,17,52]. Conceptually, [84] trains a generatortogether with a pixel-wise discriminator network to performsegmentation, while [41,17,52] learn to generate synthetic3D scenes by matching distributions of real and renderedimagery.

9 In parallel work, [97] exploit GANs to synthesizelarge labeled data datasets using very few labeled contrast, to the best of our knowledge, our method isthe first fully Generative approach to Semantic segmentationthat uses only adversarial objectives and no cross entropyterms and in which the generator models thejointp(x, y)and directly synthesizes images together with pixel-wise la-bels. We further use the Generative model as a decoder ofsemantic outputs at test time, which we show leads to bettergeneralization than prior and parallel Inversion: A critical part of our method isthe effective inversion of the GAN generator at test-time toinfer the latent embedding of a new image to be labeled. Weare building on previous works that have studied this taskbefore. Optmization-based methods iteratively optimize areconstruction objective [99,92,56,1,36,16,75,74] orperform Markov chain Monte Carlo [21], while encoder-8301 + W+x<latexit sha1_base64="hL+FaLtOT9luwfLW3Ut08xl3 Pcw=">AAAB6 HicbVDLTgJBEOzFF+IL9ehlIjHxRHbRRI9 ELx4hkUcCGzI79 MLI7 OxmZtZICF/gxYPGePWTvPk3 DrAHBSvppFLVne6uIBFcG9f9dnJr6xubW/ntws7u 3v5B8fCoqeNUMWywWMSqHVCN gktsGG4 EthOFNAoEtoLR7cxvPaLSPJb3 ZpygH9GB5 CFn1 Fip/tQrltyyOwdZJV5 GSpCh1it+dfsxSyOUhgmqdcdzE+NPqDKcCZwWuqn GhLIRHWDHUkkj1P5kfuiUnFmlT8JY2 ZKGzNXfExMaaT2 OAtsZUTPUy95M/M/rpCa89idcJqlByRaLwlQQE5P Z16 TPFTIjxpZQpri9lbAhVZQZm03 BhuAtv7xKmpWyd1Gu1C9L1 ZssjjycwCmcgwdXUIU7qEEDGCA8wyu8OQ/Oi/Puf Cxac042cwx/4Hz+AOeHjQA=</latexit>x0<latexit sha1_base64="fE3 ASVSeVbJpP9+EYmbYOaSFZRg=">AAAB6 XicbVDLTgJBEOzFF+IL9ehlIjF6 IrtookeiF49o5 JEAIbNDL0yYnd3 MzBrJhj/w4kFjvPpH3vwbB9iDgpV0 UqnqTneXHwuujet+O7mV1bX1jfxmYWt7Z3evuH/Q 0 FGiGNZZJCLV8qlGwSXWDTcCW7 FCGvoCm/7oZuo3H1 FpHskHM46xG9KB5 AFn1 Fjp/um0 Vyy5 ZXcGsky8jJQgQ61X/Or0I5aEKA0 TVOu258amm1 JlOBM4 KXQSjTFlIzrAtqWShqi76ezSCTmxSp8 EkbIlDZmpvydSGmo9Dn3bGVIz1 IveVPzPaycmuOqmXMaJQcnmi4

10 JEEBOR6dukzxUyI8aWUKa4vZWwIVWUGRtOwYbgLb 68 TBqVsndertxdlKrXWRx5 OIJjOAMPLqEKt1 CDOjAI4 Ble4c0 ZOS/Ou/Mxb8052cwh/IHz+QNIDo0x</latexit>y0<latexit sha1_base64="0rMzGhO6lakG6u8bnTFO/7ooUVY=">AAAB6 XicbVBNS8 NAEJ34 WetX1aOXxSJ6 KkkV9Fj04rGK/YA2lM120y7dbMLuRCih/8 CLB0W8+o+8+W/ctjlo64 OBx3szzMwLEikMuu63s7K6tr6xWdgqbu/s7u2 XDg6bJk414w0Wy1i3A2q4 FIo3 UKDk7 URzGgWSt4LR7dRvPXFtRKwecZxwP6 IDJULBKFrpYXzWK5 XdijsDWSZeTsqQo94rfXX7 MUsjrpBJakzHcxP0M6pRMMknxW5qeELZiA54x1 JFI278bHbphJxapU/CWNtSSGbq74mMRsaMo8B2 RhSHZtGbiv95nRTDaz8 TKkmRKzZfFKaSYEymb5O+0 JyhHFtCmRb2 VsKGVFOGNpyiDcFbfHmZNKsV76 JSvb8s127yOApwDCdwDh5cQQ3uoA4 NYBDCM7zCmzNyXpx352 PeuuLkM0fwB87nD0mTjTI=</latexit>Ls<latexit sha1_base64="WJ4+ma/b/ayQw59Ju91t78XA9qw=">AAAB6nicbVA9 SwNBEJ3zM8avqKXNYhCswl0 UtAzaWFhENB+QHGFvM5cs2ds7dveEcOQn2 FgoYusvsvPfuEmu0 MQHA4/3 ZpiZFySCa+O6387K6tr6xmZhq7i9s7u3 Xzo4bOo4 VQwbLBaxagdUo+ASG4 Ybge1 EIY0 Cga1gdDP1W0+oNI/loxkn6Ed0 IHnIGTVW erjr6V6p7 FbcGcgy8 XJShhz1 Xumr249 ZGqE0 TFCtO56bGD+jynAmcFLsphoTykZ0gB1 LJY1Q+9ns1Ak5tUqfhLGyJQ2 Zqb8nMhppPY4C2xlRM9SL3lT8z+ukJrzyMy6T1 KBk80 VhKoiJyfRv0ucKmRFjSyhT3N5K2 JAqyoxNp2hD8 BZfXibNasU7r1 TvL8q16zyOAhzDCZyBB5dQg1uoQwMYDOAZXuHNEc 6L8+58zFtXnHzmCP7A+fwBLf6 Nug==</latexit>Lu<latexit sha1_base64="cqXXK+kgVE/xlkca2 Ksox5 Uhf+M=">AAAB6nicbVA9 SwNBEJ3zM8avqKXNYhCswl0 UtAzaWFhENB+QHGFvM5cs2ds7dveEcOQn2 FgoYusvsvPfuEmu0 MQHA4/3 ZpiZFySCa+O6387K6tr6xmZhq7i9s7u3 Xzo4bOo4 VQwbLBaxagdUo+ASG4 Ybge1 EIY0 Cga1gdDP1W0+oNI/loxkn6Ed0 IHnIGTVW erjrpb1S2a24M5Bl4uWkDDnqvdJXtx+zNEJpmKBa dzw3MX5 GleFM4 KTYTTUmlI3oADuWShqh9rPZqRNyapU+CWNlSxoyU 39 PZDTSehwFtjOiZqgXvan4n9dJTXjlZ1wmqUHJ5ov CVBATk+nfpM8 VMiPGllCmuL2 VsCFVlBmbTtGG4C2+vEya1Yp3 XqneX5Rr13kcBTiGEzgDDy6hBrdQhwYwGMAzvMKb I5wX5935mLeuOPnMEfyB8/kDMQaNvA==</latexit>E<latexit sha1_base64="oW7 OyFl8fZOuY5qYrNXg+guGckE=">AAAB6 HicbVDLSgNBEOyNrxhfUY9eBoPgKexGQY9 BETwmYB6 QLGF20puMmZ1dZmaFEPIFXjwo4tVP8ubfOEn2oIk FDUVVN91dQSK4Nq777eTW1jc2t/LbhZ3dvf2D4uF RU8epYthgsYhVO6