Denoising Vision Transformers
Authors: name: Jiawei Yang; name: Katie Z Luo; name: Jiefeng Li; name: Kilian Q Weinberger; name: Yonglong Tian; name: Yue Wang
Published: 2024-01-05 18:59:52

Abstract: We delve into a nuanced but significant challenge inherent to Vision
Transformers (ViTs): feature maps of these models exhibit grid-like artifacts,
which detrimentally hurt the performance of ViTs in downstream tasks. Our
investigations trace this fundamental issue down to the positional embeddings
at the input stage. To address this, we propose a novel noise model, which is
universally applicable to all ViTs. Specifically, the noise model dissects ViT
outputs into three components: a semantics term free from noise artifacts and
two artifact-related terms that are conditioned on pixel locations. Such a
decomposition is achieved by enforcing cross-view feature consistency with
neural fields in a per-image basis. This per-image optimization process
extracts artifact-free features from raw ViT outputs, providing clean features
for offline applications. Expanding the scope of our solution to support online
functionality, we introduce a learnable denoiser to predict artifact-free
features directly from unprocessed ViT outputs, which shows remarkable
generalization capabilities to novel data without the need for per-image
optimization. Our two-stage approach, termed Denoising Vision Transformers
(DVT), does not require re-training existing pre-trained ViTs and is
immediately applicable to any Transformer-based architecture. We evaluate our
method on a variety of representative ViTs (DINO, MAE, DeiT-III, EVA02, CLIP,
DINOv2, DINOv2-reg). Extensive evaluations demonstrate that our DVT
consistently and significantly improves existing state-of-the-art
general-purpose models in semantic and geometric tasks across multiple datasets
(e.g., +3.84 mIoU). We hope our study will encourage a re-evaluation of ViT
design, especially regarding the naive use of positional embeddings.