CVPR | 2025 | MAP：通過掩碼自回歸預訓練釋放混合 Mamba - Transformer 視覺骨干網絡的潛力

• 論文：https://arxiv.org/pdf/2410.00871
• 代碼：https://github.com/yunzeliu/MAP
• （代碼）鏡像：https://gitee.com/apuppyliu-cong/MAP.git
• 會議：CVPR
• 年份：2025

創新點

• 貢獻主要有以下三點：
  • 第一，首次提出了一種用于預訓練混合 Mamba - Transformer 視覺骨干網絡的新方法，在統一框架下提升了混合骨干網絡以及純 Mamba 和純 Transformer 骨干網絡的性能。
  • 第二，為了便于分析，我們還對純 Mamba 在自回歸預訓練中的關鍵組件進行了深入分析，發現其有效性取決于預訓練順序與 Mamba 掃描順序的一致性，以及合適的令牌掩碼比例。
  • 第三，我們證明了所提出的 MAP 方法在各種 2D 和 3D 數據集上，顯著提升了 Mamba - Transformer 和純 Mamba 骨干網絡的性能。大量的消融實驗證實了每個設計選擇的合理性和有效性。

初步研究

表 1. 初步研究。我們使用 ViT - B 和 Vim - B 作為默認配置。AR 策略按行優先順序處理圖像令牌，而 MAE 按照默認設置運行。對于對比學習，我們僅使用裁剪和縮放的數據增強，并使用 MoCov2 進行預訓練。所有實驗都在 224x224 的分辨率下進行。AR 的掩碼令牌數量設置為 40 個令牌（20%）。

1. 實驗表明，MAE 更適合 Transformer 的預訓練，而 AR 更適合 Mamba 的預訓練。
   1.1 對于 ViT 而言，應用 MAE 策略對于建立令牌之間的雙向關聯、從而提高性能至關重要
   1.2 對于 Vim，更重要的是對前后令牌之間的連續性進行建模
2. 經過深入分析發現，與掃描順序一致的自回歸預訓練以及適當的掩碼比例是 Mamba 預訓練的關鍵
   

表 2. AR 順序對下游任務的影響。Vim（R）指采用行優先掃描的 Vim。Vim（C）指采用列優先掃描的 Vim。AR（R）指行優先自回歸預訓練。AR（C）指列優先自回歸預訓練。結果表明，當 AR 預訓練設計與 Mamba 的掃描順序一致時，可獲得最佳性能。

圖 3. AR 預訓練和 Mamba 掃描的不同順序。行優先和列優先的順序使網絡能夠以不同的方式和序列感知局部信息。

3. 自回歸預訓練的掩碼比例。
   

表 3. 掩碼比例對 AR 預訓練的影響。分別掩碼 1 個令牌（0.5%）、20 個令牌（10%）、40 個令牌（20%）、60 個令牌（30%）、100 個令牌（50%）和 140 個令牌（70%），同時記錄在下游任務上的微調結果。實驗表明，適當的掩碼比例對于自回歸預訓練很重要。

初步結論

得出以下三個結論，作為設計混合骨干網絡預訓練的參考：

• MAE 更適合 Transformer，而 AR 更適合 Mamba。
• 對于 Transformer 的 MAE 預訓練，非對稱結構和適當的掩碼比例很重要。
• 對于 Mamba 的 AR 預訓練，適當的 AR 順序和掩碼比例很重要。

方法

本文的重點是研究如何預訓練混合 Mamba - Transformer 框架，而非設計混合框架本身。

確定一個混合網絡

結果表明，采用 MMMTMMMT 的混合方式效果最佳。在比較 Mamba - R * 與 MMMMMMTT 時，發現在 Mamba 之后添加 Transformer 增強了其長上下文建模能力，從而提升了性能。在比較 MMMMMMTT 與 TTMMMMMM 時，我們觀察到僅僅在 Mamba 之后附加 Transformer 并不能充分發揮該架構的潛力。

這表明在開始部分融入 Transformer 對于提取足夠的局部特征至關重要。我們認為 MMMTMMMT 方法有效地平衡了局部特征提取和上下文建模增強，因此將其作為我們的默認配置。

圖 4. 不同的混合模型設計。（d）取得了最佳結果，被設置為默認模型，并稱為 HybridNet。

