Towards Agnostic and Holistic Universal Image Segmentation with Bit Diffusion

TL;DR

This work introduces a diffusion-based framework for universal image segmentation that is agnostic to label vocabularies and aims for holistic scene understanding. It adapts discrete diffusion to segmentation via analog bit diffusion, a location-aware palette with a 2D gray-code layout, an input-scaled noise schedule, a final $\tanh$ activation, and $x$-prediction with sigmoid loss. Empirical results on EntitySeg show that analog-bit encodings with LAP outperform RGB and one-hot encodings, with near-optimal sampling achieved around 8 denoising steps and guidance weight around 1.0, but there remains a gap to state-of-the-art mask-based models. The authors argue that combining these techniques with large-scale pretraining or promptable conditioning could yield competitive, agnostic, and holistic universal segmenters in the future, while also enabling principled ambiguity modeling.

Abstract

This paper introduces a diffusion-based framework for universal image segmentation, making agnostic segmentation possible without depending on mask-based frameworks and instead predicting the full segmentation in a holistic manner. We present several key adaptations to diffusion models, which are important in this discrete setting. Notably, we show that a location-aware palette with our 2D gray code ordering improves performance. Adding a final tanh activation function is crucial for discrete data. On optimizing diffusion parameters, the sigmoid loss weighting consistently outperforms alternatives, regardless of the prediction type used, and we settle on x-prediction. While our current model does not yet surpass leading mask-based architectures, it narrows the performance gap and introduces unique capabilities, such as principled ambiguity modeling, that these models lack. All models were trained from scratch, and we believe that combining our proposed improvements with large-scale pretraining or promptable conditioning could lead to competitive models.

Figures (10)

• Figure 1: The modifications to a base diffusion model and their performance gains, visualized along with samples from our model.
• Figure 2: The cosine noise schedule with latent diffusion samples $x_t$ for various values of $t$. The latent samples use 3 bits (up to 8 masks) to make them viewable as RGB images.
• Figure 3: Performance for the three encoding types as the number of representable classes are varied.
• Figure 4: Mean performance on the validation set as the number of timesteps is varied for different guidance weights (gw).
• Figure 5: Mean performance on the validation set as the guidance weight is varied.