[小样本图像分割]Deep Gaussian Processes for Few-Shot Segmentation

深度高斯过程的Few-Shot分割

论文地址

阅读本文需要一定的统计知识！

摘要

一种常见的小样本分割方法是用一个原型为每个类建模。虽然这些方法概念简单，但当目标外观分布在特征空间中是多模态或不是线性可分时，这些方法就会受到影响。

为了解决这一问题，我们提出了一种基于高斯过程(GP)回归的few-shot学习公式。通过GP的表达能力，我们的方法能够对分布在深层特征空间上的复杂外观进行建模。GP提供了一种捕捉不确定性的原则性方法，这是CNN解码器获得的最终分割的另一个强大线索。我们进一步利用我们方法的端到端学习能力来学习GP学习器的输出空间，确保分割掩码能得到一个更加丰富的编码表示。

存在的问题及解决方案

大多数近期的小样本分割方法是通过从支持集中提取一个原型使用点积运算和从查询图像中提取的特征进行比较。本质上，这导致了线性分类过程，因此依赖于特征空间在类上是线性可分的。事实上，最近的研究表明，即使训练了这种线性可分性，在许多情况下，不可见类的结果特征空间也不是线性可分的。一些工作认为，即使是单个物体也可能需要不同的原型来对应不同的部分。有鉴于此，一个有吸引力的替代方案是创建一个更灵活的、能够模拟复杂决策边界的few-shot学习模块。然而，必须小心地确保这样的模块不会过拟合支持集，另外它必须在计算上易于处理并且完全可微。

我们探索使用高斯过程回归来堆支持集进行建模以对查询集进行预测。在每个片段（episode）中，支持集通过编码器提供一组特征和掩码的相应值。基于这些，我们构建了一个GP。GP概率性地建模特征和掩码值之间的映射，并且在噪声的特定假设下足够灵活地充当插值器。在推理过程中，GP作用于查询特征，并根据每个特征的掩码值推断出概率分布，即预测后验。在GP的学习和推理过程中所做的所有计算都是可微分的，并且依赖于很容易集成为神经网络模块的操作。

将GPs与深度神经网络相结合是很有吸引力的，原因有几个。神经网络为GP提供了强大的视觉特征。反过来，GP预测后验分布给出了结合潜在高度相关样本的原则性方法。此外，预测后验的协方差提供了额外的不确定性信息。这一信息在few-shot学习中尤为重要。因为我们只得到几个训练的例子，我们不可避免地需要查询没见过(unseen)的外观。每当GP查询没见过的特征时，它就会以高协方差估计的形式传递这一信息。

由于高斯过程并不局限于一维输出空间，因此我们不仅仅只预测每个查询特征的单个掩码值。我们用一个轻量级的神经网络对给定的支持掩码进行编码。得到的掩码编码是一个多维特征向量，包含了更丰富的掩码表示。在推理过程中，我们的高斯过程few-shot学习器对这个输出空间进行预测。预测的掩码编码，连同不确定性估计，被输入到随后的解码器。这个解码器是一个神经网络，它结合这些线索来产生最终的分割。

方法

首先，通过图像编码器输入支持图像和查询图像，生成两组深层特征，即支持集${\left\{x_S^i\right\}}_i$和查询集${\left\{x_Q^j\right\}}_j$，每个特征$x^i\in {\mathbb{R}}^D$对应一幅图像中的一个空间位置。我们将支持和查询特征堆叠成一个矩阵中的行，分别产生$x_S$和$x_Q$。支持掩码通过另一个编码器输入以产生掩码编码${\left\{y_S^i\right\}}_i$，每个支持特征都有一个。这些被堆叠到一个向量$y_Q$以匹配$x_Q$的行。现在的目标是预测每个查询特征$x_Q^j$对应的掩码编码$y_Q^j$。本工作的思想是利用高斯过程推断条件概率分布$p\left({\mathbf{y}}_{\mathcal{Q}}\mid {\mathbf{x}}_{\mathcal{Q}},{\mathbf{x}}_{\mathcal{S}},{\mathbf{y}}_{\mathcal{S}}\right)$，即给定支持集下查询图像掩码编码的概率分布。然后将这种分布的表示形式输入解码器。该解码器将表示与从查询图像中提取的浅层特征相结合，以预测最终的分割。

