Deep ViT Features as Dense Visual Descriptors

In plain words

A vision transformer cuts a picture into a grid of small squares and gives each square a list of numbers, called features, that describe it. The authors ask what those per-square features contain when the model was trained without labels (a method called DINO). Their answer: the features mark object parts, not just the main subject, and the same kind of part gets similar features across related object types, such as body parts of different animals.

Because of that, they do not train anything new. They group similar squares with a standard clustering method, decide which groups belong to the object by checking an attention map, and match squares between two photos by keeping only pairs that pick each other as best match. With these steps they find the object that several photos share, split it into shared parts, and find matching points between two photos.

The appeal is simplicity: no labelled data for the task and no extra training. The authors report results that are close to supervised methods on some benchmarks and better than earlier unsupervised ones.

TL;DR

• Keys (one of four attention projections) from the last block, or from a middle block for correspondences, serve as dense patch descriptors §3 §4.
• Co-segmentation clusters the patches of all images together, keeps the clusters that a saliency map marks as foreground in most images, and refines the masks with GrabCut §4.
• Part co-segmentation clusters only the foreground patches into a chosen number of parts and smooths the labels with a dense CRF §4.
• Correspondences use log-binned keys, mutual nearest neighbours, and a clustering step that picks well spread pairs §4 eq. 5.
• Results are author-reported and use no training; the code (MIT) differs from the paper in several details listed under Paper vs code.

Problem

Earlier work on pretrained DINO ViTs used only the final layer, and mainly as a global or foreground-aware feature §2. Supervised co-segmentation and correspondence methods learn the notion of "common object" from their training data, so they struggle when images come from different classes §2. CNN features are coarse at the deepest layers and biased toward texture, which hurts localisation §2. The method assumes a pretrained DINO ViT is available and that the input images share an object (co-segmentation) or are semantically related (correspondence) §1 §4.

Method

The paper first studies the features, then uses them for three tasks. All three use only clustering, thresholding, binning and nearest-neighbour matching on fixed features.

Feature extraction.

1. Tokens and keys. Each image is split into n non-overlapping patches, linearly embedded with learned positional embeddings, and given a CLS token. L transformer blocks follow, each applying normalisation, multi-head self-attention and an MLP with residual connections §3 eq. 1. In each attention block every token gets a query, a key and a value eq. 2. The paper uses keys as the descriptor, and justifies that with a visual comparison and with ablations §3.1 fig. 4 tab. 3 tab. 4.
2. Overlapping patches. At test time the paper extracts overlapping patches (a stride smaller than the patch size) and interpolates the positional embedding to the new grid. This raises the feature resolution without retraining §4, Resolution Increase. The stride value is not stated in the paper body.
3. Saliency. For each patch the paper uses the mean CLS attention over selected heads of the last block §4 fig. 5(c) eq. 3. Which heads are selected is not stated in the paper.

Co-segmentation §4 fig. 5(a–e).

4. Extract the keys of every patch of every image from the last block and pool them into one bag of descriptors. Cluster the bag with k-means. The number of clusters is chosen with the elbow method §4.
5. For each cluster and image, average the saliency of the cluster's patches eq. 3.
6. Cluster votes: a cluster is foreground when it is salient in more than a fraction p of the images, using a threshold tau on the saliency average §4 eq. 4. The paper leaves tau and p as symbols.
7. Refine each binary mask with GrabCut §4.

Part co-segmentation §4 fig. 5(f–h).

8. Repeat the clustering on the foreground descriptors only, with the requested number of parts k §4.
9. Smooth the part labels with a multi-label CRF over the image §4.
10. For images with a small set, apply random crop and flip augmentations to improve clustering stability §4.

Point correspondences §4 eq. 5.

11. Compute log-binned descriptors from a middle block. Binning pools each patch with neighbours at two scales §4 tab. 4. The bin layout is shown only in code (see Paper vs code).
12. Compute cosine similarity between all patches of image A and image B. Keep a pair (i, j) when j is the best match of i and i is the best match of j (a best-buddies pair) eq. 5.
13. Choose a fixed number of well spread pairs. The paper describes the matching but not how the final pairs are chosen; see Paper vs code.

Pseudocode. These blocks follow the code's behaviour where the paper is silent, and say so in comments. The keys used here are per image, so the blocks do not depend on a shared layout.

