Unsupervised Keypoints from Pretrained Diffusion Models

TL;DR

This work tackles unsupervised keypoint discovery by leveraging pre-trained text-to-image diffusion models. It optimizes text embeddings so cross-attention maps in Stable Diffusion become localized Gaussians, yielding dataset-wide keypoints without labels. The approach achieves strong results across CelebA, CUB-200-2011, Tai-Chi-HD, DeepFashion, and Human3.6m, notably excelling on unaligned, in-the-wild data and sometimes surpassing supervised baselines, while demonstrating cross-domain generalization. The method offers a practical, scalable route to robust landmarks with public code available, enabling broader deployment in pose estimation and related tasks.

Abstract

Unsupervised learning of keypoints and landmarks has seen significant progress with the help of modern neural network architectures, but performance is yet to match the supervised counterpart, making their practicability questionable. We leverage the emergent knowledge within text-to-image diffusion models, towards more robust unsupervised keypoints. Our core idea is to find text embeddings that would cause the generative model to consistently attend to compact regions in images (i.e. keypoints). To do so, we simply optimize the text embedding such that the cross-attention maps within the denoising network are localized as Gaussians with small standard deviations. We validate our performance on multiple datasets: the CelebA, CUB-200-2011, Tai-Chi-HD, DeepFashion, and Human3.6m datasets. We achieve significantly improved accuracy, sometimes even outperforming supervised ones, particularly for data that is non-aligned and less curated. Our code is publicly available and can be found through our project page: https://ubc-vision.github.io/StableKeypoints/

Figures (4)

• Figure 1: Example attention maps -- we show example attention maps for a selected learned keypoint for the CUB-200-2011 dataset, on the CUB-aligned subset. As shown, our keypoint attention map responds consistently across varying images.
• Figure 2: Overview -- we pass a randomly initialized text embedding into Stable Diffusion stableDiffusion and extract the attention maps. We then optimize the text embedding to have localized attention maps, by supervising them to become a single-mode Gaussian distribution, drawn at the location of their maxima. We also enforce attention maps to be transformation equivariant to small affine transformations on images. We repeat this process over a set of training images, which after optimization provides a set of $K$ keypoints.
• Figure 3: Qualitative examples of unsupervised keypoints -- we show our learned keypoints for the CelebA, CUB-200-2011, Tai-Chi-HD, DeepFashion, and Human 3.6M datasets (both for cropped and masked as well as our relaxed version). Note how our keypoints are consistent despite the variability. Our method significantly outperforms other baselines, especially for the challenging Tai-Chi-HD dataset and the CUB subsets.
• Figure 4: Generalization -- we apply our learned text tokens (keypoints) to images from other datasets, including those that are of completely different domains. Our tokens generalize well for data of similar type, and surprisingly well even for some extreme cases.