MixFormerV2：基于Transformer的高效跟踪器

本文介绍我们组在单目标跟踪任务上的新工作：MixFormerV2: Efficient Fully Transformer Tracking。本工作主要解决了目前基于 transformer 架构的跟踪模型过于沉重，导致难以在 GPU 和 CPU 等边缘设备上实时部署的问题。   我们通过简洁有效的模型结构设计和高效的基于知识蒸馏的模型压缩，对于现有的 MixFormer 模型进行了大幅度的轻量化提出了 MixFormerV2，同时依旧保持了稳定的跟踪精度性能。我们分别发布了 GPU 和 CPU 两个版本的模型，在 LaSOT，TrackingNet，TNL2K 等 benchmark 上，其跟踪性能和推理速度均能够大幅度超越目前同量级的 SOTA 模型。

我们的工作已被 NeurIPS 2023 接收，论文、代码和模型均已开源：  

论文地址：

https://arxiv.org/abs/2305.15896

代码地址：

https://github.com/MCG-NJU/MixFormerV2

基于 MixFormerV2，结合现有的 SAM，E2FGVI 等方法，我们开发了一个支持 video object tracking, segmentation 以及 inpainting 的 demo，下面是一些 demo 结果和界面演示。

研究动机

在 MixFormer 被提出之后（详见我们组的另一篇工作，本作的前身 MixFormer [1]），该全新的跟踪范式和模型结构的有效性被不断验证，基于这种范式的新方法也不断涌现（如 OSTrack，SimTrack 等），提升了目标跟踪任务的性能上限。   简单来说，如图所示传统的跟踪模型呈现双流三阶段的结构，对于模板和当前帧两个输入分别进行特征提取、信息交互和定位预测三个模块，而 MixFormer-style 的模型利用了 transformer 注意力机制的灵活性，整合了前两个阶段，使用一个通常是预训练好的 ViT backbone，对于两个输入同时进行特征提取和融合，形成了目前主流的单流两阶段的跟踪器范式。   然而，这种单流跟踪模型的缺点也很明显，因为使用预训练的 ViT，同时进行特征提取和融合的建模会使得计算开销非常大，导致模型的推理速度降低，尤其是在 CPU 设备上，这也使得模型的实际部署成为问题。  

性能分析

既然是轻量化的工作，我们有必要对于现有的模型进行一下 profile。以现有的 MixFormer 为例，我们使用在期刊扩展中提出的以 vanilla ViT 为 backbone 版本的 MixViT，我们分别测试了不同模型层数，输入图像尺寸以及 MLP 维度比率的在 GPU 和 CPU 上的推理速度，如图所示。

很自然的，当模型深度越小时，模型的推理速度几乎以线性提高。尤其的，MLP ratio 成为制约 CPU 速度的一个重要瓶颈，因为 GPU 的并行度非常高，所以简单地增加维度并不会大幅度明显地降低运行效率，但是对于 CPU 却会影响非常大。所以为了实现 CPU 上的实时运行，MLP 的隐藏层维度需要从常规的 x4 降低到 x1。 

然后我们测试了 MixViT 的两个预测头，一个密集 pyramid 卷积定位头，和一个基于 attention 的质量分数预测头对于速度的影响，在 8 层的 backbone 基础上，我们发现两个预测模块的计算开销也是不能忽略的。当将两个预测头都替换成我们提出的基于 token 的非常轻量的预测头（两个简单的 MLP）时，GPU 上的推理速度能够提升 84.4%。而当 backbone 部分被进一步压缩时，预测头的替换所带来性能提升的比重将会更大。

所以通过各个模块分析，模型压缩的思路已经很自然了，首先我们需要想办法将两个笨重的预测头替换为更加轻量的实现，然后将标准的 ViT backbone 从深度和宽度上都进行减小。

方法

基于预测token的模型结构

我们的 MixFormerV2 是一个完全基于 transformer 的跟踪架构，没有任何的卷积操作和复杂的分数预测模块，跟踪框和置信度的预测都通过一个简单的 MLP 实现，整个模型非常得简洁统一，如图所示。这得益于我们的核心设计——四个特殊预测 token。  

预测token参与的混合注意力机制

和原始的 MixViT 中的混合注意力模块（Mixed Attention Module, MAM）不同，我们引入了四个特殊预测 token，用来在 backbone 中捕获目标模板和搜索区域之间的相关性。这些 token 可以在网络每层中逐步地压缩目标信息，形成紧凑且稀疏的表征，用于之后的分类和回归任务。   具体来说，模板、搜索区域的图像特征 token 以及四个预定义的可学习 prediction token，会输入到我们提出的预测 token 参与的混合注意力模块（Prediction-Token-Invoked Mixed Attention, P-MAM）。下标 t, s, e 分别表示模板，搜索区域以及 prediction token，P-MAM 的操作定义为：  

和原始的 MixFormer 相同，我们会使用不对称混合注意力机制来提高推理时效率。我们引入的四个可学习 token 和标准 ViT 中的 CLS token 实际上非常类似，它们会在跟踪数据集上自动地学习来压缩目标和搜索区域的信息并建模其交互关系。  

