[论文解读] Harnessing Diffusion Models for Visual Perception with Meta Prompts

Qiang Wan, Zilong Huang+3 more|arXiv (Cornell University)|Dec 22, 2023

Generative Adversarial Networks and Image Synthesis被引用 5

一句话总结

该论文提出了一种新颖的方法，通过可学习的元提示（meta prompts）将预训练的文本到图像扩散模型适配于视觉感知任务——这些可学习嵌入能够激活特定任务的特征并重新排列层次化特征，以提升关注重点。该方法在 NYU Depth V2 和 KITTI 上实现了深度估计的最先进性能，并在 CityScapes 上实现了语义分割的新 SOTA。

ABSTRACT

The issue of generative pretraining for vision models has persisted as a long-standing conundrum. At present, the text-to-image (T2I) diffusion model demonstrates remarkable proficiency in generating high-definition images matching textual inputs, a feat made possible through its pre-training on large-scale image-text pairs. This leads to a natural inquiry: can diffusion models be utilized to tackle visual perception tasks? In this paper, we propose a simple yet effective scheme to harness a diffusion model for visual perception tasks. Our key insight is to introduce learnable embeddings (meta prompts) to the pre-trained diffusion models to extract proper features for perception. The effect of meta prompts are two-fold. First, as a direct replacement of the text embeddings in the T2I models, it can activate task-relevant features during feature extraction. Second, it will be used to re-arrange the extracted features to ensures that the model focuses on the most pertinent features for the task on hand. Additionally, we design a recurrent refinement training strategy that fully leverages the property of diffusion models, thereby yielding stronger visual features. Extensive experiments across various benchmarks validate the effectiveness of our approach. Our approach achieves new performance records in depth estimation tasks on NYU depth V2 and KITTI, and in semantic segmentation task on CityScapes. Concurrently, the proposed method attains results comparable to the current state-of-the-art in semantic segmentation on ADE20K and pose estimation on COCO datasets, further exemplifying its robustness and versatility.

研究动机与目标

• 为解决在不依赖外部字幕模型或数据集特定标签的情况下，将预训练的文本到图像扩散模型适配于视觉感知任务的挑战。
• 开发一种轻量级、端到端可训练的提示机制，能够动态激活并重组特征以实现特定任务的感知。
• 通过利用扩散模型固有的去噪过程，采用循环细化训练策略来提升特征质量和模型适应性。
• 以一种高效且可泛化的方式实现有效特征提取，适用于深度估计和语义分割等多样化感知任务。

提出的方法

• 引入一种可学习的嵌入，称为“元提示”，替代扩散模型中的文本嵌入，作为任务特定的条件输入，以激活相关特征。
• 在多尺度上使用元提示与特征图之间的点积操作，重新排列并强调当前任务最相关的特征。
• 实施一种循环细化训练策略，其中相同的 UNet 参数在多个去噪步骤中重复使用，并通过调制的时间步嵌入来适应特征分布的变化。
• 应用共享的 VQVAE 编码器将输入图像压缩到潜在空间，随后使用共享的 UNet 和任务特定的元提示进行迭代去噪。
• 使用标准的下游任务损失函数（如深度估计的 L1 损失、分割任务的交叉熵损失）对整个系统进行端到端训练。
• 使用双线性插值上采样最终的注意力加权特征，并生成可视化结果（热力图），以解释每个元提示的关注区域。

实验结果

研究问题

• RQ1 在不微调主干网络的情况下，预训练的扩散模型能否被有效重用于视觉感知任务的特征提取？
• RQ2 如何设计元提示，使其能够自适应地激活并重组多样化感知任务（如深度估计和语义分割）的特征？
• RQ3 采用共享 UNet 参数的循环细化策略是否能提升感知任务中的特征质量？与单次前向特征提取相比表现如何？
• RQ4 调制的时间步嵌入在多轮去噪过程中在多大程度上缓解了分布偏移？这对性能有何影响？

主要发现

• 所提方法在深度估计任务中实现了新的 SOTA，刷新了 NYU Depth V2 和 KITTI 基准的记录。
• 在 CityScapes 上实现了语义分割的新 SOTA，表明其在不同数据集间具有强大的泛化能力。
• 在 ADE20K 上，该方法在语义分割任务中达到了与当前 SOTA 方法相当的性能，显示出其鲁棒性与通用性。
• 该方法在 COCO 上的姿态估计任务中也表现出竞争力，进一步验证了其在多样化感知任务中的适应能力。
• 随着细化步骤的增加，性能持续提升，证实了循环去噪策略在增强特征质量方面的有效性。
• 最优元提示数量分别为：深度估计任务 50 个，语义分割任务 150 个，表明提示效率具有任务依赖性。