Visualizing ViT Patch Features with PCA

Google Colab

Vision Transformers (ViTs) split an image into a grid of patches and produce one feature vector per patch. These patch features (or dense features) capture what the model understands about each region of the image. But each vector has hundreds of dimensions, so we can’t look at them directly.

A simple and surprisingly effective trick, popularized by DINO, is to run PCA on the patch features and map the top three principal components to the red, green, and blue channels of an image. The result paints each patch by its dominant features: semantically similar regions get similar colors, so objects pop out from the background without the model ever being trained for segmentation. Because of this, PCA maps have become a standard qualitative check of feature quality for self-supervised and distilled backbones, and the DINOv2 and DINOv3 papers both use them to showcase their models.

In this tutorial we:

  1. Load a pretrained DINOv3 backbone with LightlyTrain and extract its patch features.

  2. Turn those features into an RGB image with PCA.

  3. Compare the feature maps of DINOv3, DINOv2, and EUPE on the same image.

  4. Meet Lightly’s distilled ViT-Tiny models and see how well their features hold up against much larger backbones.

Install Dependencies

We only need lightly-train. It already pulls in torch, numpy, and matplotlib, which is everything we use here. The PCA is computed directly with torch, so there is no need for scikit-learn.

!pip install lightly-train

Get a Sample Image

Let’s download a single image to visualize. Feel free to swap in your own.

import io
import urllib.request

from PIL import Image

url = "https://raw.githubusercontent.com/pytorch/hub/master/images/dog.jpg"
image = Image.open(io.BytesIO(urllib.request.urlopen(url).read())).convert("RGB")
image.resize((512, 512))

Preprocess the Image

Before feeding the image to the model we resize it to a multiple of the model’s patch size and normalize it with the standard ImageNet statistics that the DINO models were trained with. Resizing to a multiple of the patch size means the image divides evenly into patches. A larger size gives a finer feature grid (and a sharper visualization) at the cost of more compute.

import torch
import torchvision.transforms.v2 as transforms

# Standard ImageNet normalization used by the DINOv2/DINOv3 backbones.
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)


