Denoising Vision Transformers: 消除ViT特征图中的伪影

引言

Vision Transformer (ViT)自问世以来在计算机视觉领域取得了巨大成功,在多个任务上超越了卷积神经网络(CNN)的表现。然而,研究人员发现ViT的特征图中存在一些不可忽视的问题 - 网格状的伪影,这些伪影严重影响了ViT在一些下游密集预测任务(如语义分割、深度估计等)中的性能。为了解决这一问题,来自南加州大学、康奈尔大学等机构的研究人员提出了Denoising Vision Transformers (DVT)方法。

DVT:消除ViT特征图中的伪影

DVT的核心思想是通过一个两阶段的去噪过程,有效地消除ViT特征图中的网格状伪影。这种方法不需要重新训练已有的预训练ViT模型,可以直接应用于任何ViT架构。

第一阶段:单图像去噪

在第一阶段,DVT采用了一种基于神经场(neural field)的单图像优化方法。其核心思想是将ViT的输出特征分解为三个部分:

1. F(x): 表示去噪后的语义特征
2. G(position): 表示跨视图共享的伪影
3. h(x, position): 建模位置和语义内容之间的相互依赖关系

通过最小化重建误差和跨视图一致性损失,DVT可以有效地分离出干净的语义特征F(x)。这种单图像优化过程为离线应用提供了无伪影的特征估计。

第二阶段:训练通用去噪器

在第二阶段,研究人员利用第一阶段得到的大量去噪前后的特征对{y, F},训练了一个轻量级的Transformer块作为去噪器网络D。该网络可以直接从原始ViT输出y预测无噪声特征F,即F = D(y)。这个通用去噪器具有良好的泛化能力,可以应用于未见过的图像。

DVT的优势和应用

DVT方法具有以下几个显著优势:

1. 无需重新训练:DVT可以直接应用于现有的预训练ViT模型,无需昂贵的重新训练过程。

2. 通用性强:DVT可以应用于各种ViT架构,包括DINO、DeiT-III、EVA02、CLIP、DINOv2等。

3. 性能提升显著:在多个下游任务中,DVT consistently提升了ViT的性能。

研究人员在多个具有代表性的数据集上评估了DVT的效果,包括:

• PASCAL VOC 2012 (语义分割)
• ADE20K (语义分割)
• NYU Depth V2 (深度估计)
• PASCAL VOC 2007+2012 (目标检测)

实验结果表明,DVT在所有这些任务上都取得了显著的性能提升。以DINOv2为例:

• 在PASCAL VOC 2012语义分割任务上,mIoU从83.60%提升到84.84%
• 在ADE20K语义分割任务上,mIoU从47.29%提升到48.66%
• 在NYU Depth V2深度估计任务上,相对误差从0.1238降低到0.1200
• 在PASCAL VOC目标检测任务上,mAP从81.4%提升到81.9%

这些结果充分证明了DVT在消除ViT特征图伪影方面的有效性,以及它在提升ViT下游任务性能方面的巨大潜力。

DVT的工作原理可视化

为了更直观地理解DVT的工作原理,研究人员提供了一系列可视化结果。以DINOv2 ViT-Base模型为例:

从左到右,各列分别表示:

1. 输入裁剪图像
2. 原始DINOv2输出
3. 原始输出的KMeans聚类结果
4. 原始输出的L2特征范数
5. 原始输出中中心patch与其他patch的相似度
6. DVT去噪后的输出
7. 去噪输出的KMeans聚类结果
8. 去噪输出的L2特征范数
9. 去噪输出中中心patch与其他patch的相似度
10. 分解出的共享伪影
11. 共享伪影的L2范数
12. 真实残差误差
13. 预测的残差项
14. 共享伪影和预测残差项的组合

从这些可视化结果中,我们可以清晰地看到DVT如何有效地消除了特征图中的网格状伪影,使得特征表示更加清晰和语义一致。

DVT的潜在应用

DVT的成功不仅仅局限于提升ViT在特定任务上的性能,它还揭示了ViT特征中潜在的一些有趣属性:

1. 对象突出性:通过对去噪后特征进行PCA分析,研究人员发现第二主成分能很好地捕捉图像中主要对象的突出性。这一发现可能对无监督目标检测和分割任务有重要启示。

2. 特征范数作为对象指示器:去噪后特征的L2范数被发现可以作为一个有效的对象指示器,这在原始ViT特征中并不明显。

3. 改进的特征聚类:去噪后的特征在KMeans聚类时表现出更好的语义一致性,这对于无监督语义分割等任务可能有重要价值。

这些发现为ViT在更广泛的计算机视觉任务中的应用开辟了新的可能性。

结论与展望

Denoising Vision Transformers (DVT)为解决ViT特征图中的伪影问题提供了一个简单而有效的解决方案。通过消除这些伪影,DVT不仅显著提升了ViT在多个下游任务中的性能,还揭示了ViT特征中一些潜在的有趣属性。

这项研究工作提醒我们,尽管ViT在许多任务上取得了令人印象深刻的性能,但其内部表示仍然存在一些问题。DVT的成功也鼓励我们重新评估ViT的设计,特别是在位置编码的使用方面。

未来的研究方向可能包括:

1. 探索DVT在更多视觉任务中的应用,如图像生成、视频理解等。
2. 研究如何将DVT的思想整合到ViT的训练过程中,从源头上减少伪影的产生。
3. 利用DVT揭示的ViT特征属性,开发新的无监督或自监督学习方法。

总的来说,DVT为提升ViT的性能和可解释性开辟了一个新的研究方向,有望推动计算机视觉领域的进一步发展。研究人员已经公开了DVT的代码和模型检查点,这将有助于更多研究者和开发者在此基础上进行进一步的探索和应用。