Repurposing Stable Diffusion Attention for Training-Free Unsupervised Interactive Segmentation

TL;DR

This paper introduces a novel, training-free unsupervised image segmentation method that leverages the self-attention mechanism of Stable Diffusion to produce high-quality, interactive segmentation results with fewer manual clicks.

Contribution

The authors propose a new approach that interprets Stable Diffusion's self-attention as a Markov transition operator, enabling unsupervised segmentation without training.

Findings

Outperforms many training-based unsupervised methods in click efficiency

Produces sharper semantic boundaries and less noisy segmentation maps

Effective in interactive point prompt segmentation tasks

Abstract

Recent progress in interactive point prompt based Image Segmentation allows to significantly reduce the manual effort to obtain high quality semantic labels. State-of-the-art unsupervised methods use self-supervised pre-trained models to obtain pseudo-labels which are used in training a prompt-based segmentation model. In this paper, we propose a novel unsupervised and training-free approach based solely on the self-attention of Stable Diffusion. We interpret the self-attention tensor as a Markov transition operator, which enables us to iteratively construct a Markov chain. Pixel-wise counting of the required number of iterations along the Markov chain to reach a relative probability threshold yields a Markov-iteration-map, which we simply call a Markov-map. Compared to the raw attention maps, we show that our proposed Markov-map has less noise, sharper semantic boundaries and more uniform values within semantically similar regions. We integrate the Markov-map in a simple yet effective truncated nearest neighbor framework to obtain interactive point prompt based segmentation. Despite being training-free, we experimentally show that our approach yields excellent results in terms of Number of Clicks (NoC), even outperforming state-of-the-art training based unsupervised methods in most of the datasets. Code is available at https://github.com/mkarmann/m2n2.