高斯过程（Gaussian Processes）

我们寻找一个把输入和输出联系起来的函数，即$f(x)=y$，给定一个支持样本${\left\{x^i,y^i\right\}}_i$。高斯过程推断出这些函数的概率分布。利用这个分布，我们可以根据先前未见的数据$x$推断出$y$的条件分布。在我们的例子中，不可见的数据是查询特征$x_Q^j$，我们推理出对应掩码编码$y_Q^j$上的分布。高斯过程回归的思想是假设输出是某个均值$\mu$和协方差$\Sigma$的高斯联合分布。两个输出$y$和$y^{\prime }$之间的协方差$Cov(y,y^{\prime })$由协方差函数或者核κ定义，用于评估两个输出的对应关系。也就是说$Cov(y,y^{\prime })=\kappa (x,x^{\prime })$。核是一个正定函数，如果它的输入特征相似，相应的输出也应该相似。

高斯过程提供了一个具有概率输出的有原则的模型$f$，包含点估计和不确定性信息。GPs在输入空间中是非参数的，并存储所有给定的输入-输出对。因此，它们可以灵活地建模多模态输入数据以及复杂的决策边界。和函数的选择将控制决策边界的特性。这些特征仅在κ范围内使用，特征维数的时间复杂度为线性，这使得GPs适合基于深度特征的预测。重要的是，在看不见的数据上推断分布是一个完全可微的操作。基于高斯过程的Few-Shot学习器因此很容易集成为深度神经网络中的一个模块。

定义$y_S$和$y_Q$分别为所有的支持掩码和查询掩码向量。前者是可以观察到的，我们的目标是预测后者的分布。这里我们假设掩码编码是一维的，以方便标记。GP假设支持样本和查询样本是联合分布的：

通常情况下，先验的方法是假设${\mu }_S$和${\mu }_Q$为零。同时，协方差${\mathbf{K}}_{\mathcal{S}\mathcal{S}}=\kappa \left({\mathbf{x}}_{\mathcal{S}},{\mathbf{x}}_{\mathcal{S}}\right)$，${\mathbf{K}}_{\mathcal{Q}\mathcal{Q}}=\kappa \left({\mathbf{x}}_{\mathcal{Q}},{\mathbf{x}}_{\mathcal{Q}}\right)$，${\mathbf{K}}_{\mathcal{S}\mathcal{Q}}=\kappa \left({\mathbf{x}}_{\mathcal{S}},{\mathbf{x}}_{\mathcal{Q}}\right)$。这些由核$\kappa :{\mathbb{R}}^D\times {\mathbb{R}}^D\to \mathbb{R}$定义，在我们的实验中，我们采用了常用的平方指数核：

接下来，我们使用支持掩码编码$y_S$推断查询掩码编码$y_Q$的后验分布。也就是说，我们计算条件的参数:

我们假设对$y_S$的观察存在一定的噪声，例如方差为${\sigma }_y^2$的高斯噪声。这相当于在支持协方差矩阵$K_{SS}$上加上一个单位矩阵${\sigma }_y^{2}I$。那么，后验参数为：

元学习的高斯过程（Gaussian Processes for Meta Learning）

