NeurIPS 2021 ｜ Twins：重新思考高效的視覺注意力模型設計

Twins 是美團和阿德萊德大學合作提出的視覺注意力模型，相關論文已被 NeurIPS 2021 會議接收。本文主要講述 Twins 解決的難點、設計和實現思路，以及在美團場景的探索落地，希望能對從事視覺演算法研發的同學有所幫助和啟發。

概要

Twins ^[1] 是美團和阿德萊德大學合作提出的視覺注意力模型，相關論文已被 NeurIPS 2021 會議接收，程式碼也已在GitHub上進行開源。NeurIPS（Conference on Neural Information Processing Systems）是機器學習和計算神經科學相關的學術會議，也是人工智慧方向的國際頂級會議。

Twins 提出了兩類結構，分別是 Twins-PCPVT 和 Twins-SVT：

• Twins-PCPVT 將金字塔 Transformer 模型 PVT ^[2] 中的固定位置編碼（Positional Encoding）更改為團隊在 CPVT ^[3] 中提出的條件式位置編碼 （Coditional Position Encoding, CPE），從而使得模型具有平移等變性（即輸入影像發生平移後，輸出同時相應發生變化），可以靈活處理來自不同空間尺度的特徵，從而能夠廣泛應用於影像分割、檢測等變長輸入的場景。
• Twins-SVT 提出了空間可分離自注意力機制（Spatially Separable Self-Attention，SSSA）來對影像特徵的空間維度進行分組，分別計算各區域性空間的自注意力，再利用全域性自注意力機制對其進行融合。這種機制在計算上更高效，效能更優。

Twins 系列模型實現簡單，部署友好，在 ImageNet 分類、ADE20K 語義分割、COCO 目標檢測等多個經典視覺任務中均取得了業界領先的結果。

背景

2020 年 9 月，谷歌的視覺注意力模型 （Vision Transformer， ViT） ^[4] 成功將原本用於自然語言處理的 Transformer ^[5] 應用到視覺的分類任務中。ViT 將一幅輸入影像切分為若干個影像塊（Patch），並把一個影像塊類比為一個文字（Word）作為 Transformer 編碼器的輸入（如圖 1 所示），經過 L 層的編碼器處理後使用普通的多層感知機（Multilayer Perceptron, MLP）對映到類別空間。ViT 的模型效能大幅超過了卷積神經網路，此後迅速發展成為了當前視覺領域研究的主要熱點。

