Emergent Correspondence from Image Diffusion

In plain words

Matching points between two photos, such as the same cat's ear in two pictures, is easy for people and hard to teach a program without labelled examples. This paper observes that image generators already contain this knowledge. A diffusion model learns to clean up noisy images, and while it does, its inner layers organise pixels by what they mean.

The method is simple. Take a real photo, blur it with a chosen amount of noise, and pass it through the generator's denoising network once. Keep the values from one of the inner layers and treat them as a description of each pixel. Run this on two photos, and for a point in the first, pick the pixel in the second whose description looks most alike. Repeating the noise and averaging the results makes the descriptions steadier.

Nothing is trained for this task. The choice of noise level and inner layer matters and is tuned per task. The authors report that these generic features beat self-supervised image features on matching tasks, and sometimes match methods trained with labels.

TL;DR

• DIFT reads features out of a pretrained diffusion U-Net, with no fine-tuning and no correspondence labels §1 §4.
• A real image is first noised to a level t, then passed once through the network. Features come from one upsampling block, averaged over several noise draws §4.2 app. C.
• Matching is nearest neighbour by cosine distance §3.
• Author-reported: on SPair-71k, DIFT from Stable Diffusion gets 52.9 PCK per image, against 33.9 for DINO and 38.4 for OpenCLIP tab. 1.
• Several settings are not in the paper or differ from the code. Read Paper vs code and Pitfalls before implementing.

Problem

Dense correspondence between two images is needed for reconstruction, tracking, segmentation and editing §1. Supervised methods need annotated pairs, which are costly to collect §2. Self-supervised contrastive features give useful correspondences, but the training objective is image-level app. B. The question is whether diffusion models, trained to denoise, learn correspondence too §1.

Two obstacles exist. The network is trained on noisy inputs, and a real image is clean, so the authors add noise before feeding it in §4.2. Also, the intermediate noisy images that the generator sees when producing an image are not available for a real photo §4.2.

Method

Correspondence as nearest neighbour §3 eq. 1. Given feature maps F_1 and F_2 and a pixel p_1, the match is the pixel p_2 whose feature has the smallest distance to F_1(p_1). Features are read between pixels by bilinear interpolation. The distance is cosine by default §3.

p_2 = \arg\min_{p} \; d\big(F_1(p_1),\, F_2(p)\big) eq. 1

Diffusion forward process §4.1 eq. 2. A clean image x_0 is mixed with Gaussian noise to give x_t, with a fixed schedule \alpha_t. A network f_\theta takes x_t and t and predicts the noise. Denoising runs the forward steps in reverse §4.1.

x_t = \sqrt{\alpha_t}\, x_0 + \sqrt{1-\alpha_t}\, \epsilon, \qquad \epsilon \sim \mathcal{N}(0, \mathbf{I}) eq. 2

Generated images §4.2 fig. 2. For an image the model generated, the authors have every intermediate state. They take feature maps at one time step and find correspondences between generated images. This works across categories and across domains such as photo to painting fig. 2.

Real images §4.2. Procedure, in order:

1. Add noise to the real image to reach level t, using the forward process of eq. 2 §4.2.
2. Pass the noised image and t through the network once.
3. Keep the output of one upsampling block, the block index n app. C.
4. Repeat steps 1 to 3 with several noise draws, and average the feature maps §4.2 app. C.

The result is DIFT. Both the time step and the block index are chosen per task by grid search §4.2 app. C. Larger t and earlier blocks give more semantic features. Smaller t and later blocks keep more low-level detail §4.2.

# Extract DIFT for one real image  [§4.2] [app. C]
# Input:  image (3, H, W), scaled to [-1, 1]; timestep t; block index n
# Output: features (C, h, w)

def extract_dift(image, t, n, copies):
    # Stack copies of the image so each gets its own noise  [app. C]
    batch = repeat(image, copies)                      # (K, 3, H, W)
    # Noise every copy to level t with the forward process  [eq. 2] [§4.2]
    noisy = add_noise(batch, t)                        # (K, 3, H, W)
    # One denoiser pass; keep the output of upsampling block n  [app. C]
    activations = unet_up_blocks(noisy, t)[n]          # (K, C, h, w)
    # Average over the copies to get one feature map  [§4.2]
    return mean(activations, axis=0)                   # (C, h, w)

