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 proposes to treat deep activations of denoising diffusion probabilistic models trained on image datasets as unsupervised pixel features. To extract features for an image, one performs a fixed number of steps of the diffusion process, passes the resulting noisy image through the U-Net denoiser, and upsamples activations from a chosen layer. These pixel representations are used as inputs to simple classifiers that perform semantic segmentation. This approach is validated on three image datasets, where it shows strong results compared to baselines in a few-shot setting.