接下来，我们将一个GP元学习器整合到一个Few-Shot的分割框架中。首先，$K$个支持图像${\left\{I_{\mathcal{S}}^k\right\}}_k^K$输入图像编码器$F^{feat}:{\mathbb{R}}^{3\times H\times W}\to {\mathbb{R}}^{D\times \frac{H}{16}\times \frac{W}{16}}$来产生深度特征图$X_S^k$。这些深度特征图中的特征向量被收集来构建我们的支持集$x_S$。相应的注释，或着说支持掩码，输入到掩码编码器Fmask :{0,1}H×W→RE×H16×W16F^{\text {mask }}:\{0,1\}^{H \times W} \rightarrow \mathbb{R}^{E \times \frac{H}{16} \times \frac{W}{16}}Fmask :{ 0,1}H×W→RE×16H?×16W?中产生相应的掩码编码。这些被收集和堆叠以产生支持掩码编码向量$y_S$。注意，与之前的工作不同，我们同时为背景和前景建模。所有特征都被利用。使用相同的图像编码器，我们计算了查询特征$x_Q$。我们通过(3)-(5)获得查询掩码的后验，因此首先计算协方差矩阵块$K_{SS}$、$K_{SQ}$和$K_{QQ}$。然后通过(4)和(5)得到均值${\mathbf{\mu }}_{\mathcal{Q}\mid \mathcal{S}}$和协方差${\mathbf{\Sigma }}_{\mathcal{Q}\mid \mathcal{S}}$，这涉及到矩阵逆与另一个矩阵的乘法运算。在实践中，我们使用Cholesky分解并求解由此产生的三角形线性方程组。后验(3)产生掩码编码的联合概率分布。我们最感兴趣的是$y_Q^j$的边缘分布，它包含了期望的掩码编码及其不确定性。我们通过取均值和协方差的对角线得到这个，形成一个分布的表示形式：

其中[·]表示连接，$j$表示均值和协方差索引。注意，在GP中，空间顺序是无关紧要的。最终查询分布在支持集元素的顺序上是不变的，在查询特性的排列上是等变的。然而，在构造$z^j$之后，我们将它们进行排列，以匹配从图像中提取的原始特征图。然后将结果输入解码器$F^{\mathrm{d}\mathrm{e}\mathrm{c}}:{\mathcal{X}}_{\mathcal{Q},\text{ shallow }}\times {\mathbb{R}}^{(E+1)\times \frac{H}{16}\times \frac{W}{16}}\to [0,1{]}^{H\times W}$产生最终的分割。

掩码编码器（Mask Encoder）

到目前为止，我们假设掩码编码为一维值。在最好的情况下，GP模块的输出将是一个低分辨率的掩码。我们相信，解码器将受益于额外的指导，例如可能是关于局部形状或边缘的信息。为此，我们通过一个轻量级的神经网络输入掩码，该网络以与特征图相同的分辨率输出多维掩码编码。这个想法并不新鲜。

高斯过程能够模拟多维输出。为了简单和计算的原因，我们假定协方差函数是独立的，并且在输出维度上是等距的，这意味着协方差变成了块对角矩阵。在实践中，当GP被用来推理$p\left({\mathbf{y}}_{\mathcal{Q}}\mid {\mathbf{x}}_{\mathcal{Q}},{\mathbf{x}}_{\mathcal{S}},{\mathbf{y}}_{\mathcal{S}}\right)$我们通过一系列矩阵-向量乘法(4)找到平均值。如果掩码编码是多维向量，我们将它们作为行堆叠在一个矩阵中。列表示不同的维度，结果的均值是一个相同大小的矩阵。实质上，(4)是按掩码编码的每个维度计算一次。结果表示$z^j$包含一个多维均值和一个方差值。请注意，虽然在多维情况下有一个一般协方差是可能的，但它很快就变得难以计算。

实验结果

总结

我们提出了一种基于高斯过程回归的Few-Shot学习器用于Few-Shot分割任务。GP在深度特征空间对支持集进行建模，其灵活性允许它捕获复杂的特征分布。它对查询图像进行概率预测，提供点估计和额外的不确定性信息。这些预测被输入到一个CNN解码器，用来预测最终的分割。