本文介紹我們組在單目標(biāo)跟蹤任務(wù)上的新工作：MixFormerV2: Efficient Fully Transformer Tracking。本工作主要解決了目前基于 transformer 架構(gòu)的跟蹤模型過于沉重，導(dǎo)致難以在 GPU 和 CPU 等邊緣設(shè)備上實(shí)時(shí)部署的問題。 我們通過簡潔有效的模型結(jié)構(gòu)設(shè)計(jì)和高效的基于知識(shí)蒸餾的模型壓縮，對(duì)于現(xiàn)有的 MixFormer 模型進(jìn)行了大幅度的輕量化提出了 MixFormerV2，同時(shí)依舊保持了穩(wěn)定的跟蹤精度性能。我們分別發(fā)布了 GPU 和 CPU 兩個(gè)版本的模型，在 LaSOT，TrackingNet，TNL2K 等 benchmark 上，其跟蹤性能和推理速度均能夠大幅度超越目前同量級(jí)的 SOTA 模型。

我們的工作已被 NeurIPS 2023 接收，論文、代碼和模型均已開源：

論文地址：

https://arxiv.org/abs/2305.15896

代碼地址：

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

基于 MixFormerV2，結(jié)合現(xiàn)有的 SAM，E2FGVI 等方法，我們開發(fā)了一個(gè)支持 video object tracking, segmentation 以及 inpainting 的 demo，下面是一些 demo 結(jié)果和界面演示。

研究動(dòng)機(jī)

在 MixFormer 被提出之后（詳見我們組的另一篇工作，本作的前身 MixFormer [1]），該全新的跟蹤范式和模型結(jié)構(gòu)的有效性被不斷驗(yàn)證，基于這種范式的新方法也不斷涌現(xiàn)（如 OSTrack，SimTrack 等），提升了目標(biāo)跟蹤任務(wù)的性能上限。 簡單來說，如圖所示傳統(tǒng)的跟蹤模型呈現(xiàn)雙流三階段的結(jié)構(gòu)，對(duì)于模板和當(dāng)前幀兩個(gè)輸入分別進(jìn)行特征提取、信息交互和定位預(yù)測三個(gè)模塊，而 MixFormer-style 的模型利用了 transformer 注意力機(jī)制的靈活性，整合了前兩個(gè)階段，使用一個(gè)通常是預(yù)訓(xùn)練好的 ViT backbone，對(duì)于兩個(gè)輸入同時(shí)進(jìn)行特征提取和融合，形成了目前主流的單流兩階段的跟蹤器范式。 然而，這種單流跟蹤模型的缺點(diǎn)也很明顯，因?yàn)槭褂妙A(yù)訓(xùn)練的 ViT，同時(shí)進(jìn)行特征提取和融合的建模會(huì)使得計(jì)算開銷非常大，導(dǎo)致模型的推理速度降低，尤其是在 CPU 設(shè)備上，這也使得模型的實(shí)際部署成為問題。

性能分析

既然是輕量化的工作，我們有必要對(duì)于現(xiàn)有的模型進(jìn)行一下 profile。以現(xiàn)有的 MixFormer 為例，我們使用在期刊擴(kuò)展中提出的以 vanilla ViT 為 backbone 版本的 MixViT，我們分別測試了不同模型層數(shù)，輸入圖像尺寸以及 MLP 維度比率的在 GPU 和 CPU 上的推理速度，如圖所示。

很自然的，當(dāng)模型深度越小時(shí)，模型的推理速度幾乎以線性提高。尤其的，MLP ratio 成為制約 CPU 速度的一個(gè)重要瓶頸，因?yàn)?GPU 的并行度非常高，所以簡單地增加維度并不會(huì)大幅度明顯地降低運(yùn)行效率，但是對(duì)于 CPU 卻會(huì)影響非常大。所以為了實(shí)現(xiàn) CPU 上的實(shí)時(shí)運(yùn)行，MLP 的隱藏層維度需要從常規(guī)的 x4 降低到 x1。

然后我們測試了 MixViT 的兩個(gè)預(yù)測頭，一個(gè)密集 pyramid 卷積定位頭，和一個(gè)基于 attention 的質(zhì)量分?jǐn)?shù)預(yù)測頭對(duì)于速度的影響，在 8 層的 backbone 基礎(chǔ)上，我們發(fā)現(xiàn)兩個(gè)預(yù)測模塊的計(jì)算開銷也是不能忽略的。當(dāng)將兩個(gè)預(yù)測頭都替換成我們提出的基于 token 的非常輕量的預(yù)測頭（兩個(gè)簡單的 MLP）時(shí)，GPU 上的推理速度能夠提升 84.4%。而當(dāng) backbone 部分被進(jìn)一步壓縮時(shí)，預(yù)測頭的替換所帶來性能提升的比重將會(huì)更大。