# Match one source pixel to the target image  [§3] [eq. 1]
# Input:  source_feat (C, h1, w1), target_feat (C, h2, w2), point (x, y)
# Output: matched target pixel (x', y')

def match_point(source_feat, target_feat, point):
    # Feature at the source point, bilinearly interpolated  [§3]
    query = sample_bilinear(source_feat, point)        # (C,)
    # Cosine similarity with every target location  [§3]
    scores = cosine(target_feat, query)                # (h2 * w2,)
    # Nearest neighbour is the best-scoring location  [eq. 1]
    return pixel_of(argmax(scores), target_feat)

Two backbones §5.1. DIFT_{sd} uses Stable Diffusion 2-1, trained on LAION. DIFT_{adm} uses Ablated Diffusion Model, trained on ImageNet without labels §5.1. The baselines are DINO (ViT-B/8, ImageNet without labels) and OpenCLIP (ViT-H/14, LAION), used without fine-tuning §5.1. Each baseline is also grid-searched over block index per task app. C.

Prompts app. C. For Stable Diffusion, SPair-71k uses the text "a photo of a [class]" with the class from the dataset. CUB uses a single bird prompt. HPatches and the video tasks use an empty prompt app. C. The authors report that an empty prompt on SPair-71k costs very little.

Hyperparameters

Values as stated in the paper. Feature map size is the stride relative to the input. Dimension is the channel count of the chosen block.

Results

Author-reported. Setups: SPair-71k test split, 12,234 pairs over 18 categories §5.2; PF-WILLOW, 900 pairs §5.2; CUB, 14 splits of 25 images §5.2; HPatches, 116 sequences §6.1; DAVIS-2017 validation videos and JHMDB §6.2. Baselines in the same tables are weakly supervised, supervised, or self-supervised. The paper labels each row by group.

SPair-71k, PCK per image, \alpha_{bbox}=0.1 tab. 1.

SPair-71k, PCK per point, \alpha_{bbox}=0.1 tab. 2.

The authors also report that DIFT_{sd} is above the best supervised method in 9 of 18 SPair-71k categories §5.2 tab. 1. Our count from tab. 1 agrees (aeroplane, bicycle, bird, cat, chair, horse, motorbike, plant and sheep). The overall per-image score is 2.4 points below SCorrSAN tab. 1, which the introduction describes as remaining on par §1.

PF-WILLOW and CUB tab. 3.

HPatches homography accuracy, all sequences tab. 4.

DAVIS-2017 and JHMDB tab. 5.

PF-PASCAL, PCK per image tab. 6. The paper leaves this out of the main results, because 302 of 506 test images also appear in the training set app. D. DIFT_{sd} scores 69.4 at α=0.05 and 84.6 at α=0.10. PWarpC scores 64.2 and 84.4.

Cost app. B. One denoiser pass per copy. On one A6000 GPU, DIFT_{sd} takes 203 ms per 768×768 image and OpenCLIP 231 ms. DIFT_{adm} takes 110 ms per 512×512 image and DINO 154 ms app. B.

Ablations. Dropping the copy averaging costs DIFT_{sd} 1.6 points of per-point PCK on SPair-71k (59.5 to 57.9), and DIFT_{adm} 0.9 points (52.0 to 51.1) app. B. Switching to an empty prompt costs DIFT_{sd} 1.9 points (59.5 to 57.6) app. C. Features from the VAE encoder of Stable Diffusion are an order of magnitude worse than the U-Net features app. B. Accuracy stays high over a wide range of t §5.2 fig. 5 fig. 11 fig. 12.

Paper vs code

The code is Tsingularity/dift at 9421eb2034396c5b66f1aff37f03e540c264e52f. The project page links it as the paper's code [project page]. Its LICENSE is MIT and names Luming Tang, the first author LICENSE.

