We introduce ODISE: Open-vocabulary DIffusion-based panoptic SEgmentation, which unifies pre-trained text-image diffusion and discriminative models to perform open-vocabulary panoptic segmentation. Text-to-image diffusion models have the remarkable ability to generate high-quality images with diverse open-vocabulary language descriptions, suggesting that their internal representations correlate well with open-world visual concepts. Text-image discriminative models like CLIP are adept at classifying images into open-vocabulary labels. We leverage the frozen internal representations of both model types for panoptic segmentation of any category in the wild.

Our approach substantially outperforms prior state-of-the-art on open-vocabulary panoptic and semantic segmentation benchmarks. Specifically, with only COCO training data, our method achieves 23.4 PQ and 30.0 mIoU on the ADE20K dataset, with 8.3 PQ and 7.9 mIoU absolute improvement over previous approaches. Our code and models are publicly available.

Humans effortlessly recognize an unlimited number of object categories and understand fine-grained distinctions between similar items. Recent computer vision research has focused on open-vocabulary recognition, enabling systems to identify unlimited categories rather than only training-set categories. However, very few unified frameworks address both instance-level and scene-level semantic understanding simultaneously through panoptic segmentation — the task that jointly resolves "what objects are present" and "where each pixel belongs."

Current open-vocabulary approaches primarily rely on text-image discriminative models trained on internet-scale data, which excel at classifying individual proposals but struggle with spatial relationships and dense prediction tasks. We hypothesize that the lack of spatial and relational understanding in text-image discriminative models is a bottleneck for open-vocabulary panoptic segmentation. These models were designed for image-level classification rather than pixel-level understanding, limiting their effectiveness when applied directly to segmentation.

Text-to-image diffusion models trained on internet-scale data represent an alternative approach. These models condition image generation through cross-attention between text embeddings and internal visual representations, suggesting that these internal representations correlate with high-level and mid-level semantic concepts expressible through language. Clustering the diffusion model's internal features reveals semantically distinct, localized information where objects naturally group together, indicating that diffusion models encode rich spatial-semantic structure well-suited for segmentation.

We propose ODISE to leverage both large-scale text-image diffusion and discriminative models for state-of-the-art open-vocabulary panoptic segmentation. Our approach consists of three stages: extracting internal features from a frozen text-to-image diffusion model using input images and captions; generating class-agnostic panoptic masks using these features; and categorizing masks into open-vocabulary categories by associating mask diffusion features with text embeddings of object category names. Our key contributions include: the first exploration of large-scale text-to-image diffusion models for open-vocabulary segmentation; a novel pipeline that effectively combines diffusion and discriminative models; and significant performance advances across multiple benchmarks.

Panoptic segmentation encompasses both instance segmentation and semantic segmentation tasks, requiring the model to assign every pixel a semantic label while simultaneously distinguishing individual object instances. Previous works follow closed-vocabulary assumptions, recognizing only training-set categories. This limitation restricts segmentation to finite-sized vocabularies that are substantially smaller than the real-world diversity of object descriptions. Methods like Mask2Former achieve strong closed-vocabulary performance but cannot generalize to novel categories unseen during training.

Open-vocabulary segmentation prior work addresses either instance segmentation with object detection or semantic segmentation independently, but not unified panoptic segmentation. Existing approaches rely on large-scale discriminative pre-training for image classification or image-text contrastive learning. The concurrent work MaskCLIP also uses CLIP for open-vocabulary segmentation. However, discriminative models' internal representations prove suboptimal for dense segmentation tasks compared to text-to-image diffusion model representations, as they were originally trained for image-level rather than pixel-level objectives.

Generative models for segmentation has been explored using GANs and diffusion models for semantic segmentation. Previous works train generative models on small-vocabulary datasets (e.g., cats, faces, ImageNet) and use few-shot examples to classify internal representations into semantic regions. DDPMSeg demonstrates state-of-the-art label-efficient semantic segmentation by leveraging diffusion model features. However, these approaches differ fundamentally from ours: they tackle small closed-vocabulary label-efficient segmentation, while ODISE addresses open-vocabulary panoptic segmentation of many unseen categories at scale.

The approach trains models using base training categories C_train, which may differ from test categories C_test, meaning C_train is not equal to C_test. Test categories may include novel categories unseen during training. During training, binary panoptic mask annotations for each category are provided. Additionally, either category labels for each mask or image text captions are available as supervision signals. During testing, neither labels nor captions are provided; only test category names are available.

