DVT: Denoising Vision Transformers

Abstract

We study a crucial yet often overlooked issue inherent to Vision Transformers (ViTs): feature maps of these models exhibit grid-like artifacts (“Original features” in the teaser), which hurt the performance of ViTs in downstream dense prediction tasks such as semantic segmentation, depth prediction, and object discovery. We trace this issue down to the positional embeddings at the input stage. To mitigate this, we propose a two-stage denoising approach, termed Denoising Vision Transformers (DVT). In the first stage, we separate the clean features from those contaminated by positional artifacts by enforcing cross-view feature consistency with neural fields on a per-image basis. This per-image optimization process extracts artifact-free features from raw ViT outputs, providing clean feature estimates for offline applications. In the second stage, we train a lightweight transformer block to predict clean features from raw ViT outputs, leveraging the derived estimates of the clean features as supervision. Our method, DVT, does not require re-training the existing pre-trained ViTs, and is immediately applicable to any Vision Transformer architecture. We evaluate our method on a variety of representative ViTs (DINO, DeiT-III, EVA02, CLIP, DINOv2, DINOv2-reg) and demonstrate that DVT consistently improves existing state-of-the-art general-purpose models in semantic and geometric tasks across multiple datasets (Fig. 1, right, Tabs. 2 to 4). We hope our study will encourage a re-evaluation of ViT design, especially regarding the naive use of positional embeddings. Our code and checkpoints are publicly available.

Why are these artifacts are harmful?

Despite the significant strides made by ViTs, our work reveals a crucial yet often overlooked challenge: the presence of persistent noise artifacts in ViT outputs, observable across various training algorithms. These artifacts not only compromise visual clarity but also hinder feature interpretability and disrupt semantic coherence. For example, the below examples (2nd to 4th columns) demonstrate that applying clustering algorithms directly on the raw ViT output results in noisy clusters, and the patch feature similarity is less reliable.

Jump to the method section to understand what each column means.

*Right click and open image in new tab for zoomed-in views.

---

Denoised DINOv2

Emerged Properties from denoised DINOv2

We empirically identify several properties that may be hidden in the original DINOv2 model but are revealed in our denoised version.

First, we collect features from randomly sampled images from the DAVIS dataset and compute a global PCA reduction matrix to map high-dimensional features to three-dimensional color values. This reduction matrix is then applied to the features of the denoised DINOv2 model.

For visualizations of the PCA results, refer to (h) for 3-channel visualizations and to (c), (d), and (e) for individual PCA components. These components highlight different aspects of the image, with the second component capturing the object's prominence. We thus use it as a foreground mask to separate the object from the background and visualize the foreground's PCA feature map separately in (b). This foreground disentanglement is not observed in the first three principal components of the original DINOv2 model. Following this, we examine up to 10 PCA components and find the 5th component to be the most informative, though still noisy, for the object (check out this). It is arguably unfair for us to directly compare against these noisy foreground objects. Therefore, we use the officially provided `standard_array` to compute the foreground mask for the original DINOv2 model. However, it remains unclear how this `standard_array` is generated. We also provide the foreground PCA results by relying on the 5th PCA component for reference.

Additionally, as discussed in the paper, the feature norm of the denoised features can also serve as an object indicator. We visualize these feature norms in (h), noting how they complement the PCA results. Lastly, we visualize the KMeans clustering results of the denoised features in (g) to show how the features are clustered in the feature space.

The same procedure is applied to the original DINOv2 model for comparison. We find that the top 3 PCA components of the original DINOv2 model do not capture the object's prominence as well as the denoised DINOv2 model. They mainly indicate the positional patterns. The feature norm of the original DINOv2 model also does not serve as a good object indicator. The KMeans clustering results of the original DINOv2 model are adversely affected by the noise in the features.

*All results are derived from our generalizable denoiser (to be released) applied to unseen frames during denoiser training.
*Best viewed in Chrome.

---

DINOv2-reg ViT-base

DVT denoises the DINOv2 models trained with register tokens. Appending learnable register tokens to the inputs of ViTs is a recently proposed method aimed at removing artifacts in ViTs and enhancing their performance. However, we have found that models trained with registers continue to exhibit certain artifacts, albeit to a lesser extent. Our method is capable of further reducing these artifacts.

*All results are derived from our generalizable denoiser applied to unseen frames during denoiser training.
*Best viewed in Chrome.

---

How do we remove them?

Stage-1: Per-Image Denoising

Our proposed per-image denoising approach effectively removes these artifacts through per-image optimization. The 2nd to 4th columns in the visualizations display the original outputs from the DINOv2 ViT-base model, clearly illustrating how positional artifacts negatively affect semantic coherence.

The remaining columns, from the 5th to the last, showcase outputs from our neural field-based denoising approach. At the core of our method is the model:

ViT(x) ≈ F(x) + [G(position) + h(x, position)],

where:

• F(x): Represents the denoised semantic features.

• G(position): Denotes the shared artifacts across all views.

• h(x, position): Models the interdependency between position and semantic content.

Stage-2: Learn a generalizable denoiser

After applying per-image denoising to many images, we accumulate a dataset of pairs consisting of noisy ViT outputs y and their denoised counterparts F, denoted as {y, F}, We then train a denoiser network D to predict noise-free features from raw ViT outputs, i.e., F = D(y). The denoiser network is lightweight, consisting of a single Transformer block, and it generalizes well.

---

Addressing these artifacts can fully unleash the potential of pre-trained ViTs and lead to substantial performance improvements. See Table 1 below.

Table 1: Quantitative performance of DVT. DVT improves differently pre-trained ViTs for dense prediction tasks. We report performance on semantic segmentation (VOC2012, ADE20K), depth prediction (NYUd-v2), and object detection tasks (VOC2007+2012).

Citation

@article{yang2024denoising,
  author = {Yang, Jiawei and Luo, Katie Z and Li, Jiefeng and Deng, Congyue and Guibas, Leonidas J. and Krishnan, Dilip and Weinberger, Kilian Q and Tian, Yonglong and Wang, Yue},
  title = {DVT: Denoising Vision Transformers},
  journal = {arXiv preprint arXiv:2401.02957},
  year = {2024},
}

Acknowledgement

We are grateful to many friends, including Jiageng Mao, Junjie Ye, Justin Lovelace, Varsha Kishore, and Christian Belardi, for their fruitful discussions on this work and follow-ups. Katie Luo is supported by an Nvidia Graduate Fellowship. Leonidas Guibas acknowledges the support from a Vannevar Bush Faculty Fellowship. We also acknowledge an unrestricted gift from Google in support of this project.