• Stable Diffusion works in latent space. The paper describes noising the real image §4.2 eq. 2. The code first encodes the image with the VAE and samples the latent, then noises the latent dift_sd.py. ADM works in pixel space dift_adm.py.
• Noise copies are independent. The code repeats the image and draws fresh noise for each copy, then averages the block outputs over the batch dift_sd.py dift_sd.py dift_adm.py. This matches §4.2.
• ADM model size. The network is built with image_size=256 [code: src/models/dift_adm.py#L17], but the paper and README feed 512×512 inputs app. C README.md. The code does not explain this.
• "Per image" is per pair. The paper defines PCK per image as the average of per-image values §5.2. The code appends one value per image pair, then averages those eval_spair.py. The paper's table uses the code's version, so compare per-pair numbers, not per-picture numbers.
• HPatches threshold. RANSAC and LMEDS use a reprojection threshold of 3 pixels eval_homography.py. The paper does not state it.
• DAVIS context cap. For the video "shooting", the code caps previous frames at 10 to avoid running out of memory eval_davis.py. The paper's DAVIS setting is 28 frames app. C. The cap is code-only.
• DAVIS copy default. The DAVIS script defaults to 4 copies [code: eval_davis.py#L267], while the paper's default is 8 app. C. The README uses 4 for DIFT_{adm} and 8 for DIFT_{sd} [README.md#L140-L192].
• Missing evaluation code. The repository has scripts for SPair-71k, HPatches, DAVIS-2017 and homography only. JHMDB, PF-WILLOW, CUB and PF-PASCAL have no evaluation script [file listing]. Their numbers cannot be rerun from this code.
• Settings that match. The README commands reproduce the paper's values for SPair-71k (DIFT_{sd} t=261, block 1; DIFT_{adm} t=101, block 4), HPatches (DIFT_{sd} t=0, block 2; DIFT_{adm} t=41, block 11) and DAVIS (both models, with the temperature, top-k, radius and 28-frame settings) [README.md#L74-L195] app. C. The code's block indices run from 0 [code: src/models/dift_sd.py#L127].
• Scripts run on GPUs. The featurizers call .cuda() directly [code: src/models/dift_sd.py#L197] dift_adm.py. CPU use needs edits.

Pitfalls

• Tuning on the test split. The time step and block index are grid-searched on SPair-71k §5.2. SPair-71k is also the benchmark reported. The same settings are reused for PF-WILLOW and CUB §5.2. Treat the SPair numbers as tuned.
• Best setting per method. For HPatches the authors report, per method, the best of cosine or L2 and RANSAC or LMEDS app. C. The same best-of rule applies to every method in that table app. C. Expect some gain from selection.
• Noise level is the main knob. A small t keeps detail and a large t keeps semantics. The best t depends on the task, so use the table values for that task rather than a single default §4.2 app. C.
• Results change with noise. Each run draws new noise, so features vary. Use the copy average and keep the seed fixed when you need reproducible output app. B dift_sd.py.
• Memory and time. Each copy is a full U-Net pass at 768×768 for DIFT_{sd}. The paper's fallback is 4 copies app. C. It measured the cost of no averaging at all, not of 4 copies app. B.
• Grid resolution. Features are at 1/8 or 1/16 of the input. The code bilinearly upsamples the feature map to image size before matching eval_spair.py, so the match lands on an image-size grid while its features stay as coarse as the native map app. C.
• Keypoint sampling convention. HPatches features are sampled with grid_sample after normalising coordinates with (w-1)/2 [code: eval_hpatches.py#L53-L54] eval_hpatches.py. The call does not set align_corners, so check your PyTorch default before porting.
• Threshold choice in PCK. The SPair tables use the bounding-box threshold eval_spair.py. The PF-PASCAL and some baseline rows use image-size thresholds tab. 6. Do not compare the two columns directly.
• Weights are separate. The Stable Diffusion 2-1 and ADM checkpoints have their own licence terms. This paper's CC BY 4.0 licence covers the paper, not those weights. Not checked here; check before commercial use.
• Compute is fixed. Only one denoiser pass per copy is needed, so no sampling loop runs. Diffusion inversion gives similar results and is several times slower app. B.
• Failures. The authors show failures when the keypoint's semantic definition is ambiguous or the viewpoint change is extreme app. E fig. 19.
• Video is label propagation. The temporal results use a propagation scheme with a softmax temperature, a top-k filter and a spatial radius §6.2 eval_davis.py. These settings are part of the method. Using DIFT features alone does not reproduce the numbers.

Lineage

• Builds on: none stated. It reuses the Stable Diffusion 2-1 and ADM checkpoints as they are §5.1.
• Compared against: DINO and OpenCLIP, which share training data with the diffusion models §5.1 app. B; supervised SPair-71k methods (CATs, MMNet, TransforMatcher, SCorrSAN) tab. 1; weakly supervised methods (CAPS, PWarpC, ASIC, GANgealing) tab. 1 tab. 2 tab. 3 tab. 4; video self-supervised methods tab. 5.
• Superseded by: not known from content already in this wiki.