表 4. Mamba - Transformer 骨干網絡的混合設計。所有實驗都是從零開始訓練的。Mamba - R表示 24 個 Mamba - R [40] Mamba 層加上 8 個額外的 Mamba 層。DeiT表示 24 個 DeiT [39] Transformer 層加上 8 個額外的 Transformer 層。MMMMMMTT 表示 24 個 Mamba 層之后跟隨 8 個 Transformer 層。TTMMMMMM 表示 8 個 Transformer 層之后跟隨 24 個 Mamba 層。TMMMTMMM 表示由 1 個 Transformer 層和 3 個 Mamba 層組成的單元，重復 8 次。MMMTMMMT 表示由 3 個 Mamba 層之后跟隨 1 個 Transformer 層組成的單元，重復 8 次。

方法

圖2.（a）MAE預訓練。其核心在于基于未掩碼的令牌重建被掩碼的令牌，以建立全局雙向上下文理解。（b）AR預訓練。它著重于建立上下文之間的關聯，其可擴展性在大型語言模型領域已得到充分驗證。（c）MAP預訓練（我們提出的方法）。我們的方法首先對輸入圖像進行隨機掩碼處理，然后以逐行自回歸的方式重建原始圖像。這種預訓練方法在對局部特征的上下文特征以及局部特征之間的關聯進行建模方面展現出顯著優勢，使其與Mamba-Transformer混合架構具有高度的兼容性。（d）在ImageNet-1K上不同預訓練策略下的性能提升。我們發現MAE預訓練更適合Transformer，而AR更適配Mamba。另一方面，MAP更適合Mamba-Transformer骨干網絡。此外，MAP在對純Mamba或純Transformer骨干網絡進行預訓練時也表現出令人矚目的性能，這體現了我們方法的有效性和廣泛適用性。

如圖 2（c）所示，對于給定圖像，HybridNet 首先將經過隨機掩碼處理的圖像映射到特征空間，隨后借助 Transformer 解碼器按行對原始圖像進行解碼。

假設將圖像 $\mathbf{I}$ 按行劃分：
I={r1,r2,...,rM}I=\left\{r_{1}, r_{2}, ..., r_{M}\right\}I={r1?,r2?,...,rM?}
每行 $r_i$ 由以下令牌組成：
ri={xi1,xi2,...,xiN}r_{i}=\left\{x_{i 1}, x_{i 2}, ..., x_{i N}\right\}ri?={xi1?,xi2?,...,xiN?}
在每行中選取一部分令牌進行掩碼處理。

令 Mi?{1,2,...,N}M_{i} \subset\{1,2, ..., N\}Mi??{1,2,...,N} 表示行 $r_i$ 中被掩碼令牌的索引。

對于給定的行 $r_i$ ，需同時預測所有被掩碼的令牌：
$$p\left(x_{ij}|x_{i,j\notin M_i},r_{<i}\right)$$

其中，$r_{<i}$ 指的是第 $i$ 行之前的所有行。

第 $i$ 行中令牌的預測取決于所有先前的行以及該行內可見的令牌。這可以表示為：
$$p\left(r_i|r_{<i}\right)=\prod _{j=1}^{N}p\left(x_{ij}|x_{i,j\notin M_i},r_{<i}\right)$$
整體的損失函數為預測令牌的負對數似然之和：
$$\mathcal{L}=?\sum _{i=1}^M\sum _{j\in M_i}\log p\left(x_{ij}|x_{i,j\notin M_i},r_{<i}\right)$$

掩碼機制

實驗了不同的掩碼策略，包括隨機掩碼、順序掩碼和對角線掩碼。實驗表明，隨機掩碼能產生最佳結果。

圖 5. 不同的掩碼策略。隨機掩碼策略產生最佳結果。

掩碼比例

MAP的transformer解碼器

使用 Transformer 解碼器而非 Mamba 解碼器的原因是，Transformer 解碼器可以通過應用解碼器掩碼，基于編碼器的特征進行區域重建。相比之下，Mamba 解碼器由于其單向掃描特性，難以同時重建整個局部區域。

解碼器采用獨特的行解碼策略，允許一次對一行令牌進行自回歸解碼，增強了網絡捕獲局部特征和區域間上下文關系的能力。

圖 6. 不同的解碼器掩碼。綠色表示激活。白色表示非激活。

表 6. 解碼器掩碼設計。我們的 MAP 解碼器策略取得最佳結果。

重建目標

重建目標與 MAE 一致，我們將歸一化的原始像素作為重建目標，并采用 MSE 損失。

實驗

ImageNet-1k 上的 2D 分類

預訓練階段，我們采用 50% 掩碼比例的隨機掩碼策略，僅使用隨機裁剪作為數據增強策略。在所有設置下訓練 1600 個 epoch。在微調階段，我們直接微調 400 個 epoch 并報告結果。

表 8. ImageNet-1k 分類結果。吞吐量在 A100 GPU 上計算。內存開銷以 128 的批處理大小衡量。我們的結果以藍色突出顯示。