Label-Efficient Semantic Segmentation with Diffusion Models

1 INTRODUCTION . Denoising diffusion probabilistic models ( DDPM ) ( Sohl-Dickstein et al. , 2015 ; Ho et al. , 2020 ) have recently outperformed alternative approaches to model the distribution of natural images both in the realism of individual samples and their diversity ( Dhariwal & Nichol , 2021 ) . These advantages of DDPM are successfully exploited in applications , such as colorization ( Song et al. , 2021 ) , inpainting ( Song et al. , 2021 ) , super-resolution ( Saharia et al. , 2021 ; Li et al. , 2021b ) , and semantic editing ( Meng et al. , 2021 ) , where DDPM often achieve more impressive results compared to GANs . So far , however , DDPM were not exploited as a source of effective image representations for discriminative computer vision problems . While the prior literature has demonstrated that various generative paradigms , such as GANs ( Donahue & Simonyan , 2019 ) or autoregressive models ( Chen et al. , 2020a ) , can be used to extract the representations for common vision tasks , it is not clear if DDPM can also serve as representation learners . In this paper , we provide an affirmative answer to this question in the context of semantic segmentation . In particular , we investigate the intermediate activations from the U-Net network that approximates the Markov step of the reverse diffusion process in DDPM . Intuitively , this network learns to denoise its input , and it is not clear why the intermediate activations should capture semantic information needed for high-level vision problems . Nevertheless , we show that on the certain diffusion steps , these activations do capture such information , therefore , can potentially be used as image representations for downstream tasks . Given these observations , we propose a simple semantic segmentation method , which exploits these representations and successfully works even if only a few labeled images are provided . On several datasets , we show that our DDPM-based segmentation method outperforms the existing baselines for the same amount of supervision . To sum up , the contributions of our paper are : 1 . We investigate the representations learned by the state-of-the-art DDPM and show that they capture high-level semantic information valuable for downstream vision tasks . 1anonymous URL 2 . We design a simple segmentation approach that exploits these representations and outperforms the alternatives in the few-shot operating point . 3 . We compare the DDPM-based representations with their GAN-based counterparts on the same datasets and demonstrate the advantages of the former in the context of semantic segmentation . 2 RELATED WORK . In this section , we briefly describe the existing research lines relevant to our work . Diffusion models ( Sohl-Dickstein et al. , 2015 ; Ho et al. , 2020 ) are a class of generative models that approximate the distribution of real images by the endpoint of the Markov chain that originates from a simple parametric distribution , typically a standard Gaussian . Each Markov step is modeled by a deep neural network that effectively learns to invert the diffusion process with a known Gaussian kernel . ( Ho et al. , 2020 ) highlighted the equivalence of diffusion models and score matching ( Song & Ermon , 2019 ; 2020 ) , showing them to be two different perspectives on the gradual conversion of a simple known distribution into a target distribution via the iterative denoising process . Very recent works ( Nichol & Dhariwal , 2021 ; Dhariwal & Nichol , 2021 ) have developed more powerful model architectures as well as different advanced objectives for DDPM , which led to the “ victory ” of DDPM over GANs in terms of generative quality and diversity . DDPM have been widely used in several applications , including image colorization ( Song et al. , 2021 ) , super-resolution ( Saharia et al. , 2021 ; Li et al. , 2021b ) , inpainting ( Song et al. , 2021 ) , semantic editing ( Meng et al. , 2021 ) . In our work , we demonstrate that one can also successfully use them for semantic segmentation . Image segmentation with generative models is an active research direction at the moment , however , existing methods are primarily based on GANs . The first line of works ( Voynov & Babenko , 2020 ; Voynov et al. , 2021 ; Melas-Kyriazi et al. , 2021 ) is based on the evidence that the latent spaces of the state-of-the-art GANs have directions corresponding to effects that influence the foreground/background pixels differently , which allows producing synthetic data to train segmentation models . However , the approaches from ( Voynov & Babenko , 2020 ; Voynov et al. , 2021 ; MelasKyriazi et al. , 2021 ) are currently able to perform binary segmentation only , and it is not clear if they can be used in the general setup of semantic segmentation . The second line of works ( Zhang et al. , 2021 ; Tritrong et al. , 2021 ; Xu & Zheng , 2021 ) is more relevant to our study since they are based on the intermediate representations obtained in GANs . In particular , the method proposed in ( Zhang et al. , 2021 ) trains a pixel class prediction model on these representations and confirms their label efficiency . In the experimental section , we compare the method from ( Zhang et al. , 2021 ) to our DDPM-based one and demonstrate several distinctive advantages of our solution . Representations from generative models for discriminative tasks . The usage of generative models as representation learners has been widely investigated for global prediction ( Donahue & Simonyan , 2019 ; Chen et al. , 2020a ) , and dense prediction problems ( Zhang et al. , 2021 ; Tritrong et al. , 2021 ; Xu & Zheng , 2021 ; Xu et al. , 2021 ) . While previous works highlighted the practical advantages of these representations , such as out-of-distribution robustness ( Li et al. , 2021a ) , generative models as representation learners receive less attention compared to alternative unsupervised methods , e.g. , based on contrastive learning ( Chen et al. , 2020b ) . The main reason is probably the difficulty of training a high-quality generative model on a complex , diverse dataset . However , given the recent success of DDPM on Imagenet ( Deng et al. , 2009 ) , one can expect that this direction will attract more attention in the future . 3 REPRESENTATIONS FROM DIFFUSION MODELS . In the following section , we investigate the image representations learned by diffusion models . First , we provide a brief overview of the DDPM framework . Then , we describe how to extract features from them and investigate what kind of semantic information these features might capture . Background . Diffusion models transform noise xT∼N ( 0 , I ) to the sample x0 by gradually denoising xT to less noisy samples xt . Formally , we are given a forward diffusion process : q ( xt|xt−1 ) : = N ( xt ; √ 1− βtxt−1 , βtI ) , ( 1 ) for some fixed variance schedule β1 , . . . , βt . Importantly , a noisy sample xt can be obtained directly from the data x0 : xt = √ ᾱtx0 + √ 1− ᾱt , ( 2 ) where ∼ N ( 0 , 1 ) and αt : = 1− βt , ᾱt : = ∏t s=1 αs . Pretrained DDPM approximates a reverse process : pθ ( xt−1|xt ) : = N ( xt ; µθ ( xt , t ) , Σθ ( xt , t ) ) . ( 3 ) In practice , rather than predicting the mean of the distribution in Equation ( 3 ) , the noise predictor network θ ( xt , t ) predicts the noise component at the step t ; the mean is then a linear combination of this noise component and xt . The covariance predictor Σθ ( xt , t ) can be either a fixed set of scalar covariances or learned as well ( the latter was shown to improve the quality of models ( Nichol & Dhariwal , 2021 ) ) . The denoising model θ ( xt , t ) is typically parameterized by different variants of the UNet architecture ( Ronneberger et al. , 2015 ) , and in our experiments we investigate the state-of-the-art one proposed in ( Dhariwal & Nichol , 2021 ) . Extracting representations . For a given real image x0 ∈ RH×W×3 , one can compute T sets of activation tensors from the noise predictor network θ ( xt , t ) . Formally for t = T . . . 1 we first corrupt x0 by adding a Gaussian noise : xt = √ ᾱtx0+ √ 1− ᾱt . The noisy xt is used as an input of θ ( xt , t ) parameterized by the UNet model . The UNet ’ s intermediate activations are then upsampled to H×W with bilinear interpolation . This allows treating them as pixel-level representations of x0 . 3.1 REPRESENTATION ANALYSIS . We analyze the representations produced by the noise predictor θ ( xt , t ) for different t. We consider the state-of-the-art DDPM checkpoints trained on the LSUN-Bedroom and FFHQ datasets2 . The intermediate activations from the noise predictor capture semantic information . For this experiment , we take a few images from the LSUN-Bedroom and FFHQ datasets and manually assign each pixel to one of the 28 and 34 semantic classes , respectively . Our goal is to understand whether the pixel-level representations produced by DDPM effectively capture the information about semantics . To this end , we train a multi-layer perceptron ( MLP ) to predict the pixel semantic label from its features produced by one of the 18 UNet decoder blocks on a specific diffusion step t. Note that we consider only the decoder activations because they also aggregate the encoder ones through the skip connections . MLP is trained on 20 images and evaluated on 20 hold-out ones . The predictive performance is measured in terms of mean IoU . 2https : //github.com/openai/guided-diffusion The evolution of predictive performance across the different blocks and diffusion steps t is presented in Figure 2 . The blocks are numbered from the deep to shallow ones . Figure 2 shows that the discriminability of the features produced by the noise predictor θ ( xt , t ) varies for different blocks and diffusion steps . In particular , the features corresponding to the later steps of the reverse diffusion process typically capture semantic information more effectively . In contrast , the ones corresponding to the early steps are generally uninformative . Across different blocks , the features produced by the layers in the middle of the UNet decoder appear to be the most informative on all diffusion steps . Additionally , we split the LSUN-Bedroom semantic classes into small-sized and big-sized based on the average area in the annotated dataset . Then , we evaluate mean IoU for these classes independently across the different UNet blocks and diffusion steps . The results are presented in Figure 3 . As expected , for big-sized objects , the predictive performance starts growing earlier in the reverse process . The shallow blocks ( 10 , 12 ) are more informative for small-sized objects , while the deeper ones ( 2 , 4 ) are for big-sized ones . Figure 2 implies that for certain UNet blocks and diffusion steps , similar DDPM-based representations correspond to the pixels of the same semantics . Figure 1 shows the k-means clusters ( k=5 ) formed by the features extracted from the blocks { 6 , 8 , 10 , 12 } on the diffusion steps { 50 , 250 , 450 , 650 , 850 } , and confirms that clusters can span coherent semantic objects and objectparts . In the block B=6 , the features correspond to coarse semantic masks , e.g. , they can not discriminate between various face parts . On the other extreme , the features from B=12 do not exhibit semantic meaningness , reaffirming the predictive performance behavior from Figure 2 . The block B=8 appears to be a “ sweet spot ” demonstrating fine-grained semantic fragmentation . Across different diffusion steps , the most meaningful features correspond to the later ones . We attribute this behavior to the fact that on the earlier diffusion steps , the global structure of a DDPM sample has not emerged yet , therefore , it is hardly possible to predict segmentation masks at this stage . This intuition is qualitatively confirmed by the masks in Figure 1 . For t=850 , the masks poorly reflect the content of actual images , while for smaller values of t , the masks and images are semantically coherent .

Summary: This paper explores to what extent Denoising Diffusion Probabilistic Models (DDPMs) serve as good representational learners for transfer or semi-supervised learning on downstream tasks. They are particularly interested in the semantic segmentation as a prototypical dense computer vision task.  They find that DDPM do provide useful representations.  In particular they show that the middle U-NET layers, at intermediate steps in the reverse diffusion process produce activations from which a simple MLP can infer good segmentations results (as measured by IoU).  They demonstrate very good results in a substantial number of experiments in which a DDPM is learned from a large set of unlabeled images, and then the MLP is trained on a relatively small set of labelled in-domain images, and then evaluated. They show that a simple segmentation network trained on top of DDPM activations outperforming a wide array of baseline models, including DatasetGAN and other recent SOTA methods.