论文地址
代码地址：暂无
论文名称：MobileVOS: Real-Time Video Object Segmentation Contrastive Learning meets Knowledge Distillation

摘要

本文研究了手机等资源受限设备上的半监督视频对象分割问题。我们将该问题建模为一个蒸馏任务，通过该任务，我们证明了具有有限内存的小型时空记忆网络可以取得与现有技术相媲美的结果，但计算开销仅为一部分(在三星Galaxy S22上每帧32毫秒)。具体而言，我们提供了一个理论基础框架，将知识蒸馏与有监督的对比表示学习统一起来。这些模型能够共同受益于像素级的对比学习和来自预训练教师的蒸馏。我们通过在标准的DAVIS和YouTube基准测试中实现具有竞争力的J&F来验证这一损失，其中运行速度快了5倍，参数少了32倍。

引言

背景：memory-based的方法在半监督视频分割任务中很流行，但是内存网络往往在内存使用和准确性之间进行权衡。同时，文献中没有一种方法能够在手机上实时执行并保留最新的结果。

本文的方法：本文关注基于记忆的网络，并开始弥合无限记忆和有限记忆网络之间的鸿沟，这是通过知识蒸馏实现的。

(1)本文提出了一个逐像素的表征蒸馏损失，旨在将一个大的、无限的记忆教师的结构信息传递给一个较小的、有限的记忆学生。

(2)使用了一个简单的边界感知采样策略，它能提高训练收敛性

(3)本文提供了一个自然的泛化来实现一个有监督的像素级对比表示目标。

(4)使用真实标签作为结构信息的另一个场所，本文使用超参数在表征蒸馏损失和对比学习损失之间进行插值，并将其作为一个统一的损失。

贡献：

(1)本文将知识蒸馏和有监督的对比学习结合起来，设计了一种新的基于表示的损失，它弥合了大的无限记忆模型和小的有限记忆模型之间的差距。此外，本文证明了一个简单的边界感知像素采样策略可以进一步提高结果和模型的收敛性。

(2)利用这个统一的损失，本文证明了一个通用的网络设计可以获得与最先进的网络设计相媲美的结果，同时运行速度比之前快达5倍，参数少达32倍。

(3)在没有复杂的架构或内存设计变更的情况下，本文提出的损失可以在移动设备(三星Galaxy S22)上解锁实时性能(30FPS+)，同时保持与先进技术的竞争性能。

方法

说明：教师模型利用了无限的记忆，而学生模型只保留了前一帧和第一帧的图像和掩码的记忆。由于教师有更多的记忆，蒸馏的任务是学习跨多帧一致的特征。对比表示学习的辅助目标是在物体边界附近鼓励判别性的像素级特征。

Representation Distillation

本文在最后的全连接层之前利用得到的表征，更早之前可能会丢失大量有利于知识迁移的结构信息。

具体来说，本文提出了一个逐像素的表示蒸馏损失，可以在这两个模型之间传递结构信息。为了提取结构信息，我们提出构建相关矩阵：Cs, Ct，它们的维度都是HW*HW，分别来自学生模型和教师模型，这两个矩阵可以捕获所有像素对之间的关系。公式如下：

其中Zs，ZT表示在最后的逐点卷积和上采样层之前提供的d维的L2归一化的学生模型和教师模型得到的表示。

受[1]的启发，得到的最终的损失函数如下：

这种损失的两个组成部分可以解释为正则化项，以及两个模型之间的相关性对齐，而这里使用log2算子来提高对虚假相关性的鲁棒性。它相当于最大化学生和教师表征之间的像素级别的相互信息。

蒸馏损失函数为：(这其实是上面损失的一个特例，后面会说)

其中H2和I2分别基于矩阵的估计量，分别类似于Renyi熵和2阶互信息，参考文献:[2]
这是通过推导得到的，更多的可以看论文的附件。

Unification with contrastive learning

问题：知识蒸馏的有效性取决于学生和教师模型的相对能力差距。随着能力差距的缩小，这种训练制度的规模变得很小。

已有的结论：使用自监督作为辅助任务可以改进常规蒸馏

解决办法：(1)首先构建另外一个目标管关联矩阵：

其中Y是含有2个类别(对象与背景)的空间降采样、独热编码标签。

(2)将两个目标相关矩阵进行耦合，以提供在这两个训练模式之间进行插值的方法：

其中Cty可以代入到上面的Lrepr损失函数中去。

当考虑表征Z和w=0时，Lrepr损失函数就变成了有监督的对比学习：

其中Sim是表示中两个独立像素之间的余弦相似度，Pi为第i个像素的正指标集合。

对w的讨论：

w=0时，退化为有监督的对比损失；当w=1时，退化为蒸馏损失。并不是说相等，只是等价。

Boundary-aware sampling

存在的问题：(1)大多数预测误差发生在物体的边界上，下图是采用随机像素和边界像素采样策略训练的两个模型的验证精度，黑色是随机采样的。

(2)为所有像素构造相关矩阵的计算量太大

解决办法：提出只对物体边界附近的像素进行采样。

好处：这种采样策略限制了蒸馏梯度只流经导致下游预测误差的像素，同时也使公式的计算效率更高；这种修改可以提高模型的整体收敛性。

Logit distillation

在两个模型之间的logit空间引入一个额外的KL散度项：

pS，pT分别为学生和教师模型预测的概率，用一个温度项τ参数化，以软化(τ>1)或锐化(τ<1)这两个预测。与传统logit蒸馏的唯一区别是本文只使用物体边界周围或附近的像素。

最后的损失为：

伪代码为：

实验

相关设置

对于所有的实验，我们用一个额外的线性嵌入层并且使用一个预训练的STCN[3]来训练学生作为老师。

学生模型对STCN做了一些修改：value encoder被替换成了MobileNetV2，query encoder被替换成了ResNet18；同时在解码器之前使用了ASPP模块；我们选取固定的内存队列长度为2，当队列满了遵循先进先出原则。

实验结果

在DAVIS2016上，最好的模型只比STCN低0.3 J&F和比RDE-VOS低0.2 J&F，但比它们运行快4倍和3倍。真的很快，达到100FPS了，精度也没下降多少。

在DAVIS2017上，本文的模型只比STCN教师模型稍差0.3 J&F。但是，本文比STCN和RDE-VOS快了近5倍和4倍。

在YouTube2019上，本文的方法在没有使用BL30K预训练的情况下优于RDE-VOS，并且在使用它的情况下具有竞争力。不同于RDE-VOS为YouTube评估训练不同的损失权重以避免在未见类上的过拟合，我们展示了使用与DAVIS评估相同的权重集的预测。

在模型尺寸和推理速度上，本文的MobileVOS模型明显更小，同时在短视频和长视频中保持恒定的低延迟：

结论

在这项工作中，本文提出了一个简单的损失，将知识蒸馏和对比学习结合起来。通过在半监督视频对象分割的背景下应用该损失，本文以很小的计算成本和模型大小取得了与当前最好的结果竞争的结果。这些模型足够紧凑，可以适应标准移动设备，同时保留实时延迟。

参考文献

[1]Roy Miles, Adrian Lopez Rodriguez, and Krystian Mikolajczyk. Information Theoretic Representation Distillation. BMVC, 2022. 3, 13

[2]Luis G. Sanchez Giraldo and Jose C. Principe. Information theoretic learning with infinitely divisible kernels. ICLR, 2013. 3

[3]Ho Kei Cheng, Yu-Wing Tai, and Chi-Keung Tang. Rethinking space-time networks with improved memory coverage for efficient video object segmentation. NeurIPS. 1, 2, 5, 6, 7, 12