ODISE extracts features from a frozen text-to-image diffusion model using input images and captions. With these extracted features and provided training mask annotations, a trainable mask generator produces panoptic masks for all possible image categories. An open-vocabulary mask classification module, trained using training images' category labels or text captions, categorizes each predicted mask by associating diffusion features with text embeddings of training category names. After training, open-vocabulary panoptic inference employs both text-image diffusion and discriminative models in a complementary manner.

Text-to-image diffusion models generate high-quality images from input text prompts, trained with millions of internet-crawled image-text pairs. Text input is encoded into text embeddings using pre-trained text encoders like T5 or CLIP. Before diffusion network input, images undergo Gaussian noise distortion. The diffusion network learns to undo this distortion given the noisy input and paired text embedding. During inference, the model takes image-shaped pure Gaussian noise and user-provided text embeddings, progressively de-noising through multiple iterations to produce the final image.

Prevalent text-to-image diffusion models typically use UNet architecture for learning denoising. The UNet includes convolution blocks, upsampling/downsampling blocks, skip connections, and attention blocks performing cross-attention between text embeddings and UNet features. At every de-noising step, diffusion models use text to infer the de-noising direction of noisy images. Since text injection occurs via cross-attention layers, the visual features become correlated with rich semantic text descriptions. The UNet output feature maps thus provide rich, dense features for panoptic segmentation.

Our method requires only a single forward pass through the diffusion model to extract visual representations, rather than the entire multi-step generative diffusion process. Given input image-text pairs (x, s), we first sample a noisy image at timestep t as x_t, where alpha values represent pre-defined noise schedules. The caption is encoded with a pre-trained text encoder, and diffusion UNet internal features are extracted by feeding the timestep and text encoding. We extract feature maps from every three UNet blocks and resize them using FPN-style approaches to create multi-scale feature pyramids suitable for dense prediction.

The diffusion model's visual representation depends on paired caption input. This becomes problematic when extracting visual representations for images without paired captions — the common application case. Using empty text proves suboptimal. We introduce an Implicit Captioner that generates implicit text embeddings directly from images rather than relying on off-the-shelf captioning networks. The implicit captioner uses frozen CLIP image encoders to encode input images into their embedding space. A learned MLP projects the image embedding into an implicit text embedding fed into the diffusion UNet. During training, image encoder and UNet parameters remain unchanged; only the MLP parameters are fine-tuned.

The mask generator accepts the visual representation as input and outputs N class-agnostic binary masks and corresponding N mask embedding features. The architecture is flexible — any panoptic segmentation network generating mask predictions works. We instantiate it with both bounding box-based and direct segmentation mask-based methods. Using box-based approaches, ROI-Aligned features of predicted mask regions compute mask embedding features. With mask-based methods, masked pooling on final feature maps computes mask embedding features.

Since our approach emphasizes dense pixel-wise predictions, we select a direct segmentation-based architecture using Mask2Former as our primary mask generator. Following established practices, predicted class-agnostic binary masks receive supervision through pixel-wise binary cross-entropy loss with corresponding ground-truth masks treated as class-agnostic ones. This loss ensures the mask generator learns to produce high-quality region proposals regardless of semantic category, which is essential for open-vocabulary generalization.

Each predicted binary mask receives a category label from open vocabularies using text-image discriminative models. These models, trained on internet-scale image-text pairs, demonstrate strong open-vocabulary classification capabilities, consisting of image encoders and text encoders. We employ two commonly used supervision signals for predicting mask category labels: category label supervision and image caption supervision.

Category Label Supervision: Training assumes access to each mask's ground-truth category label, similar to traditional closed-vocabulary training. With K_train training categories, each mask embedding feature has a corresponding ground-truth category. All training category names are encoded with frozen text encoders to form category embeddings. We compute mask embedding feature probabilities belonging to training classes through classification loss using softmax-normalized dot products with category embeddings and a learnable temperature parameter.

Image Caption Supervision: This setting assumes no category labels for each annotated mask, but access to natural language captions for each image. We extract nouns from captions and treat them as grounding category labels. A grounding loss supervises mask category label predictions: given image-caption pairs with K_word nouns extracted from captions, we compute similarity between each image-caption pair through masked attention mechanisms, encouraging nouns to be grounded by one or few masked image regions while avoiding penalizing ungrounded regions. The grounding loss follows image-text contrastive loss patterns with a temperature parameter.

During inference, test category names are available, and no captions or labels exist for test images. Images pass through the implicit captioner for implicit captions, then through the diffusion model to obtain UNet features, and finally through mask generators to predict all possible binary masks. To classify each predicted mask into test categories, probabilities are computed using the same approach as training, now with test category embeddings replacing training category embeddings.

