動態視覺Backbone！港大團隊開源TransXNet，入選頂刊TNNLS2025

©作者 |婁蒙、俞益洲

單位 |香港大學計算機系

香港大學計算機系俞益洲教授及其研究團隊開發了一種新型的通用視覺 backbone—TransXNet，該模型同時融合了全域性和動態性。其創新之處在於引入了一個即插即用的模組，名為 Dual Dynamic Token Mixer（D-Mixer）。

與以往的卷積和 self-attention 混合模組或全域性與區域性資訊融合模組不同，D-Mixer 專注於動態性的捕捉。具體來說，D-Mixer 能夠利用網路上一層的 feature maps 資訊，生成用於提取當前 feature maps 的全域性和區域性動態權重，從而顯著增強模型對多尺度特徵的動態表徵能力。

為了驗證 TransXNet 的效能，研究團隊分別在 ImageNet-1K、COCO 2017 和 ADE20K 資料集上進行了廣泛的實驗，結果表明，TransXNet 作為一個通用視覺模型，展現出了卓越的效能。

論文標題：

TransXNet: Learning both global and local dynamics with a dual dynamic token mixer for visual recognition

論文地址：

https://arxiv.org/abs/2310.19380

https://ieeexplore.ieee.org/document/10948454

程式碼地址：

https://github.com/LMMMEng/TransXNet

團隊主頁：

https://i.cs.hku.hk/~yzyu/index.html

引言

在本文中，作者深入探討了 self-attention 和卷積在權重計算方式上的不同，並基於此提出了改進方案。作者指出，self-attention 之所以效能卓越，部分原因在於其權重（即 attention matrix）會根據不同的輸入動態變化，這種 input-dependent 的特性使得模型能夠更好地捕捉輸入資料的特徵。

相比之下，傳統的卷積操作中的權重是靜態的，與輸入資料無關，即 input-independent，這限制了其表徵能力。據此，作者進一步分析了將 self-attention 和卷積簡單結合時可能遇到的問題：

表徵能力差異：由於卷積缺乏類似 self-attention 的動態性，這導致在卷積-attention混合模組中，self-attention的語義建模能力可能會被削弱，因為卷積部分的靜態特性限制了整體模組的表徵能力。
深層模型的挑戰：在深層模型中，self-attention 能夠利用先前特徵圖的全域性和區域性資訊來生成動態的 attention matrix，從而提高模型的效能。然而，卷積核的靜態特性使其無法利用這些資訊來動態調整，這限制了模型在深層結構中的表現。

方法

為了解決上述問題，作者提出了一個即插即用模組 D-Mixer。該模組能夠動態地捕獲全域性和區域性資訊，從而使網路兼具大感受野和穩定的歸納偏置。

如圖 1（a）所示，對於給定的輸入特徵圖，D-Mixer 首先會將其沿通道拆分為兩部分，並分別將這兩部分輸入到兩個不同的模組中，即 Input-dependent Depthwise Convolution（IDConv）以及 Overlapping Spatial Reduction Attention（OSRA）。

輸出結果會進一步進行拼接，並且由一個輕量級的 Squeezed Token Enhancer（STE）進行整合。以下為不同模組的具體描述：

Input-dependent Depthwise Convolution (IDConv)：如圖 1（b）所示，對於任意輸入特徵圖，作者首先使用一個自適應平均池化層聚合空間上下文，並將空間維度壓縮為 K×K 大小，其中 K 為待生成的動態卷積核的大小。

緊接著，輸出的特徵圖被送入兩個 1×1 卷積層中，並生成一個多組空間注意力圖，維度為 A∈G×C×K×K，其中 G、C 分別表示注意力圖的組數以及輸入特徵圖的通道大小。並且在 G 維度上使用 softmax 函式對注意力圖進行歸一化，使其具備自適應的選擇能力。

最後，透過將注意力圖和一組維度同為 G×C×K×K 的可學習引數進行逐元素相乘來生成動態卷積核，並將其用於輸入特徵圖的特徵提取。IDConv 可以根據上一層特徵圖攜帶的語義資訊來動態地生成卷積核的權重，進而可以與 self-attention 進行深度嵌合。

Overlapping Spatial Reduction Attention (OSRA)：如 1 圖（c）所示，OSRA 的整體流程和 PVT 中提出的 Spatial Reduction Attention（SRA）類似，都採用了 token-to-region 的全域性建模方式。

主要不同之處在於，SRA 在降取樣的過程中使用的是 Non-overlapping 的卷積，即卷積核大小和步長大小一致。因此，一些連續的區域性資訊會在這個過程中被破壞，從而降低 key/value 分量包含的空間結構資訊。

據此，OSRA 引入了 Overlapping 卷積來進行降取樣，即卷積核的大小大於步長。此外，OSRA 還引入了一個輕量的區域性運算元（以 3×3 的 depthwise 卷積實現）來進一步提煉 key/value 分量的區域性資訊。

Squeezed Token Enhancer (STE)：為了高效地將 IDConv 和 OSRA 處理後的特徵進行融合，本文引入了一個輕量級的特徵融合模組 STE。如圖1（d）所示，STE 透過一個 3×3 的 depthwise 卷積和兩個通道壓縮的 1×1 卷積來實現特徵融合。值得一提的是，STE 相較於直接使用一個全連線層進行特徵融合具有更好的效能以及更優的計算複雜度。

TransXNet 整體架構：如圖 2 所示，TransXNet 的基本構成部分主要包含 D-Mixer 和 Multiscale Feed-forward Network（MS-FFN）。其中，MS-FFN 使用了基於不同卷積核大小的 depthwise 卷積來提取多尺度資訊。由於 TransXNet 同樣沿用了 ResNet 和 Swin Transformer 中的金字塔設計，因此，該網路可以應用到各類視覺任務中。