# Dense keys of one image from one block, CLS removed  [§3] [§4]
# Input:  image (3, H, W); layer index
# Output: keys (n, D), n patches, D = heads * head_dim
def extract_keys(image, layer):
    # Overlapping 8 px patches with stride 4 (code default)  [§4]
    # [code: extractor.py#L27]
    tokens = embed_patches(image, patch=8, stride=4)
    # Keys are projected from the block's attention input  [eq. 2]
    # [code: extractor.py#L178-L182]
    keys = key_projection(block[layer], tokens)    # (1 + n, D)
    # Drop the CLS token  [code: extractor.py#L294-L295]
    return keys[1:]                                # (n, D)

# Co-segmentation of N images  [§4]  [fig. 5]
# Input:  images (N); keys[i] (n, D); saliency[i] (n,) in [0, 1]
# Output: N binary masks
def cosegment(images, keys, saliency):
    # Pool all patches of all images into one bag  [§4]
    bag = concatenate(keys)                        # (N * n, D)
    # L2-normalise, then cluster; k from the elbow rule  [§4]
    # [code: cosegmentation.py#L124-L143]
    labels, k = kmeans_with_elbow(normalise(bag))
    # Vote per image: mean saliency of each cluster above tau
    # [code: cosegmentation.py#L157-L162] (tau = 0.065)
    votes = zeros(k)
    for i in range(N):
        for c in range(k):
            score = mean(saliency[i][labels[i] == c])
            votes[c] += score > tau
    # Foreground if at least ceil(p * N) votes, p = 75%  [§4]
    # [code: cosegmentation.py#L163]
    foreground = votes >= ceil(p * N)
    masks = []
    for i in range(N):
        # Refine the patch mask with GrabCut  [§4]
        # [code: cosegmentation.py#L170-L186]
        patch_mask = foreground[labels[i]]         # (n,) boolean
        masks.append(grabcut(images[i], patch_mask))
    return masks

# Part co-segmentation  [§4]  [fig. 5]
# Input:  foreground keys of all images; num_parts = k
# Output: per image, a part label per pixel (label k = background)
def part_cosegment(images, foreground_keys, num_parts):
    # Cluster only the foreground patches into k parts  [§4]
    # [code: part_cosegmentation.py#L162-L175]
    centroids = kmeans(normalise(foreground_keys), k=num_parts)
    parts = []
    for image in images:
        # Squared distance to each part centroid, plus the distance
        # to the nearest background centroid as an extra column
        # [code: part_cosegmentation.py#L190-L204]
        energies = patch_distances(image, centroids)   # (n, k + 1)
        # Upsample to pixels, then smooth with a dense CRF  [§4]
        # [code: part_cosegmentation.py#L205-L216]
        smoothed = dense_crf(upsample(energies), image)
        parts.append(argmax(smoothed, axis=label))
    return parts

# Point correspondences between two images  [§4]  [eq. 5]
# Input:  image_a, image_b
# Output: up to 10 matched point pairs
def correspondences(image_a, image_b):
    # Log-binned keys of block 9 (0-based) for both images  [§4] [tab. 4]
    # [code: correspondences.py#L35-L41]
    desc_a = log_binned(keys(image_a, layer=9))    # (n_a, 17 * D)
    desc_b = log_binned(keys(image_b, layer=9))    # (n_b, 17 * D)
    # Cosine similarity between every patch of A and every patch of B
    sim = cosine_matrix(desc_a, desc_b)           # (n_a, n_b)
    # Best match in each direction  [eq. 5]
    a_to_b = argmax(sim, axis=1)
    b_to_a = argmax(sim, axis=0)
    # Keep best-buddies pairs  [eq. 5]
    pairs = [(i, a_to_b[i]) for i in range(n_a) if b_to_a[a_to_b[i]] == i]
    # Code only: saliency filter with threshold 0.05
    # [code: correspondences.py#L44-L66]
    pairs = saliency_filter(pairs, sal_a, sal_b, b_to_a)
    # Code only: k-means on concatenated pair descriptors, 10 clusters,
    # then keep the most salient pair of each cluster
    # [code: correspondences.py#L69-L90]
    return select_per_cluster(pairs, k=min(10, len(pairs)))

Key equations, as written in the paper:

\hat{T}^{l} = \mathsf{MSA}(\mathsf{LN}(T^{l-1})) + T^{l-1}, \quad T^{l} = \mathsf{MLP}(\mathsf{LN}(\hat{T}^{l})) + \hat{T}^{l} eq. 1

q^{l}_i = W_{q}^{l}\, t_i^{l-1}, \quad k^{l}_i = W_{k}^{l}\, t_i^{l-1}, \quad v^{l}_i = W_{v}^{l}\, t_i^{l-1} eq. 2

