扩散模型再发力！图灵奖得主Hinton团队提出：图像和视频全景分割新框架

来源：机器之心

作者：Ting Chen 等

编辑：赵阳

本文的创新点一方面在于能够在大型全景掩码上训练扩散模型，一方面在于可以同时适配图片和流媒体场景。

全景分割是一项基本的视觉任务，该任务旨在为图像的每个像素指定语义标签和实例标签。语义标签描述每个像素的类别（例如天空、竖直物体等），实例标签为图像中的每个实例提供唯一的 ID（以区分同一类别的不同实例）。该任务结合了语义分割和实例分割，提供了关于场景的丰富语义信息。

虽然语义标签的类别是先验固定的，但分配给图像中对象的实例 ID 是可以互相交换同时不影响识别的。例如，交换两辆车的实例 ID 不会影响结果。因此，经过训练以预测实例 ID 的神经网络应该能够学习从单个图像到多个实例 ID 分配的一对多映射。一对多映射的学习具有挑战性，传统方法通常利用多个阶段的管道，包括对象检测、分割、合并多个预测。最近，基于可微二分图匹配，一些学者提出了端到端方法，能够有效地将一对多映射转换为基于识别匹配的一对一映射。然而，这些方法仍然需要定制的架构和专门的损失函数，以及用于全景分割任务的内置归纳偏置。

最近的通用视觉模型，例如 Pix2Seq、OFA、UViM 和 Unified I/O，提倡通用的、不限制任务的框架来实现泛化任务，同时还能比以前的模型简单得多。例如，Pix2Seq 根据图像生成一系列具有语义意义的序列，来完成一些核心的视觉任务，并且这些模型基于 Transformers 来训练自回归模型。

在一篇新的论文中，谷歌大脑的 Ting Chen、Geoffrey Hinton 等研究者遵循相同的理念，从条件离散数据生成的角度理解全景分割任务问题。

A Generalist Framework for Panoptic Segmentation of Images and Videos

论文链接 https://arxiv.org/abs/2210.06366

如图 1 所示，研究者为全景掩码设计了一个生成模型，并为输入到模型的每一张图片都生成一组离散 token。使用者只需将过去帧的预测作为附加条件信号，就可以将此模型应用于视频数据（在线数据 / 流媒体）。这样一来，模型就可以自动地学习跟踪和分割对象。

全景分割的生成式建模非常具有挑战性，因为全景掩码是离散的，或者说是有类别的，并且模型可能非常大。例如，要生成 512×1024 的全景掩码，模型必须生成超过 1M 的离散标记（语义标签和实例标签）。这对于自回归模型来说开销还是比较昂贵的，因为 token 本质上是顺序的，很难随着输入数据的规模变化而变化。扩散模型更擅长处理高维数据，但它们最常应用于连续域而不是离散域。通过用模拟位表示离散数据，本文作者表明可以直接在大型全景掩码上训练扩散模型，而无需学习潜在空间。

通过广泛的实验，研究者们证明了他们的通用法可以在类似环境中与最先进的专家方法一较高下。

模型架构

扩散模型采样是迭代进行的，因此在推理过程中必须多次运行网络的前向传播。因此，如图 2 所示，研究者有意将网络分成两个组件：1）图像编码器；2) 掩码解码器。前者将原始像素数据映射到高级表示向量，然后掩码解码器迭代地读出全景掩码。

像素 / 图像编码器

编码器是将原始图像映射到中的特征图的网络，其中 H’和 w’是全景遮罩的高度和宽度。全景遮罩可以与原始图像的大小相同或小一点。在这项工作中，研究者们使用 ResNet 作为骨干网络，然后使用 Transformer 的编码器层作为特征提取器。为了确保输出的特征图具有足够的分辨率，并包含不同尺度的特征，受 U-Net 和特征金字塔网络的启发，研究者使用具有双边连接和上采样操作的卷积从不同的分辨率来合并特征。虽然可以使用更复杂的编码器，这样可以使用一些架构设计方面的最新进展，但这不是网络模型的主要关注点，所以研究者只是使用较为简单的特征提取器来说明其在模型中的作用。

掩码解码器

解码器在模型推理过程中基于图像特征，迭代地细化全景掩码。具体来说，研究者使用的掩码解码器是 TransUNet。该网络将来自编码器的图像特征图和噪声掩码（随机初始化或迭代地来自编码过程）的连接作为输入，并输出对掩码的精确预测。解码器与用于图像生成和图像到图像转换的标准 U-Net 架构之间的一个区别是，在上采样之前，本文使用的 U-Net 顶部使用了带有交叉注意力层的 transformer 解码器层来合并编码的图像特征。

在视频模态下的应用

研究者将图像条件下的全景遮罩建模为：p(m|x)。基于给定视频的三维掩码（有一个额外的时间维度），本文的模型就可以直接适用于视频全景分割。为了适应在线 / 流媒体的视频设置，可以改用 p(m_t|x_t,m_(t-1),m_(t-k))建模，从而基于当前的图像和上一时刻的掩码生成新的全景掩码。如图 5 所示，这一变化可以通过将过去的全景掩码 (m_(t-1),m_(t-k)) 与现有的噪声掩码连接起来来实现。除了这个微小的变化之外，其他方面都与视频基础模型（p(m|x)）相同。这个模型很简单，对图像全景模型进行微调就可以应用到视频场景里。