Learning to Annotate Part Segmentation with Gradient Matching

1 INTRODUCTION . In recent years , deep neural networks have shown a remarkable ability to learn complex visual concepts from large quantities of labeled data . However , collecting manual annotations for a growing body of visual concepts remains an extremely expensive practice , especially for pixelwise prediction tasks such as semantic segmentation . A common and effective way of reducing the dependence on manual labels is first learning representations from unlabeled data ( Pathak et al. , 2016 ; Noroozi & Favaro , 2016 ; Gidaris et al. , 2018 ; He et al. , 2020 ) and then deploying learned representations to a supervised learning task ( Chen et al. , 2020 ; Zhai et al. , 2019 ) . The insight of this paradigm is that some of the transferred representations are informative , easily identified , and disentangled , such that small amount of labeled data suffices to train the target task . A promising direction for learning such representations is using generative models ( Kingma & Welling , 2014 ; Goodfellow et al. , 2014 ) , since they attain the capability of synthesizing photorealistic images and disentangling factors of variation in dataset . Recent methods ( Goetschalckx et al. , 2019 ; Shen et al. , 2020 ; Karras et al. , 2019 ; Plumerault et al. , 2020 ; Jahanian et al. , 2020 ; Voynov & Babenko , 2020 ; Spingarn-Eliezer et al. , 2021 ; Härkönen et al. , 2020 ) show that certain factors such as object shape and position in the images synthesized by generative adversarial networks ( GANs ) ( Goodfellow et al. , 2014 ) can be individually controlled by manipulating latent features . Building on the success of powerful GAN models , Zhang et al . ( 2021 ) and Tritrong et al . ( 2021 ) demonstrate that feature maps in StyleGAN ( Karras et al. , 2019 ; 2020b ; a ) can be mapped to semantic segmentation masks or keypoint heatmaps through a shallow decoder , denoted as annotator in this paper . Remarkably the annotator can be trained after few synthesized images are labeled , and the joint generator and trained annotator can be used as an infinite labeled data generator . Nonetheless , labeling synthesized images with “ human in the loop ” has two shortcomings : ( i ) it can be difficult to label the generated images of low quality ; ( ii ) it requires continual human effort whenever the generator is changed . To mitigate these limitations , a natural idea is to learn annotator from a collection of image-label pairs . This is non-trivial because annotator only receives as input generator features which are lacked in labeled image datasets . One straightforward way to address this issue is to inverse the generation process to obtain generator features for given images . However , as suggested by our experimental results , inversion methods can not ensure exact recovery of the generator features , leading to degraded performance of annotator . Another way to address this issue is to align the joint distribution of generated images and labels with that of real ones using adversarial learning as in SemanticGAN ( Li et al. , 2021 ) . However , adversarial learning is notoriously unstable ( Brock et al. , 2018 ) and requires sufficient data to prevent discriminator from overfitting ( Karras et al. , 2020a ) . In this paper , we formulate the annotator learning as a learning-to-learn problem – the annotator learns to label a set of randomly generated images such that a model , denoted as student , trained on these automatically labeled images obtains superior performance evaluated on a small set of manually labeled images . This problem involves solving a nested-loop optimization , where the inner loop optimizes the student and the outer loop optimizes the annotator based on the solution of optimization of student . In contrast to existing learning-to-learn algorithms ( Li et al. , 2019 ; Pham et al. , 2021 ) that employs end-to-end gradient-based meta learning techniques , we reduce this complex and expensive optimization problem into a simple gradient matching problem and efficiently solve it with an iterative algorithm . We show that our method obtains performance comparable to the supervised learning counterparts ( Zhang et al. , 2021 ; Tritrong et al. , 2021 ) yet overcomes their shortcomings such that labeled data within a broader range can be utilized . Experiments also show that our method achieves better performance than the naive inversion method and SemanticGAN ( Li et al. , 2021 ) . Our method requires a large quantity of unlabeled data for pre-training GANs and a relatively small amount of labeled data for training annotator , which suits a semi-supervised learning setting . Furthermore , considering that state-of-the-art unconditional GANs only produce appealing results on single-class images but struggle to model complex scenes , we evaluate our method on a variety of part segmentation datasets of particular classes such as human face , horse , aeroplane , car , and cat . Our contribution can be summarized as follows . ( i ) We formulate the learning of annotations for GANgenerated images as a learning-to-learn problem and propose an algorithm based on gradient matching to solve it . As a consequence , a broad range of labeled data including real data , synthetic data , and even out-of-domain data is applicable . ( ii ) We empirically show that our method significantly outperforms other semi-supervised segmentation methods under the circumstances of extremely limited labeled data . 1 2 RELATED WORK . Semi-supervised learning Semi-supervised learning ( SSL ) augments the training of neural networks from a small amount of labeled data with a large-scale unlabeled dataset . A rich body of work regularize networks on unlabeled data with consistency regularization . In particular , networks are required to make consistent prediction over previous training iterations ( Laine & Aila , 2017 ; Tarvainen & Valpola , 2017 ; Izmailov et al. , 2018 ) , noisy inputs ( Miyato et al. , 2018 ) , or augmented inputs ( Berthelot et al. , 2019b ; a ; Sohn et al. , 2020 ) . This paradigm is also interpreted as a teacherand-student framework where teacher produces pseudo labels for unlabelled data to supervise the training of students . Pseudo labels can be enhanced with heuristics ( Lee et al. , 2013 ; Berthelot et al. , 2019a ) or optimized towards better generalization ( Pham et al. , 2021 ; Li et al. , 2019 ) . Apart from image classification tasks , the above ideas are also successfully applied to semantic segmentation tasks ( Hung et al. , 2018 ; Ouali et al. , 2020 ; Ke et al. , 2020 ; French et al. , 2020 ) . Our work can also be interpreted from teacher-and-student perspective . In contrast to the above methods that adopt homogeneous network structures for teacher and student , our work employ generative models as teacher and discriminative models as student . Other SSL approaches explore the utility of unlabelled data from representation learning perspective . In particular , deep neural networks are first optimized on the unlabeled data with self-supervised learning tasks ( Pathak et al. , 2016 ; Noroozi & Favaro , 2016 ; Gidaris et al. , 2018 ; He et al. , 2020 ) to gain powerful and versatile representations , and then finetuned on the labeled data to perform target tasks ( Chen et al. , 2020 ; Zhai et al. , 2019 ) . In addition , efficient and transferable representations can also be learned with generative models that are typically trained without labels . To this end , one promising direction is to decode generative features into dense annotations such as segmentation 1Code will be made public . masks and keypoint heatmap to facilitate SSL , as done in SemanticGAN ( Li et al. , 2021 ) , DatasetGAN ( Zhang et al. , 2021 ) , RepurposeGAN ( Tritrong et al. , 2021 ) and our work . However , our work differs from these works in the way that we formulate the annotator learning as a learning-to-learn problem which is solved with gradient matching . Particualrly , unlike SemanticGAN ( Li et al. , 2021 ) that learns annotator and data generator jointly , our method presumes a fixed pre-trained GAN and learns annotator only , which circumvents the complication of joint learning . Semantic Part Segmentation Semantic part segmentation provides dense form of representations for both rigid and non-rigid objects , which is of great significance in tremendous computer-vision tasks such as human parsing ( Dong et al. , 2014 ; Nie et al. , 2018 ; Fang et al. , 2018 ; Gong et al. , 2017 ) , pose estimation ( Yang & Ramanan , 2011 ; Shotton et al. , 2011 ; Xia et al. , 2017 ) and 3D object understanding ( Yi et al. , 2016 ; Song et al. , 2017 ) . The challenge of part segmentation arises from viewpoint changes , occlusion , and lack of clear boundaries for certain parts . Existing datasets manage to alleviate the problem by enrich the data and annotations . CelebAMask-HQ ( Lee et al. , 2020a ) is a large-scale face image dataset annotated with masks of face components . Chen et al . ( 2014 ) provides additional part segmentation upon PASCAL VOC dataset ( Everingham et al. , 2010 ) . Those datasets meet the demands of latest data-driven and deep neural network based method ( Lee et al. , 2020b ; Wang et al. , 2015 ; Wang & Yuille , 2015 ) . However , a well-curated part segmentation dataset containing fine-grained pixel-wise annotation usually requires time-consuming manual labeling work , which holds back the development of this field . This problem also motivates the latest work ( Zhang et al. , 2021 ; Li et al. , 2021 ; Tritrong et al. , 2021 ) to leverage a generative model as an infinite data generator in a semi-supervised learning setting , which gain promising results . Gradient matching Gradient matching is an important technique in meta learning ( Li et al. , 2018 ; Sariyildiz & Cinbis , 2019 ) , which is shown to be effective in domain generalization ( Li et al. , 2018 ) , zero-shot classification ( Sariyildiz & Cinbis , 2019 ) , and dataset condensation ( Zhao et al. , 2021 ; Zhao & Bilen , 2021 ) etc . Although the general principle and gradient matching loss function in our work resemble those in the above works , we are motivated to solve a totally different problem . In particular , Gradient Matching Network ( GMN ) ( Sariyildiz & Cinbis , 2019 ) employs gradient matching to train a conditional generative model towards application on zero-shot classification . Dataset Condensation ( DC ) ( Zhao et al. , 2021 ; Zhao & Bilen , 2021 ) condenses a large-scale dataset into a small synthetic one with gradient matching . Our work shows that gradient matching is also an effective way to learn annotations of GAN-generated images from limited labeled examples . Moreover , in contrast to DC where the synthesized images are typically not realistic at all , our work are concerned about photo-realistic synthetic images . 3 METHOD . 3.1 PRELIMINARY . Let G denote a generative model trained from a dataset Dtot without labels . For example , G can be trained via adversarial learning ( Goodfellow et al. , 2014 ) such that it can generate photo-realistic images . The generation process is formulated as ( h , x ) = G ( z ) , z ∼ Pz ( 1 ) where z denotes a random variable drawn from a prior distribution Pz which can be a uniform or Gaussian distribution , x denotes the generated image , and h denotes the hidden features in G. An annotator Aω is constructed to decode h into labels ŷ = Aω ( h ) , ( 2 ) where ω parameterizes the annotator . Let Dl denote a manually labeled dataset . Given a fixed G , our problem is to learn annotator Aω from Dl such that Aω can annotate images generated from G. Revisiting supervised learning manner DatasetGAN ( Zhang et al. , 2021 ) and RepurposeGAN ( Tritrong et al. , 2021 ) train annotators in a supervised learning manner . First , a set of synthetic images { xi } Ni=1 generated from G are selected and annotated by humans and their generator features { hi } Ni=1 are reserved . These features and manual labels { yi } Ni=1 constitute the labeled dataset Dl = { ( xi , hi , yi ) } Ni=1 . Second , an annotator is trained by minimizing the loss computed between annotator prediction and ground truth labels , ω = argmin ω E ( · , h , y ) ∼Dlf ( Aω ( h ) , y ) , ( 3 ) where f denotes the loss function for specific tasks , e.g . per-pixel cross-entropy function for segmentation tasks . This method , however , requires Dl to be a synthetic and generator-specific set , which has the following two defects . ( i ) Details might be illegible in synthetic images , which increases the difficulty of annotation . ( ii ) Data needs to be re-labeled whenever the generator is changed . Although projecting labeled images into GAN latent space can alleviate the above drawbacks , it leads to degenerate performance due to inexact inversion ( see Section C.1 in the appendix ) .

Summary: This work deals with the problem of training a part segmentation network by automatically synthesizing pairs of images and annotations. The authors propose a training method of an annotator given the well-pretrained generator model. The formulation is well-motivated and converged to the gradient matching loss.   This work extends [Zhao et al., ICLR 2021] by combining the idea of DatasetGAN and RepurposeGAN in an interesting way. The proposed approach can favor the unlimited number of annotated data with only limited supervised real data + the pre-trained GAN.  The proposed method shows noticeable improvement over the other baselines and the competing methods in scarcely labeled data regimes (especially semi-supervised learning).