\mathsf{Sal}\left(S^{\mathcal{I}}_k\right) = \frac{1}{|S^{\mathcal{I}}_k|} \sum_{i \in S^{\mathcal{I}}_k} \mathsf{Attn}^{\mathcal{I}}_i eq. 3

\mathsf{Votes}(k) = \mathbb{1}\left[\sum_{\mathcal{I}} \mathsf{Sal}\left(S^{\mathcal{I}}_k\right) \ge \tau\right] eq. 4

\mathsf{BB}(M, Q) = \left\{ (m, q) \mid m \in M,\ q \in Q,\ \mathsf{NN}(m, Q) = q \wedge \mathsf{NN}(q, M) = m \right\} eq. 5

where NN is the nearest neighbour under cosine similarity §4.

Hyperparameters

The paper body gives no backbone size, patch size, stride, input size, k-means settings or hardware. The code's defaults are listed where they exist. Code defaults are not paper claims.

Results

Author-reported. Setup: DINO ViT features with no training. Co-segmentation is evaluated on MSRC, Internet300 and PASCAL-VOC (intra-class) and on PASCAL-CO, which the authors built with 40 sets of 6 images each from related classes §5.2. Part co-segmentation is evaluated on CUB, whose test set has 5K bird images §5.1. Correspondences are evaluated on 360 random Spair-71k pairs, 20 per category, with PCK at alpha 0.1 §5.3 tab. 4. Hardware is not stated. Baselines are listed in each table. Table values are copied from the paper.

Part co-segmentation, CUB tab. 1. Three key-point regression errors (lower is better) and four clustering scores (higher is better). All methods predict k = 4 parts.

Source: tab. 1. The paper states this method beats unsupervised methods by a large margin and is competitive with the supervised Choudhury et al. method §5.1.

Co-segmentation tab. 2. Jaccard index J_m and precision P_m, mean over all sets in each dataset.

Source: tab. 2. The first four rows are supervised; the others are unsupervised. The paper states the method is competitive with supervised methods and that on PASCAL-CO it beats all listed methods §5.2.

Ablation on PASCAL-CO, saliency baselines and facets tab. 3. This shows the keys are the best facet, and that the DINO ViT saliency baseline comes closest to the method.

Source: tab. 3. The paper says the ablation supports keys over other facets §5.2.

Point correspondences, Spair-71k PCK at alpha 0.1 tab. 4. Higher is better. The paper states this is mean PCK over categories.

Source: tab. 4. In the paper, the NBB and supervised columns span both rows, so the same value applies to both settings. The paper states the method beats NBB by a large margin and narrows the gap to the supervised CATs method §5.3.

Paper vs code