基于token的直接预测

在有了已经收集到丰富的模板和搜索区域信息的预测 token 之后，我们就可以直接使用这些 token 来进行目标定位和质量分数的预测了，并且通过几个简单统一的 MLP 就可以实现。   具体来说，对于定位回归任务，每个预测 token 分别用来建模目标框的一条边，输入一个共享参数的 MLP 进行包围框坐标的回归预测；对于置信度分类任务，我们首先要将四个 token 的信息进行整合，在这里我们简单地将四个 token 取均值就可以达到不错的效果，然后输入一个 MLP 输出一个预测的置信度分数。   值得一提的是，对于回归任务，我们并不是直接去回归包围框的绝对坐标，而是预测包围框边界的概率分布，这在实验中被验证能够有效地提升模型精度，并且非常有利于下面会介绍的知识蒸馏过程。    这些基于 token 的预测头对于分类和回归任务都大大降低了计算复杂度，相比原始 MixViT 中密集卷积定位头和使用 attention 的分数模块，形成了简单统一并且 Fully Transformer 的模型架构。  

基于知识蒸馏的模型压缩

为了进一步提升模型的效率，我们提出了一种基于知识蒸馏的模型压缩范式，包括两个阶段，第一阶段密集到稀疏（dense-to-sparse）蒸馏，来使得原始的教师模型更好地传递知识到我们基于 token 的预测模型；第二阶段从深层到浅层（deep-to-shallow）蒸馏，来进行高效的模型剪枝。  

Dense-to-Sparse Distillation

我们选取现有的 SOTA 模型 MixViT 作为我们的教师模型，但是存在的问题是，如何将教师模型的知识传递给学生 MixFormerV2？因为教师和学生模型的输出结构是不一样，原始 MixViT 使用的卷积定位头会预测目标框角点的二维分布，而我们的 MixFormerV2 的四个 MLP 预测头会回归预测四个边界的坐标，这时候就体现出我们上面提到的预测目标边界的概率分布而不是绝对坐标值的优势了。   具体来说，我们将目标的上下左右边界建模为四个随机变量 ，而我们的 就是用来预测该随机变量的概率分布：

最终的预测结果可以表示为该概率分布的数学期望：

而对于 MixViT 来说，其预测角点输出本质上是左上和右下坐标的二维联合概率分布，这种建模方式使得教师模型的密集预测头输出和学生模型的基于稀疏 token 的输出就可以通过边缘分布非常自然地联系起来：

这时，使用 KL 散度损失函数，教师模型的输出就可以作为软标签来监督学生模型，完成 dense-to-sparse 蒸馏阶段的知识传递，如图中 stage 1 所示。Deep-to-Shallow DistillationBackbone 始终是计算开销的大头，所以必须对 backbone 进行压缩，所以我们提出了一种基于 feature mimicking 和 logits distillation 的由深到浅的压缩方式，如图中 stage 2 所示。   对于 logits distillation 我们施加 KL 散度即可，对于 feature mimicking，记 分别为学生和教师模型的特征图，下标表示层数的索引，施加 损失：，其中为师生匹配监督的层数对。因为直接移除模型中的部分层会导致不连续和不一致的问题，所以我们探索了一种渐进式的深度剪枝方法。 具体来说，不是直接使用教师模型从头监督一个更小的学生模型，我们让学生模型初始化为一个和教师模型相同的副本，然后逐渐地撤除学生模型中的部分层，让剩余的层在训练中拟合教师模型的表征，这个设计让学生和教师模型的初始表征分布尽可能得一致，提供了一种更加平滑的迁移策略并降低了 feature mimicking 的难度。 形式地，令 为 MixFormerV2 第 层的输出，一个注意力块的计算可以表示为（为了简便表示公式中忽略了 LN 操作）：

让 为需要被删除的层数的集合，我们会在这些层上施加一个衰退系数 ：

在训练的前 个 epoch， 会以余弦函数的方式衰减，这意味着这些层在从模型中逐渐撤除并最终成为一个恒等变换，压缩后的模型可以直接将剩余层堆叠在一起得到。  

实验

SOTA对比

我们实例化了 GPU 和 CPU 两个版本的模型，分别为 MixFormerV2-B 和 MixFormerV2-S，可以看出，在保持强劲的性能指标基础上，我们的 MixFormerV2 在推理速度上都大幅超过了目前的主流跟踪模型。  

消融实验

针对我们的框架中的各个组件设计，我们都进行了详细的探究实验验证其有效性。

可视化

值得一提的是，我们对 MixFormerV2 的 prediction token 的 attention map 进行了可视化，我们能够发现四个 token 确实是在关注建模目标的四个边界，这也证明我们的方法的可靠性和可解释性。  

总结

我们的工作 MixFormerV2 基于现有的 sota 跟踪模型 MixViT，改进设计了简洁的模型架构和高效的压缩方法，在多种硬件设备上实现了跟踪模型性能精度和推理速度良好的平衡。我们希望 MixFormerV2 能够有助于模型的实际落地应用，并促进高效跟踪模型的发展。