![圖1 視覺注意力模型（ViT）將用於自然語言處理任務的 Transformer 應用於視覺任務（來源：ViT [4]）](https://p1.meituan.net/travel...)

Transformer 編碼器中多頭注意力（Multi-head attention）的基本計算方法由下式給定，其中 Q、K、V 分別為 Query（查詢）、Key（鍵）、Value（值） 的縮寫，d 為編碼維度，softmax 為歸一化函式，注意力機制可以理解為對輸入按照相關性加權的過程。

原生的視覺注意力模型做主幹網路並不能很好地適配目標檢測、語義分割等常用的稠密預測任務。此外，相比於卷積神經網路，ViT 計算量通常要更大，推理速度變慢，不利於在實際業務中應用。因此設計更高效的視覺注意力模型，並更好地適配下游任務成為了當下研究的重點。香港大學、商湯聯合提出的金字塔視覺注意力模型 PVT ^[2] 借鑑了卷積神經網路中的影像金字塔正規化來生成多尺度的特徵，這種結構可以和用於稠密任務的現有後端直接結合，支援多種下游任務，如圖 2（c）所示。但由於 PVT 使用了靜態且定長的位置編碼，通過插值方式來適應變長輸入，不能針對性根據輸入特徵來編碼，因此效能受到了限制。另外，PVT 沿用了 ViT 的全域性自注意力機制，計算量依然較大。

![圖2 PVT 將卷積神經網路（a）的金字塔正規化遷移到視覺注意力模型（b）得到（c），以適應分類、檢測、分割多種任務（來源：PVT ^[2]）](https://p0.meituan.net/travel...)

微軟亞研院提出的 Swin ^[6] 複用了 PVT 的金字塔結構。在計算自注意力時，使用了對特徵進行視窗分組的方法（如圖 3 所示），將注意力機制限定在一個個小的視窗（紅色格子），而後通過對視窗進行錯位使不同組的資訊產生互動。這樣可以避免計算全域性自注意力而減少計算量，其缺點是損失了全域性的注意力，同時由於視窗錯位產生的資訊互動能力相對較弱，一定程度上影響了效能。

![圖3 Swin 計算每一個紅色格子的區域性自注意力，通過不同層間的視窗移位來使各區域性注意力之間產生互動（來源：Swin [6]）](https://p0.meituan.net/travel...)

視覺注意力模型設計的難點

簡單總結一下，當前視覺注意力模型設計中需要解決的難點在於：

• 高效率的計算：縮小和卷積神經網路在運算效率上的差距，促進實際業務應用；
• 靈活的注意力機制：即能夠具備卷積的區域性感受野和自注意力的全域性感受野能力，兼二者之長；
• 利於下游任務：支援檢測、分割等下游任務，尤其是輸入尺度變化的場景。

Twins 模型設計

從這些難點問題出發，基於對當前視覺注意力模型的細緻分析，美團視覺智慧部重新思考了自注意力機制的設計思路，提出了針對性的解決方案。首先將 PVT ^[2] 和 CPVT ^[4] 相結合，形成 Twins-PCPVT 來支援尺度變化場景的下游任務。再從自注意機制的效率和感受野角度出發，設計了相容區域性和全域性感受野的新型自注意力，叫做空間可分離自注意力 （Spatially Separable Self-Attention，SSSA）， 形成了 Twins-SVT。

Twins-PCPVT

Twins-PCPVT 通過將 PVT 中的位置編碼（和 DeiT ^[7] 一樣固定長度、可學習的位置編碼）替換為 CPVT ^[4] 中的條件位置編碼 （Conditional Positional Encodings，CPE）。生成 CPE 的模組叫做位置編碼器（Positional Encoding Generator， PEG），PEG 在 Twins 模型中的具體位置是在每個階段的第 1 個 Transformer Encoder 之後，如下圖 4 所示：

條件位置編碼

下圖 5 展示了團隊在 CPVT ^[4] 中提出的條件位置編碼器的編碼過程。首先將 $N*d$ 的輸入序列轉為 $H*W*d$ 的輸入特徵，再用 $F$ 根據輸入進行條件式的位置編碼，而且輸出尺寸和輸入特徵相同，因此可以轉為 $N*d$ 序列和輸入特徵進行逐元素的加法融合。

其中，編碼函式 $F$ 可以由簡單的深度可分離卷積實現或者其他模組實現，PEG 部分的簡化程式碼如下。其中輸入 feat_token 為形狀為 $B*N*d$ 的張量，$B$ 為 batch，$N$ 為 token 個數，$C$ 為編碼維度（同圖 5 中 $d$）。將 feat_token 轉化為 $B*d*H*W$ 的張量 cnn_feat 後，經過深度可分離卷積 (PEG) 運算，生成和輸入 feat_token 相同形狀的張量，即條件式的位置編碼。

class PEG(nn.Module):
    def __init__(self, in_chans, embed_dim):
        super(PEG, self).__init__()
        self.peg = nn.Conv2d(in_chans, embed_dim, 3, 1, 1, bias=True, groups=embed_dim)
        
    def forward(self, feat_token, H, W):
        B, N, C = feat_token.shape
        cnn_feat = feat_token.transpose(1, 2).view(B, C, H, W) 
        x = self.peg(cnn_feat) + cnn_feat
        x = x.flatten(2).transpose(1, 2)
        return x

由於條件位置編碼 CPE 是根據輸入生成，支援可變長輸入，使得 Twins 能夠靈活處理來自不同空間尺度的特徵。另外 PEG 採用卷積實現，因此 Twins 同時保留了其平移等變性，這個性質對於影像任務非常重要，如檢測任務中目標發生偏移，檢測框需隨之偏移。實驗表明 Twins-PCPVT 系列模型在分類和下游任務，尤其是在稠密任務上可以直接獲得效能提升。該架構說明 PVT 在僅僅通過 CPVT 的條件位置編碼增強後就可以獲得很不錯的效能，由此說明 PVT 使用的位置編碼限制了其效能發揮。

Twins-SVT

Twins-SVT （如下圖 6 所示）對全域性注意力策略進行了優化改進。全域性注意力策略的計算量會隨著影像的解析度成二次方增長，因此如何在不顯著損失效能的情況下降低計算量也是一個研究熱點。Twins-SVT 提出新的融合了區域性-全域性注意力的機制，可以類比於卷積神經網路中的深度可分離卷積 （Depthwise Separable Convolution），並因此命名為空間可分離自注意力（Spatially Separable Self-Attention，SSSA）。與深度可分離卷積不同的是，Twins-SVT 提出的空間可分離自注意力（如下圖 7 所示）是對特徵的空間維度進行分組，並計算各組內的自注意力，再從全域性對分組注意力結果進行融合。

空間可分離自注意力採用區域性-全域性自注意力（LSA-GSA）相互交替的機制，分組計算的區域性注意力可以高效地傳導到全域性。LSA 可以大幅降低計算成本，複雜度從輸入的平方 $O(H^2W^2d)$ 降為線性的 $O(mnHWd)$。其中分組區域性注意力 LSA 關鍵實現（初始化函式略）如下：

class LSA(nn.Module):
    def forward(self, x, H, W):
        B, N, C = x.shape
        h_group, w_group = H // self.ws, W // self.ws # 根據視窗大小計算長（H）和寬（W）維度的分組個數
        total_groups = h_group * w_group
        x = x.reshape(B, h_group, self.ws, w_group, self.ws, C).transpose(2, 3) # 將輸入根據視窗進行分組 B* h_group * ws * w_group * ws * C
        qkv = self.qkv(x).reshape(B, total_groups, -1, 3, self.num_heads, C // self.num_heads).permute(3, 0, 1, 4, 2, 5) # 計算各組的 q， k， v 
        q, k, v = qkv[0], qkv[1], qkv[2]
        attn = (q @ k.transpose(-2, -1)) * self.scale # 計算各組的注意力
        attn = attn.softmax(dim=-1) # 注意力歸一化
        attn = self.attn_drop(attn) # 注意力 Dropout 層
        attn = (attn @ v).transpose(2, 3).reshape(B, h_group, w_group, self.ws, self.ws, C) # 用各組內的區域性自注意力給 v 進行加權
        x = attn.transpose(2, 3).reshape(B, N, C)
        x = self.proj(x) # MLP 層
        x = self.proj_drop(x) # Dropout 層
        return x

高效融合 LSA 注意力的 GSA 關鍵實現（初始化函式略）如下。相比於 ViT 原始的全域性自注意力，GSA 的 K、V 是在縮小特徵的基礎上計算的，但 Q 是全域性的，因此注意力仍然可以恢復到全域性。這種做法顯著減少了計算量。

class GSA(nn.Module):
    def forward(self, x, H, W):
        B, N, C = x.shape
        q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) # 根據輸入特徵 x 計算查詢張量 q
        x_ = x.permute(0, 2, 1).reshape(B, C, H, W)
        x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1) # 縮小輸入特徵的尺寸得到 x_
        x_ = self.norm(x_) # 層歸一化 LayerNorm
        kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) # 根據縮小尺寸後的特徵後 x_，計算 k, v 
        k, v = kv[0], kv[1]
        attn = (q @ k.transpose(-2, -1)) * self.scale # 計算全域性自注意力
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
        x = (attn @ v).transpose(1, 2).reshape(B, N, C) # 根據全域性自注意力對 v 加權
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

從上述程式碼中可以看出，SVT 系列在實現上採用現有主流深度學習框架中的已有操作，不需要額外的底層適配，因此部署起來比較方便。

實驗

ImageNet-1k 分類

Twins-PCPVT 和 Twins-SVT 在 ImageNet-1k 分類任務上，相比同等量級模型均取得 SOTA 結果，吞吐率佔優。另外，Twins 支援 TensorRT 部署，Twins-SVT-S 模型使用 NVIDIA TensorRT 7.0 推理可以有 1.6 倍的加速，吞吐率可以從 PyTorch 實現的 1059（images/s）提升到 1732。

ADE20K 分割

在語義分割任務 ADE20K 上，Twins 模型做主幹網分別使用 FPN 和 Upernet 後端，相比 PVT 和 Swin 也達到了更好結果，見下表 2：

COCO 目標檢測（RetinaNet 框架）

在經典的 COCO 目標檢測任務中，使用 RetinaNet 框架，Twins 模型大幅優於 PVT。而且 Twins-PCPVT 系列證明 PVT 在通過 CPVT 的編碼方式增強之後，可以媲美 Swin 同量級模型，見下表 3：

COCO 目標檢測（Mask-RCNN 框架）

在 Mask-RCNN 框架下，Twins 模型在 COCO 上也有很好的效能優勢，且在更長時間訓練（3x）時得以保持，見下表 4：

在高精地圖多要素語義分割場景的應用

高精地圖是自動駕駛中的關鍵組成部分，在美團無人配送、網約車等業務承擔著非常重要的作用。道路場景關鍵要素的語義提取作為高精建圖的前序流程，對建圖的質量有直接的影響。多要素語義分割是語義提取的重要一環，業界一般採用經典的語義分割演算法來實現。

此處，我們以 DeepLab 系列 ^[8] 為代表做介紹，分割模型通常分為編碼和解碼兩個階段，使用卷積神經網路提取特徵，並採用空間金字塔遲化（Spatial Pyramid Pooling），以及不同尺度空洞卷積（Atrous Conv）操作（如下圖 8a 所示）來增加全域性感受野。這種設計一方面受限於卷積神經網路的特徵提取能力，另一方面對全域性關係的建模能力有限，導致在分割任務上對細節關注不夠，邊緣往往不夠清晰。

![圖8 經典語義分割模型架構（DeepLabV3+ [8]）](https://p0.meituan.net/travel...)

Twins 雖然大幅提升了視覺注意力模型的效率和效能，但為了保持和卷積神經網路的接近的推理效率，我們仍需要對模型的後端結構作進一步的優化。不同於論文中為了與其他方法做公平對比而使用的較重的 FPN ^[9] 或 UperNet ^[10] 後端，我們設計瞭如下圖 9 所示的簡單輕量的後端，並在業務資料集上的效能和推理速度之間取得了很好的平衡。這個後端是根據 Twins 的特性而設計，由於 Twins 兼顧了全域性和區域性兩種注意力，因此後端無需採用複雜的設計來增大感受野，只通過各尺度特徵的線性變化和縮放，就直接恢復到相同尺寸並進行拼接（Concat），簡單維度變換後就可以輸出分割結果。

從下圖 10 的分割結果對比看，Twins 為主幹網的模型可以提取得到更精細的影像邊緣，如隔離帶、道路標牌、路燈杆等關鍵道路要素和標註真值（Ground Truth）之間的差異更小。

總結

視覺注意力模型是當前視覺領域的研究重點，並且已經在各類視覺任務上展示了相比經典卷積神經網路的優越性，但在效率上仍需精心優化，在效果上也需要繼續提升。探索設計更高效率的注意力模型並促進前沿的視覺研究轉向工業落地，對美團業務也具有重要的意義。

此次，美團和阿德萊德大學合作設計的 Twins 系列模型架構，有效降低了計算成本，提升了模型效能，更好地支援瞭如檢測和分割等稠密任務。此外，我們將 Twins 應用在美團高精地圖的要素語義分割場景中，帶來了更精細的分割結果，提升了高精地圖的建圖質量。後續，視覺團隊將持續探索高效的視覺注意力模型設計，並期望在美團更廣泛的業務場景中得到實踐和應用。

參考文獻

• [1] Twins: Revisiting the Design of Spatial Attention in Vision Transformers
• [2] Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions
• [3] Conditional Positional Encodings for Vision Transformers
• [4] An image is worth 16x16 words: Transformers for image recognition at scale
• [5] Attention Is All You Need
• [6] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
• [7] Training data-efficient image transformers & distillation through attention
• [8] Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation
• [9] Panoptic Feature Pyramid Networks
• [10] Unified Perceptual Parsing for Scene Understanding

作者簡介

祥祥、田值、張勃、曉林在視覺智慧部，海兵、華夏在自動車配送部。

團隊簡介及招聘資訊

美團視覺智慧部 AutoML 演算法團隊旨在通過 AutoML 及前沿的視覺技術賦能公司各項業務、加速演算法落地，涵蓋 AutoML、分割、檢測（2D、3D）、Self-training 等技術方向。歡迎感興趣的校招和社招同學傳送簡歷至：chuxiangxiang@meituan.com，實習、正式均可。

美團自動車配送部高精地圖團隊是高精地圖技術研發團隊，我們的職責是為美團自動駕駛提供高精度、高鮮度、大範圍的高精地圖服務。高精地圖是涉及多種學科的綜合技術，不僅需要依託 SLAM、地理測繪、深度學習、多感測器定位等演算法進行高精地圖的構建，而且還需要利用大資料技術、高效能運算、高併發服務等進行大規模高精地圖處理、儲存和查詢服務。高精地圖團隊長期招聘計算機視覺、SLAM、系統開發等專家，感興趣的同學可以將簡歷傳送至：tech@meituan.com（郵件主題：美團高精地圖）。

閱讀美團技術團隊更多技術文章合集

前端 | 演算法 | 後端 | 資料 | 安全 | 運維 | iOS | Android | 測試

| 在公眾號選單欄對話方塊回覆【2021年貨】、【2020年貨】、【2019年貨】、【2018年貨】、【2017年貨】等關鍵詞，可檢視美團技術團隊歷年技術文章合集。

| 本文系美團技術團隊出品，著作權歸屬美團。歡迎出於分享和交流等非商業目的轉載或使用本文內容，敬請註明“內容轉載自美團技術團隊”。本文未經許可，不得進行商業性轉載或者使用。任何商用行為，請傳送郵件至tech@meituan.com申請授權。