def preprocess(image, patch_size, size=768):
    # Round the target size down to a multiple of the patch size.
    size = (size // patch_size) * patch_size
    transform = transforms.Compose(
        [
            transforms.Resize((size, size)),
            transforms.ToImage(),
            transforms.ToDtype(torch.float32, scale=True),
            transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
        ]
    )
    return transform(image).unsqueeze(0)  # (1, 3, H, W)

Load a DINOv3 Backbone and Extract Patch Features

We load a pretrained DINOv3 backbone with LightlyTrain. get_wrapped_model accepts any model name from lightly_train.list_models() (e.g. "dinov3/vits16", "dinov2/vitb14") and downloads the pretrained weights on first use. Calling forward_features on the returned wrapper gives us the patch features as a spatial grid of shape (1, feature_dim, H, W), where H and W are the number of patches along each axis.

Note

get_wrapped_model lives under lightly_train._models and is an internal API that may change between releases. The public embed command only returns a single pooled vector per image, so we use the wrapper directly here to access the per-patch features.

from lightly_train._models.package_helpers import get_wrapped_model

# Load a pretrained DINOv3 backbone (weights download on first use).
wrapper = get_wrapped_model("dinov3/vitb16", num_input_channels=3).eval()

x = preprocess(image, patch_size=wrapper.patch_size())
with torch.no_grad():
    features = wrapper.forward_features(x)["features"]  # (1, feature_dim, H, W)

print(f"patch size:   {wrapper.patch_size()}")
print(f"feature dim:  {wrapper.feature_dim()}")
print(f"feature grid: {tuple(features.shape)}")

Reduce the Features to RGB with PCA

Now the core step. We treat every patch as a data point in feature_dim-dimensional space, center the points, and project them onto their top three principal components with torch.pca_lowrank. Those three components become the R, G, and B channels. Finally we normalize each channel to the 1st to 99th percentile range so a few outlier patches don’t wash out the colors.

import numpy as np


def pca_rgb(features):
    """Maps ViT patch features to an (H, W, 3) RGB image via PCA.

    Args:
        features: Patch features of shape (1, feature_dim, H, W).

    Returns:
        An (H, W, 3) float array with values in [0, 1].
    """
    _, feature_dim, h, w = features.shape
    # Flatten the spatial grid into a (num_patches, feature_dim) matrix.
    tokens = features[0].permute(1, 2, 0).reshape(-1, feature_dim)
    tokens = tokens - tokens.mean(dim=0, keepdim=True)

    # Project onto the top 3 principal components.
    _, _, v = torch.pca_lowrank(tokens, q=3)
    projected = (tokens @ v[:, :3]).reshape(h, w, 3).cpu().numpy()

    # Percentile-normalize each channel to [0, 1] for visualization.
    low = np.percentile(projected, 1, axis=(0, 1), keepdims=True)
    high = np.percentile(projected, 99, axis=(0, 1), keepdims=True)
    return np.clip((projected - low) / (high - low + 1e-8), 0, 1)


rgb = pca_rgb(features)
print(f"PCA image shape: {rgb.shape}")

Show the Result

The PCA image is small (one pixel per patch), so we upsample it to the input resolution with bilinear interpolation and plot it next to the original image.

import matplotlib.pyplot as plt
import torch.nn.functional as F


def upsample_to(rgb, size):
    # rgb: (H, W, 3) -> upsample to (size, size, 3).
    grid = torch.from_numpy(rgb).permute(2, 0, 1).unsqueeze(0)  # (1, 3, H, W)
    grid = F.interpolate(grid, size=size, mode="bilinear", align_corners=False)
    return grid[0].permute(1, 2, 0).numpy()


size = x.shape[-1]
fig, axes = plt.subplots(1, 2, figsize=(10, 5))
axes[0].imshow(image.resize((size, size)))
axes[0].set_title("Input")
axes[1].imshow(upsample_to(rgb, size))
axes[1].set_title("DINOv3 patch features (PCA)")
for ax in axes:
    ax.axis("off")
plt.tight_layout()
plt.show()

Why PCA Maps Are a Useful Quality Check

SSL backbones are trained without labels, so there is no validation accuracy to glance at, and proper evaluations like linear probing, k-NN, or fine-tuning on a downstream task require labeled data and GPU hours. PCA maps fill this gap as a qualitative metric, for three reasons:

  1. They probe exactly what dense downstream tasks consume. Detection, segmentation, and depth heads are trained on top of the patch features. If objects already emerge as coherent, sharply bounded regions without any task training, a small head has an easy job; if the features are noisy or entangled with the background, the head must compensate.

  2. They are label-free and instant. You can run them in seconds on images from your own domain, making them ideal as a first sanity check of a pretrained or freshly distilled checkpoint before committing to a full evaluation.

  3. They expose failure modes that global benchmarks hide. A backbone can score well on classification while its dense features are degraded: high-norm “outlier” patches show up as speckle noise, and weak localization shows up as blurry object boundaries.

Keep in mind that the maps are qualitative: they complement, but do not replace, a quantitative evaluation on your downstream task.

Compare Backbones

The exact same pipeline works for any ViT backbone that LightlyTrain ships. Let’s run it for DINOv3, DINOv2, and EUPE side by side. Each model computes its own PCA basis, so the colors are not directly comparable between panels, what to look for is how cleanly each model separates the object from the background and how much fine detail its features capture.

Run lightly_train.list_models() to see every available backbone name you can drop into this loop.

models = {
    "DINOv3": "dinov3/vitb16",
    "DINOv2": "dinov2/vitb14",
    "EUPE": "dinov3/vitb16-eupe",
}

fig, axes = plt.subplots(1, len(models) + 1, figsize=(5 * (len(models) + 1), 5))
axes[0].imshow(image.resize((size, size)))
axes[0].set_title("Input")
axes[0].axis("off")

for ax, (title, name) in zip(axes[1:], models.items()):
    wrapper = get_wrapped_model(name, num_input_channels=3).eval()
    x = preprocess(image, patch_size=wrapper.patch_size())
    with torch.no_grad():
        features = wrapper.forward_features(x)["features"]
    ax.imshow(upsample_to(pca_rgb(features), size))
    ax.set_title(title)
    ax.axis("off")

plt.tight_layout()
plt.show()

Smaller Models: Lightly’s Distilled ViT-Tiny

Large backbones like DINOv3 ViT-L produce excellent features, but they are often too heavy for real-time or embedded applications. For those cases Lightly offers tiny DINOv3 models, trained with LightlyTrain’s distillation method using DINOv3 ViT-L/16 as the teacher on ImageNet-1K. They ship ready to use as dinov3/vitt16 and dinov3/vitt16plus, see the DINOv3 model docs for details. You can also distill your own tiny model on your own dataset with the same method.

PCA maps are particularly handy for judging distilled models: the goal of distillation is to transfer the teacher’s representation into a smaller student, so a good student should reproduce the teacher’s feature structure, and any degradation is immediately visible as noisier maps or washed-out object boundaries, long before you run a full benchmark.

Let’s run the exact same PCA pipeline for DINOv3 ViT-B and the two tiny models. Despite having only a fraction of the parameters, the tiny models preserve the feature structure of the much larger models remarkably well: the object still separates cleanly from the background, so their features are a strong basis for downstream dense tasks.

models = {
    "DINOv3 ViT-B": "dinov3/vitb16",
    "ViT-Tiny+": "dinov3/vitt16plus",
    "ViT-Tiny": "dinov3/vitt16",
}

fig, axes = plt.subplots(1, len(models) + 1, figsize=(5 * (len(models) + 1), 5))
axes[0].imshow(image.resize((size, size)))
axes[0].set_title("Input")
axes[0].axis("off")

for ax, (title, name) in zip(axes[1:], models.items()):
    wrapper = get_wrapped_model(name, num_input_channels=3).eval()
    num_params = sum(p.numel() for p in wrapper.get_model().parameters())
    x = preprocess(image, patch_size=wrapper.patch_size())
    with torch.no_grad():
        features = wrapper.forward_features(x)["features"]
    ax.imshow(upsample_to(pca_rgb(features), size))
    ax.set_title(f"{title}: {num_params / 1e6:.0f}M params")
    ax.axis("off")

plt.tight_layout()
plt.show()

Conclusion

In this tutorial we loaded pretrained ViT backbones with LightlyTrain, extracted their per-patch features, and visualized them by projecting the features onto their top three principal components as RGB. This is a quick, label-free way to inspect what a foundation model has learned about your images: objects and coherent regions emerge as blocks of consistent color, even though the models were never trained for segmentation.

We also saw that the Lightly-distilled dinov3/vitt16 and dinov3/vitt16plus models retain most of the feature quality of their much larger teacher, making them a great starting point when you need a small and fast backbone, see the DINOv3 model docs to get started.

From here you can swap in any backbone from lightly_train.list_models(), try your own images, or fine-tune these backbones on a downstream task. See the LightlyTrain documentation for detection, segmentation, and classification workflows.