所以通過各個(gè)模塊分析，模型壓縮的思路已經(jīng)很自然了，首先我們需要想辦法將兩個(gè)笨重的預(yù)測頭替換為更加輕量的實(shí)現(xiàn)，然后將標(biāo)準(zhǔn)的 ViT backbone 從深度和寬度上都進(jìn)行減小。

方法

基于預(yù)測token的模型結(jié)構(gòu)

我們的 MixFormerV2 是一個(gè)完全基于 transformer 的跟蹤架構(gòu)，沒有任何的卷積操作和復(fù)雜的分?jǐn)?shù)預(yù)測模塊，跟蹤框和置信度的預(yù)測都通過一個(gè)簡單的 MLP 實(shí)現(xiàn)，整個(gè)模型非常得簡潔統(tǒng)一，如圖所示。這得益于我們的核心設(shè)計(jì)——四個(gè)特殊預(yù)測 token。

預(yù)測token參與的混合注意力機(jī)制

和原始的 MixViT 中的混合注意力模塊（Mixed Attention Module, MAM）不同，我們引入了四個(gè)特殊預(yù)測 token，用來在 backbone 中捕獲目標(biāo)模板和搜索區(qū)域之間的相關(guān)性。這些 token 可以在網(wǎng)絡(luò)每層中逐步地壓縮目標(biāo)信息，形成緊湊且稀疏的表征，用于之后的分類和回歸任務(wù)。 具體來說，模板、搜索區(qū)域的圖像特征 token 以及四個(gè)預(yù)定義的可學(xué)習(xí) prediction token，會(huì)輸入到我們提出的預(yù)測 token 參與的混合注意力模塊（Prediction-Token-Invoked Mixed Attention, P-MAM）。下標(biāo) t, s, e 分別表示模板，搜索區(qū)域以及 prediction token，P-MAM 的操作定義為：

和原始的 MixFormer 相同，我們會(huì)使用不對(duì)稱混合注意力機(jī)制來提高推理時(shí)效率。我們引入的四個(gè)可學(xué)習(xí) token 和標(biāo)準(zhǔn) ViT 中的 CLS token 實(shí)際上非常類似，它們會(huì)在跟蹤數(shù)據(jù)集上自動(dòng)地學(xué)習(xí)來壓縮目標(biāo)和搜索區(qū)域的信息并建模其交互關(guān)系。

基于token的直接預(yù)測

在有了已經(jīng)收集到豐富的模板和搜索區(qū)域信息的預(yù)測 token 之后，我們就可以直接使用這些 token 來進(jìn)行目標(biāo)定位和質(zhì)量分?jǐn)?shù)的預(yù)測了，并且通過幾個(gè)簡單統(tǒng)一的 MLP 就可以實(shí)現(xiàn)。 具體來說，對(duì)于定位回歸任務(wù)，每個(gè)預(yù)測 token 分別用來建模目標(biāo)框的一條邊，輸入一個(gè)共享參數(shù)的 MLP 進(jìn)行包圍框坐標(biāo)的回歸預(yù)測；對(duì)于置信度分類任務(wù)，我們首先要將四個(gè) token 的信息進(jìn)行整合，在這里我們簡單地將四個(gè) token 取均值就可以達(dá)到不錯(cuò)的效果，然后輸入一個(gè) MLP 輸出一個(gè)預(yù)測的置信度分?jǐn)?shù)。 值得一提的是，對(duì)于回歸任務(wù)，我們并不是直接去回歸包圍框的絕對(duì)坐標(biāo)，而是預(yù)測包圍框邊界的概率分布，這在實(shí)驗(yàn)中被驗(yàn)證能夠有效地提升模型精度，并且非常有利于下面會(huì)介紹的知識(shí)蒸餾過程。 這些基于 token 的預(yù)測頭對(duì)于分類和回歸任務(wù)都大大降低了計(jì)算復(fù)雜度，相比原始 MixViT 中密集卷積定位頭和使用 attention 的分?jǐn)?shù)模塊，形成了簡單統(tǒng)一并且 Fully Transformer 的模型架構(gòu)。

基于知識(shí)蒸餾的模型壓縮

為了進(jìn)一步提升模型的效率，我們提出了一種基于知識(shí)蒸餾的模型壓縮范式，包括兩個(gè)階段，第一階段密集到稀疏（dense-to-sparse）蒸餾，來使得原始的教師模型更好地傳遞知識(shí)到我們基于 token 的預(yù)測模型；第二階段從深層到淺層（deep-to-shallow）蒸餾，來進(jìn)行高效的模型剪枝。

Dense-to-Sparse Distillation

