扩散模型图像理解力刷新SOTA！字节联合复旦团队提出全新「元提示」策略

公众号新闻

2024-01-17 04:01

©作者 | 好困桃子

来源 | 新智元

过去一年扩散模型风头正劲，彻底改变了文生图领域！那么，扩散模型能否处理视觉感知任务？字节跳动和复旦大学技术团队在最新研究中提出了一个简单有效的方案。

Text-to-image（T2I）扩散模型在生成高清晰度图像方面显示出了卓越的能力，这一成就得益于其在大规模图像-文本对上的预训练。这引发了一个自然的问题：扩散模型是否可以用于解决视觉感知任务？

近期，来自字节跳动和复旦大学的技术团队提出了一种简单而有效的方案：利用扩散模型处理视觉感知任务。团队的关键洞察是引入可学习的元提示（meta prompts）到预训练的扩散模型中，以提取适合特定感知任务的特征。

论文标题：

Harnessing Diffusion Models for Visual Perception with Meta Prompts

论文地址：

https://arxiv.org/abs/2312.14733

代码地址：

https://github.com/fudan-zvg/meta-prompts

技术介绍

团队将 text-to-image 扩散模型作为特征提取器应用于视觉感知任务中。

输入图像首先通过 VQVAE 编码器进行图像压缩。这一步将图像分辨率降低到原始大小的 1/8，产生 latent space 中的特征表示，即。值得注意的是，VQVAE 编码器的参数是固定的，不参与后续训练。

接下来，保持未添加噪声的被送入到 UNet 进行特征提取。为了更好地适应不同任务，UNet 同时接收调制的 timestep embeddings 和多个 meta prompts，产生与形状一致的特征。

在整个过程中，为了增强特征表达，该方法进行了步的 recurrent refinement。这使得 UNet 内不同层的特征能够更好地交互融合。在第次循环中，UNet 的参数由特定的可学习的时间调制特征调节。

最终，UNet 生成的多尺度特征输入到专门为目标视觉任务设计的解码器中。

可学习的元提示（meta prompts）设计

Stable diffusion model 采用 UNet 架构，通过交叉注意力将文本提示融入图像特征中，实现了文生图。这种整合确保了图像生成在语境和语义上的准确性。

然而，视觉感知任务的多样性超出了这一范畴，因为图像理解面临着不同的挑战，往往缺乏文本信息作为指导，使得以文本驱动的方法有时显得不切实际。

为应对这一挑战，技术团队的方法采用了更为多样的策略——不依赖外部文本提示，而是设计了一种内部的可学习元提示，称为 meta prompts，这些 meta prompts 被集成到扩散模型中，以适应感知任务。

Meta prompts 以矩阵的形式表示，其中表示 meta prompts 的数量，表示维度。具备 meta prompts 的感知扩散模型避免了对外部文本提示的需求，如数据集类别标签或图像标题，也无需预训练的文本编码器来生成最终的文本提示。

Meta prompts 可以根据目标任务和数据集进行端到端的训练，从而为去噪 UNet 建立特别定制的适应条件。这些 meta prompts 包含丰富的、适应于特定任务的语义信息。比如：

在语义分割任务中，meta prompts 有效地展示了对类别的识别能力，相同的 meta prompts 倾向于激活同一类别的特征。

在深度估计任务中，meta prompts 表现出对深度的感知能力，激活值随深度变化，使 prompts 能够集中关注一致距离的物体。

在姿态估计中，meta prompts 展现出一套不同的能力，特别是关键点的感知，这有助于人体姿态检测。

这些定性结果共同突显了技术团队提出的 meta prompts 在各种任务中对任务相关激活能力的有效性。作为文本提示的替代品，meta prompts 很好地填补了了 text-to-image 扩散模型与视觉感知任务之间的沟壑。

基于元提示的特征重组

扩散模型通过其固有的设计，在去噪 UNet 中生成多尺度特征，这些特征在接近输出层时聚焦于更细致、低级的细节信息。虽然这种低级细节对于强调纹理和细粒度的任务来说足够，但视觉感知任务通常需要理解既包括低级细节的又包括高级语义解释的内容。

因此，不仅需要生成丰富的特征，确定这些多尺度特征的哪种组合方式可以为当前任务提供最佳表征也非常重要。这就是 meta prompts 的作用所在—— 这些prompts 在训练过程中保存了与所使用数据集特定相关的上下文知识。这种上下文知识使 meta prompts 能够充当特征重组的过滤器，引导特征选取过程，从 UNet 产生的众多特征中筛选出与任务最相关的特征。

团队使用点积的方式将 UNet 的多尺度特征的丰富性与 meta prompts 的任务适应性结合起来。

考虑多尺度特征 ,其中每个。和表示特征图的高度和宽度。Meta prompts 。每个尺度上重排的特征的计算为：

最后，这些经过 meta prompts 过滤的特征随后输入到特定任务的解码器中。

基于可学习的时间调制特征的recurrent refinement

在扩散模型中，添加噪声然后多步去噪的迭代过程构成了图像生成的框架。受此机制的启发，技术团队为视觉感知任务设计了一个简单的 recurrent refinement过程——没有向输出特征中添加噪声，而是直接将 UNet 的输出特征循环输入到 UNet 中。

同时为了解决随着模型通过循环，输入特征的分布会发生变化但 UNet 的参数保持不变的不一致的问题，技术团队对于每个循环引入了可学习的独特的 timestep embeddings，以调制 UNet 的参数。这确保了网络对于不同步骤中输入特征的变化性保持适应性和响应性，优化了特征提取过程，并增强了模型在视觉识别任务中的性能。

结果显示，该方法在多个感知任务数据集上都取得了最优。