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: The ultimate goal of the paper is to train an unsupervised generative model, in a semi-supervised manner, as a synthetic training example generator for training segmentation models. Prior works (eg. DatasetGAN) show that one can label a few synthetic images from a pre-trained generative model, and train a segmentation model (called annotation model) over the intermediate representation produced by the generative model. Then, every time a sample is taken from the generative model, the corresponding labels can be produced by applying the annotation model to the corresponding intermediate representations. This provides a valuable way to produce synthetic labelled training examples. However, (some of) these existing approaches require annotating (a limited set of) generated samples for training the annotation model. This is problematic as one needs to re-label such images every time the generative model changes.    This paper aims to address this problem by learning the annotation model over annotated real images instead of annotated synthetic images. To this end, the paper proposes a meta-learning approach. The main idea is to learn the annotation model such that when a (student) segmentation model is trained over the generative model outputs and their labels given by the annotation model, the resulting student model’s loss shall be low on real annotated images. This meta-learning idea is implemented as an approximation by using the Gradient Matching Loss.