Language-driven Semantic Segmentation

1 INTRODUCTION . Semantic segmentation is a core problem in computer vision , with the aim of partitioning an image into coherent regions with their respective semantic class labels . Most existing methods for semantic segmentation assume a limited set of semantic class labels that can potentially be assigned to a pixel . The number of class labels is dictated by the training dataset and typically ranges from tens ( Everingham et al. , 2015 ) to hundreds ( Zhou et al. , 2019 ; Mottaghi et al. , 2014 ) of distinct categories . As the English language defines several hundred thousand nouns ( Li et al. , 2020c ) , it is likely that the limited size of the label set severely hinders the potential recognition performance of existing semantic segmentation models . The main reason for the restricted label sets in existing methods is the cost of annotating images to produce sufficient training data . To create training datasets , human annotators must associate every single pixel in thousands of images with a semantic class label – a task that is extremely labor intensive and costly even with small label sets . The complexity of the annotation rises significantly as the number of labels increases since the human annotator has to be aware of the fine-grained candidate labels . Additionally , inter-annotator consistency becomes an issue when objects are present in an image that could fit multiple different descriptions or are subject to a hierarchy of labels . Zero- and few-shot semantic segmentation methods have been proposed as a potential remedy for this problem . Few-shot approaches ( Shaban et al. , 2017 ; Rakelly et al. , 2018 ; Siam et al. , 2019 ; Wang et al. , 2019 ; Zhang et al. , 2019 ; Nguyen & Todorovic , 2019 ; Liu et al. , 2020b ; Wang et al. , 2020 ; Tian et al. , 2020 ; Boudiaf et al. , 2021 ; Min et al. , 2021 ) offer ways to learn to segment novel classes based on only a few labeled images . However , these approaches still require labeled data that includes the novel classes in order to facilitate transfer . Zero-shot methods , on the other hand , commonly leverage word embeddings to discover or generate related features between seen and unseen classes ( Bucher et al. , 2019 ; Gu et al. , 2020 ) without the need for additional annotations . Existing works in this space use standard word embeddings ( Mikolov et al. , 2013 ) and focus on the image encoder . In this work , we present a simple approach to leveraging modern language models to increase the flexibility and generality of semantic segmentation models . Our work is inspired by the CLIP model for image classification ( Radford et al. , 2021 ) , which pairs high-capacity image and text encoders to produce robust zero-shot classifiers . We propose to use state-of-the-art text encoders that have been co-trained on visual data , such as CLIP , to embed labels from the training set into an embedding space and to train a visual encoder to produce per-pixel embeddings from an input image that are close to the corresponding label embeddings . Since the text encoder is trained to embed closely related concepts near one another ( for example , “ dog ” is closer to “ pet ” than to “ vehicle ” ) , we can transfer the flexibility of the text encoder to the visual recognition module while only training on the restricted label sets that are provided by existing semantic segmentation datasets . An example is shown in Figure 1 ( top row ) , where the model can successfully label pixels belonging to the class “ pet ” although the training set did not contain this label . Our approach enables the synthesis of zero-shot semantic segmentation models on the fly . That is , a user can arbitrarily expand , shrink , or reorder the label set for any image at test time . We further introduce an output module that can spatially regularize the predictions while maintaining this flexibility . We demonstrate several examples of the flexibility of our model in Figure 1 . LSeg is able to output different segmentation maps based on the provided label set . For instance , in the last row , output ( a ) recognizes the chair and identifies all non-chair objects as “ other ” since these are the only two labels provided to the model . When labels are added , as in ( b ) and ( c ) , the model is able to successfully segment other objects with the expanded label set . We conduct quantitative evaluation on a variety of zero- and few-shot semantic segmentation tasks . Our approach outperforms existing methods in zero-shot settings and is competitive across multiple few-shot benchmarks . Unlike the state-of-the-art baselines we compare to , our approach does not require additional training samples . Our experiments also show that introducing the text embeddings incurs only a negligible loss in performance when compared to standard fixed-label segmentation methods . 2 RELATED WORK . Generalized semantic segmentation . The majority of existing semantic segmentation models are restricted to a fixed label set that is defined by the labels that are present in the training dataset ( Minaee et al. , 2021 ) . Few-shot semantic segmentation methods aim to relax the restriction of a fixed label set when one or a few annotated examples of novel classes are available at test time . These approaches learn to find reliable visual correspondences between a query image that is to be labeled and labeled support images that may contain novel semantic classes ( Shaban et al. , 2017 ; Rakelly et al. , 2018 ; Siam et al. , 2019 ; Wang et al. , 2019 ; Zhang et al. , 2019 ; Nguyen & Todorovic , 2019 ; Liu et al. , 2020b ; Wang et al. , 2020 ; Tian et al. , 2020 ; Wang et al. , 2020 ; Tian et al. , 2020 ; Boudiaf et al. , 2021 ; Min et al. , 2021 ) . While this strategy can significantly enhance the generality of the resulting model , it requires the availability of at least one labeled example image with the target label set , something that is not always practical . Zero-shot semantic segmentation approaches aim to segment unseen objects without any additional samples of novel classes . Text embeddings of class labels play a central role in these works . Bucher et al . ( 2019 ) and Gu et al . ( 2020 ) propose to leverage word embeddings together with a generative model to generate visual features of unseen categories , while Xian et al . ( 2019 ) propose to project visual features into a simple word embedding space and to correlate the resulting embeddings to assign a label to a pixel . Hu et al . ( 2020 ) propose to use uncertainty-aware learning to better handle noisy labels of seen classes , while Li et al . ( 2020b ) introduce a structured learning approach to better exploit the relations between seen and unseen categories . While all of these leverage text embeddings , our paper is , to the best of our knowledge , the first to show that it is possible to synthesize zero-shot semantic segmentation models that perform on par with fixed-label and few-shot semantic segmentation methods . A variety of solutions have been proposed ( Zhang et al. , 2020b ; Liu et al. , 2020a ; Perera et al. , 2020 ; Zhou et al. , 2021 ) for open-set recognition ( Scheirer et al. , 2012 ; Geng et al. , 2020 ) . These aim to provide a binary decision about whether or not a given sample falls outside the training distribution , but do not aim to predict the labels of entirely new classes . Finally , a different line of work explores cross-domain adaptation methods for semantic segmentation by using feature alignment , self-training , and information propagation strategies ( Yang et al. , 2021 ; Wang et al. , 2021 ) . The target of these works is to enhance the transferability of models to novel visual domains , but they do not address the issue of a restricted label set . As such they are orthogonal to our work . Language-driven recognition . Language-driven recognition is an active area of research . Common tasks in this space include visual question answering ( Antol et al. , 2015 ) , image captioning ( Vinyals et al. , 2014 ) , and image-text retrieval ( Li et al. , 2020a ) . CLIP ( Radford et al. , 2021 ) demonstrated that classic recognition tasks that are not commonly associated with language can strongly benefit from language assistance . CLIP uses contrastive learning together with high-capacity language models and visual feature encoders to synthesize extremely robust models for zero-shot image classification . Recent works have extended this basic paradigm to perform flexible object detection . ViLD ( Gu et al. , 2021 ) introduces an advanced zero-shot object detection method that leverages CLIP , whereas MDETR ( Kamath et al. , 2021 ) proposes an end-to-end approach that modulates a transformer-based baseline detector with text features that are obtained from a state-of-the-art language model . Like CLIP , these works have shown that the robustness and generality of object detection models can be strongly improved by language assistance . Our work is inspired by these approaches and presents , to the best of our knowledge , the first approach to flexibly synthesize zero-shot semantic segmentation models by leveraging high-capacity language models . 3 LANGUAGE-DRIVEN SEMANTIC SEGMENTATION . Our approach , Language driven Semantic segmentation ( LSeg ) embeds text labels and image pixels into a common space , and assigns the closest label to each pixel . We illustrate the framework in Figure 2 and describe each part in detail below . Text encoder . The text encoder embeds the set of N potential labels into a continuous vector space RC , producing N vectors T1 , . . . , Tn ∈ RC as outputs ( blue vectors in Figure 2 ) . Multiple network architectures are possible , and we use the pretrained Contrastive Language–Image Pre-training ( CLIP ) throughout ( Radford et al. , 2021 ) . By design , the set of output vectors is invariant to the ordering of the input labels and allows their number , N , to vary freely . Image encoder . Similar to the text encoder , the image encoder produces an embedding vector for every input pixel ( after downsampling ) . We leverage dense prediction transformers ( DPT ) ( Ranftl et al. , 2021 ) as the underlying architecture . Assume H × W is the input image size and s is a user-defined downsampling factor ( s = 2 in our implementation ) . We define H̃ = Hs , W̃ = W s . The output is a dense embedding I ∈ RH̃×W̃×C ( green tensor in Figure 2 ) . We refer to the embedding of pixel ( i , j ) as Iij . Word-pixel correlation tensor . After the image and the labels are embedded , we correlate them by the inner product , creating a tensor of size H̃ × W̃ ×N ( orange tensor in Figure 2 ) , defined as fijk = Iij · Tk . ( 1 ) We refer to the N -dimensional vector of inner products between the embedding of pixel ( i , j ) and all N words as Fij ∈ RN , where Fij = ( fij1 , fij2 , ... , fijk ) T . During training , we encourage the image encoder to provide pixel embeddings that are close to the text embedding of the corresponding groundtruth class . Specifically , given the text embeddings Tk ∈ RC of N labels and the image embedding Iij ∈ RC of pixel i , j , we aim to maximize the dot product of the entry fijk that corresponds to the ground-truth label k = yij of pixel i , j . We achieve this by defining a pixelwise softmax objective over the whole image : H , W∑ i , j=1 softmaxyij ( Fij t ) , ( 2 ) where t is a user-defined temperature parameter that we set to t = 0.07 ( Wu et al. , 2018 ; Radford et al. , 2021 ) . During training , we minimize a per-pixel softmax with cross-entropy loss ( with temperature scaling ) as is standard in semantic segmentation1 . Spatial regularization . Due to memory constraints , the image encoder predicts pixel embeddings at lower resolution than the input image resolution . We use an additional post-processing module that spatially regularizes and upsamples the predictions to the original input resolution . During this process , we have to ensure that all operations stay equivariant with respect to the labels . In other words , there should be no interactions between the input channels , whose order is defined by the order of the words and can thus be arbitrary . We evaluate two functions that fulfill this property : a simple cascade of depthwise convolutions ( Chollet , 2017 ) followed by non-linear activations ( DepthwiseBlock ) , and another block that additionally augments the depthwise convolutions with the result of a max-pooling operation over the set of labels ( BottleneckBlock ) ( Li et al. , 2019 ) . In a final step we use bilinear interpolation to recover predictions at the original resolution . We refer to these functions as “ spatial regularization blocks ” and illustrate them in Figure 3 . 1In practice we implement this using the standard nn.CrossEntropyLoss from Pytorch . Depthwise Conv Activation Depthwise Conv Activation + Max ( a ) DepthwiseBlock ( b ) BottleneckBlock Figure 3 : Illustration of BottleneckBlock and DepthwiseBlock . Training details . We initialize the backbone of the image encoder with the official ImageNet pretrained weights from ViT ( Dosovitskiy et al. , 2021 ) or ResNet ( He et al. , 2016 ) 2 and initialize the decoder of DPT randomly . During training we freeze the text encoder and only update the weights of the image encoder . We provide the full label set that is defined by each training set to the text encoder for each image . Our model can be trained on any semantic segmentation dataset and supports flexible mixing of multiple datasets through the text encoder . Existing semantic segmentation models assign a fixed channel in the output to represent the probability of a pixel being the corresponding semantic class . In contrast , our approach can dynamically handle label sets with varying length , content , and order . This property allows synthesizing arbitrary zero-shot semantic segmentation models by simply changing the labels that are fed to the text encoder .

Summary: The paper proposes Language driven Semantic segmentation (LSeg) for semantic segmentation. Essentially, LSeg embeds text labels and image pixels into a common space, and assigns the closest label to each pixel. LSeg is flexible and can dynamically handle arbitrary label sets on the fly with varying length, content, and order. The paper demonstrates that LSeg achieves comparable performance as state of the art few-shot semantic segmentation networks on FSD-1000 even though it is used in zero-shot setting. When a fixed label set is used based on the data set (all labels in the training set), it also matches the accuracy of traditional segmentation algorithms on ADE-20k. LSeg uses text encoder from CLIP-ViT-B/32 which is frozen during training while the weights of the image decoder (DPT with a ViT-L/16) is updated to maximize the correlation between the text embedding and the image pixel embedding of the ground truth class of the pixel. Spatial regularization is applied at the end that also up samples the predictions to the original input resolution.