Code: official repository ShirAmir/dino-vit-features (the paper's project page links it; the owner is an author), commit 1779649078744e7d9be903325355fd481cca8103, MIT licence. The repository contains the extractor, co-segmentation, part co-segmentation, correspondence, PCA and similarity inspection scripts. It does not include the evaluation scripts for CUB, Spair-71k or the Jaccard metrics.

1. Voting rule. Paper eq. 4 sums the saliency of a cluster over all images and compares the sum with tau eq. 4. The code instead checks, per image, whether the cluster's mean saliency is above 0.065, then requires at least 75% of images to vote cosegmentation.py part_cosegmentation.py. The code does not implement eq. 4 as written.
2. Saliency scale. The paper's eq. 3 averages raw CLS attention. The code min-max normalises each image's saliency map to [0, 1] before averaging extractor.py. The threshold 0.065 is only meaningful on that scale.
3. Saliency resolution. The code computes co-segmentation saliency at half the input size with an 8 px stride, then resizes it up to the descriptor grid cosegmentation.py cosegmentation.py. The paper does not mention this.
4. Heads. The paper says "selected heads" without naming them. The code averages heads 0, 2, 4 and 5 of the dino_vits8 model and asserts that this model is in use extractor.py extractor.py.
5. Elbow rule. The code computes the within-cluster error for k = 1 to 14 and stops at the first k where the error falls by less than 2.5% relative to k − 1. It keeps that k, not k − 1 cosegmentation.py. The paper says only "elbow method".
6. Cluster count variable. The code sets num_labels = np.max(n_clusters) + 1 cosegmentation.py, which is k + 1 labels when labels run from 0 to k − 1. The extra label gets no votes, so the result is unaffected in practice.
7. Keys from the attention input. The code takes keys from the input to the attention module, which is the normalised tokens, not the raw residual stream extractor.py. Paper eq. 2 writes the projection of t^{l−1} without a normalisation step.
8. Correspondence saliency filter. The code drops best-buddies pairs that fail a saliency test with threshold 0.05 correspondences.py. The test keeps a pair when the first patch is salient and some salient patch of image B has that patch as its best match. It does not check that the matched patch itself is salient. The paper does not describe this filter.
9. Pair selection. The code clusters concatenated pair descriptors into 10 groups with sklearn k-means (random state 0) and keeps the most salient pair in each group correspondences.py. The paper does not describe how the final pairs are chosen.
10. Binning layout. The paper refers to a binning figure and to log-binning, and the binning figure is commented out in the source. The code builds 17 bins: 9 positions in the 3×3 neighbourhood (centre included) and 8 positions at offset 3, each averaged over a 3×3 window. Positions outside the grid are clamped to the border extractor.py.
11. Keypoint coordinates. The code maps a grid index to a pixel as index × stride + patch // 2 correspondences.py. The paper does not state this mapping.
12. Augmentation. The paper says small sets use crop and flip augmentation. The part co-segmentation script defaults to 0 augmentations, and the co-segmentation script has no augmentation option part_cosegmentation.py cosegmentation.py. The README says the notebook is set up for small sets, but the notebook was not read for this card.
13. Part CRF unary. The code uses squared Euclidean distances between L2-normalised patches and centroids as the CRF unary energies, upsampled with nearest neighbour part_cosegmentation.py. The paper names a multi-label CRF and gives no unary.
14. Three-stage option. The code has a three_stages option, off by default, that adds an extra clustering stage and an elbow with coefficient 0.94 part_cosegmentation.py part_cosegmentation.py. The paper does not describe it.
15. Layer numbering. The code indexes blocks from 0 extractor.py. The paper's Table 4 says "Layer 9" and "Layer 11", without saying whether it counts from 0 or 1.
16. Not in the repository. The Spair-71k PCK evaluation, the CUB key-point evaluation and the Jaccard evaluation are not in the code. The paper's supplementary material is not in the arXiv source that was fetched, so details that it contains could not be checked. Such details are marked "not stated" above.

Pitfalls

• Threshold units. tau = 0.065 applies to per-image min-max normalised saliency. A saliency map with a different spread will give different votes. Do not copy the threshold to a different saliency definition.
• Vote rule. Implement the per-image vote from the code, not eq. 4, unless the paper's author confirms eq. 4. The two give different clusters.
• Elbow surprise. The chosen k is the first k whose error drop is under 2.5%. Implementers may expect the last k that still improved enough, which is one lower. Test on a known set.
• Memory. With 224 px inputs, 8 px patches and stride 4 the grid is 55 × 55 = 3025 patches per image (1 + (224 − 8) // 4 = 55 per side) extractor.py. Assuming 384 key dimensions (the width of the ViT-S model, not stated in the paper), that is about 4.6 MB per image in float32. Binning multiplies the width by 17 (6528 dimensions), about 79 MB per image. Many images at once need a sampled k-means fit, as the code does cosegmentation.py.
• Silent GrabCut fallback. Any exception in GrabCut replaces that image's mask with the unrefined vote mask, with no warning cosegmentation.py.
• Augmented images. In the part co-segmentation script, augmented copies take part in clustering and the CRF but do not vote and are removed at the end part_cosegmentation.py part_cosegmentation.py.
• Sensitivity of k-means. Co-segmentation and part co-segmentation call faiss k-means without an explicit seed cosegmentation.py. Expect some run-to-run variation; correspondences use random state 0 correspondences.py.
• Figure and equation numbers. The numbers used here (fig. 1 to fig. 10, eq. 1 to eq. 5) follow the order in the LaTeX source. They were not checked against the typeset PDF.
• Supplementary material. The paper defers many implementation details to the supplement, which is not in the fetched source. Such values are marked "not stated" in this card.
• Cost of the method. The code runs the whole model for each image, even though it reads features from one block extractor.py.

Lineage

• Builds on: DINO self-distillation training of a ViT, which the paper uses as its feature source, and the ViT architecture §1. Tools cited in §4 are GrabCut, the elbow method, a multi-label dense CRF and best-buddies pairs, none linked here.
• Compared against: NBB (VGG19 features) and CATs (ResNet-101) for correspondences tab. 4; SCOPS, Choudhury et al. and DFF for part co-segmentation tab. 1; SSNM, DOCS, CycleSegNet, Li et al., TokenCut, Faktor et al. and Rubinstein et al. for co-segmentation tab. 2.
• Superseded by: not stated in the source. No later work is named in this card.