我們選取現(xiàn)有的 SOTA 模型 MixViT 作為我們的教師模型，但是存在的問題是，如何將教師模型的知識(shí)傳遞給學(xué)生 MixFormerV2？因?yàn)榻處熀蛯W(xué)生模型的輸出結(jié)構(gòu)是不一樣，原始 MixViT 使用的卷積定位頭會(huì)預(yù)測目標(biāo)框角點(diǎn)的二維分布，而我們的 MixFormerV2 的四個(gè) MLP 預(yù)測頭會(huì)回歸預(yù)測四個(gè)邊界的坐標(biāo)，這時(shí)候就體現(xiàn)出我們上面提到的預(yù)測目標(biāo)邊界的概率分布而不是絕對(duì)坐標(biāo)值的優(yōu)勢(shì)了。 具體來說，我們將目標(biāo)的上下左右邊界建模為四個(gè)隨機(jī)變量 ，而我們的 就是用來預(yù)測該隨機(jī)變量的概率分布：

最終的預(yù)測結(jié)果可以表示為該概率分布的數(shù)學(xué)期望：

而對(duì)于 MixViT 來說，其預(yù)測角點(diǎn)輸出本質(zhì)上是左上和右下坐標(biāo)的二維聯(lián)合概率分布，這種建模方式使得教師模型的密集預(yù)測頭輸出和學(xué)生模型的基于稀疏 token 的輸出就可以通過邊緣分布非常自然地聯(lián)系起來：

這時(shí)，使用 KL 散度損失函數(shù)，教師模型的輸出就可以作為軟標(biāo)簽來監(jiān)督學(xué)生模型，完成 dense-to-sparse 蒸餾階段的知識(shí)傳遞，如圖中 stage 1 所示。Deep-to-Shallow DistillationBackbone 始終是計(jì)算開銷的大頭，所以必須對(duì) backbone 進(jìn)行壓縮，所以我們提出了一種基于 feature mimicking 和 logits distillation 的由深到淺的壓縮方式，如圖中 stage 2 所示。 對(duì)于 logits distillation 我們施加 KL 散度即可，對(duì)于 feature mimicking，記 分別為學(xué)生和教師模型的特征圖，下標(biāo)表示層數(shù)的索引，施加 損失：，其中為師生匹配監(jiān)督的層數(shù)對(duì)。因?yàn)橹苯右瞥Ｐ椭械牟糠謱訒?huì)導(dǎo)致不連續(xù)和不一致的問題，所以我們探索了一種漸進(jìn)式的深度剪枝方法。 具體來說，不是直接使用教師模型從頭監(jiān)督一個(gè)更小的學(xué)生模型，我們讓學(xué)生模型初始化為一個(gè)和教師模型相同的副本，然后逐漸地撤除學(xué)生模型中的部分層，讓剩余的層在訓(xùn)練中擬合教師模型的表征，這個(gè)設(shè)計(jì)讓學(xué)生和教師模型的初始表征分布盡可能得一致，提供了一種更加平滑的遷移策略并降低了 feature mimicking 的難度。 形式地，令 為 MixFormerV2 第 層的輸出，一個(gè)注意力塊的計(jì)算可以表示為（為了簡便表示公式中忽略了 LN 操作）：

讓 為需要被刪除的層數(shù)的集合，我們會(huì)在這些層上施加一個(gè)衰退系數(shù) ：

在訓(xùn)練的前 個(gè) epoch， 會(huì)以余弦函數(shù)的方式衰減，這意味著這些層在從模型中逐漸撤除并最終成為一個(gè)恒等變換，壓縮后的模型可以直接將剩余層堆疊在一起得到。

實(shí)驗(yàn)

SOTA對(duì)比

我們實(shí)例化了 GPU 和 CPU 兩個(gè)版本的模型，分別為 MixFormerV2-B 和 MixFormerV2-S，可以看出，在保持強(qiáng)勁的性能指標(biāo)基礎(chǔ)上，我們的 MixFormerV2 在推理速度上都大幅超過了目前的主流跟蹤模型。

消融實(shí)驗(yàn)

針對(duì)我們的框架中的各個(gè)組件設(shè)計(jì)，我們都進(jìn)行了詳細(xì)的探究實(shí)驗(yàn)驗(yàn)證其有效性。

可視化

值得一提的是，我們對(duì) MixFormerV2 的 prediction token 的 attention map 進(jìn)行了可視化，我們能夠發(fā)現(xiàn)四個(gè) token 確實(shí)是在關(guān)注建模目標(biāo)的四個(gè)邊界，這也證明我們的方法的可靠性和可解釋性。

總結(jié)

我們的工作 MixFormerV2 基于現(xiàn)有的 sota 跟蹤模型 MixViT，改進(jìn)設(shè)計(jì)了簡潔的模型架構(gòu)和高效的壓縮方法，在多種硬件設(shè)備上實(shí)現(xiàn)了跟蹤模型性能精度和推理速度良好的平衡。我們希望 MixFormerV2 能夠有助于模型的實(shí)際落地應(yīng)用，并促進(jìn)高效跟蹤模型的發(fā)展。