We found that diffusion model internal representations generate many plausible masks but benefit from combining classification with text-image discriminative models, especially for open vocabularies. We leverage the discriminative model image encoder to further classify each predicted masked region. For each predicted mask, features within the mask region are pooled from the encoder's output using mask pooling operations. The final classification uses geometric mean fusion of category predictions from diffusion and discriminative models with a fixed balancing factor lambda between 0 and 1. This pooling approach is more efficient and equally effective compared to cropping each predicted mask's bounding box and separately encoding with the image encoder.

Architecture: We use Stable Diffusion pretrained on the LAION dataset as the text-to-image diffusion model. Feature maps are extracted from every three UNet blocks, resized using FPN-style approaches to create feature pyramids. Diffusion timesteps are set to t=0 by default. CLIP serves as the text-image discriminative model. Mask2Former provides the mask generator architecture, generating N=100 binary mask predictions. The entire model contains 28.1M trainable parameters (1.8% of full model) and 1,493.8M frozen parameters.

Training Details: ODISE trains for 90k iterations with 1024x1024 images using large-scale jittering with random scales between 0.1 and 2.0. Batch size is 64. AdamW optimizer with learning rate 0.0001 and weight decay 0.05 is used. COCO dataset provides training data, with panoptic mask annotations supervising binary mask loss. For image caption training, one caption per image is randomly selected from COCO caption annotations. Evaluation covers ADE20K for open-vocabulary panoptic, instance, and semantic segmentation; and Pascal datasets for semantic segmentation. A single checkpoint handles all tasks and datasets.

Open-Vocabulary Panoptic Segmentation: Training on COCO and testing on ADE20K, ODISE outperforms concurrent work MaskCLIP by 8.3 PQ and achieves 8.4 gains in the mAP metric. Specifically, ODISE achieves 23.4 PQ and 30.0 mIoU on ADE20K. Open-Vocabulary Semantic Segmentation: Evaluations across five datasets show ODISE outperforms existing state-of-the-art by 7.6 mIoU on A-150, 4.7 mIoU on A-847, 4.8 mIoU on PC-459 with caption supervision, and 6.2 mIoU on A-150, 4.5 mIoU on PC-459 with category label supervision. These consistent improvements across diverse benchmarks demonstrate the generalizability and robustness of the diffusion-based feature representations.

Visual Representations: We compare text-to-image diffusion model internal representations against other state-of-the-art pre-trained discriminative and generative models. All experiments freeze pre-trained model weights using identical training hyperparameters and mask generators. ODISE outperforms all compared models in PQ on both datasets. Despite both the diffusion model and CLIP being trained on equal-sized LAION datasets, the diffusion-based method outperforms CLIP(H) by a large margin on all metrics, demonstrating that diffusion model internal representations are inherently superior for open-vocabulary segmentation tasks.

Captioning Generators: We compare the implicit captioning module against baselines: providing empty strings (fixed text embeddings for all images); using two off-the-shelf image captioning networks (a heuristic approach and the BLIP model); and our proposed implicit captioning module. Results show that explicit and implicit captions both outperform empty text. The implicit captioning module generalizes best among all variants since it derives captions from the internet-scale text-image discriminative model's embedding space, avoiding the domain gap introduced by explicit captioning networks.

Diffusion Time Steps: Studies determine which diffusion timesteps best extract features. Results indicate all metrics decrease as timestep values increase, with t=0 producing optimal results. Training with learnable timesteps shows convergence to values near zero, independently validating t=0 as optimal. Mask Classifiers: Fusing class predictions from diffusion and discriminative models via geometric mean improves performance. Individual comparisons show the diffusion approach outperforms discriminative-only approaches, while fusion produces the highest values. Even without fusion, the diffusion-only method surpasses all existing approaches.

This work takes the first step in leveraging the frozen internal representations of large-scale text-to-image diffusion models for downstream recognition tasks. ODISE demonstrates that text-to-image generation models hold significant potential in open-vocabulary segmentation and establishes new state-of-the-art performance across multiple benchmarks. Our work reveals that text-to-image diffusion models learn "rich semantic representations" that go beyond image generation capabilities, opening new directions for leveraging text-to-image model internal representations in future vision tasks.

Looking forward, the synergy between generative and discriminative models for visual understanding represents a promising paradigm. As text-to-image diffusion models continue to improve in generation quality and diversity, their internal representations will likely encode even richer semantic structures. Future work may explore fine-tuning diffusion models for specific recognition tasks, extending to video segmentation, or leveraging newer diffusion architectures beyond UNet. The principle that generative pre-training produces representations useful for discriminative tasks may prove broadly applicable across computer vision.