【原】?東京大學&商湯&悉尼大學等提出融合了動態規劃、分治算法的MIM，實現綠色高效層次Transformer！已開源！

詳細信息如下：

摘要

本文提出了一種基于分層視覺Transformer(ViT)的掩蔽圖像建模(MIM)的有效方法，例如Swin Transformer，該方法允許分層ViT丟棄mask patch，只對可見patch進行操作。本文的方法包括兩個關鍵部分。

首先，對于窗口注意力，作者按照分治策略設計了一個群體窗口注意力（Group Window Attention）方案。為了減輕自注意的二次復雜性的影響，群體注意力鼓勵對任意大小的每個局部窗口內的可見patch進行統一的劃分，可以將大小相等的patch分組，然后在每個組內執行masked self-attention。

其次，作者通過動態規劃算法進一步改進分組策略，以最小化對分組patch的注意力的總體計算成本。因此，MIM現在可以以綠色高效的方式處理分層VIT。例如，作者可以將分層ViT的訓練速度提高約2.7倍，并將GPU內存使用量減少70%，同時在ImageNet分類方面仍具有競爭力，在下游COCO目標檢測基準上仍具有優勢。

Motivation

由于掩蔽語言建模（Masked Language Modeling ，MLM）在自然語言處理（NLP）領域的巨大成功以及視覺Transformer（ViT）的進步，掩蔽圖像建模（MIM）成為計算機視覺（CV）的一種有希望的自監督預訓練范式。MIM可以利用masked prediction從未標記數據中學習表示，例如，預測離散token、潛在特征或隨機掩蔽輸入圖像塊的原始像素。其中，具有代表性的工作掩蔽自動編碼器（Masked Autoencoder，MAE）表現出了有競爭力的性能和極高的效率。本質上，MAE為MIM提出了一種非對稱編碼器-解碼器架構，其中編碼器（例如，標準ViT模型）僅在可見patch上運行，而輕量級解碼器恢復所有patch以進行掩蔽預測。

一方面，非對稱編解碼結構顯著減少了預訓練的計算負擔。另一方面，MAE僅支持ViT結構作為編碼器，而大多數現代視覺模型采用層次結構，部分原因是需要處理視覺元素的尺度變化。事實上，層次結構和局部歸納偏置在各種CV任務中至關重要，這些任務需要不同級別或尺度的表示來進行預測，包括圖像分類和目標檢測。然而，如何將層次視覺Transformer（如Swin Transformer）集成到MAE框架中仍然不明確。此外，盡管SimMIM研究了用于MIM的Swin Transformer，可在可見和掩蔽patch上運行，但與MAE相比，計算成本高。

為此，作者本著綠色AI的精神，努力為MIM設計一種新的綠色方法，采用分層模型。本文的工作重點是將MAE的非對稱編碼器-解碼器結構擴展到分層視覺Transformer中，特別是代表性的模型Swin Transformer，以便僅對可見patch進行有效的預訓練。作者發現，主要障礙是局部窗口注意的局限性。盡管局部窗口注意廣泛應用于分層視覺Transformer中，但它不能很好地與隨機mask配合使用，因為它會創建各種大小的局部窗口，無法并行計算。

本文提供了解決這個問題的第一次嘗試。本文的方法在概念上很簡單，由兩個部分組成。首先，在分治原則的指導下，作者提出了一種組窗口注意力方案，該方案首先將可見patch數量不均勻的局部窗口劃分為若干個大小相等的組，然后在每個組中應用mask attention。其次，作者將前面提到的組劃分表述為一個約束優化問題，其中的目標是找到一個組劃分，使對分組token注意力的計算成本最小。受動態規劃概念和貪心原理的啟發，作者提出了一種最優分組算法，該算法自適應地選擇最優分組大小，并將局部窗口劃分為最小數量的組。

本文的方法是通用的，并且沒有對骨干模型的架構做任何修改，這樣就可以在可見和掩蔽patch上與baseline進行公平的比較。在本文的實驗評估中，作者觀察到本文的方法需要更少的訓練時間，消耗更少的GPU內存，而性能與baseline相同。具體而言，使用Swin-B，與SimMIM相比，本文的方法只需要一半的訓練時間和大約40%的GPU內存消耗，同時在ImageNet-1K上實現與SimMIM相當的83.7%的top-1微調精度。

方法

3.1 Preliminary

Notations

設表示輸入特征，其中C、H、W分別為X的通道數、高度和寬；表示訓練中隨機生成的空間mask，其中0表示一個patch對編碼器不可見，反之亦然。

Masked Image Modeling

MIM通過從其部分觀察中，預測輸入X的屏蔽部分來學習表示，即。根據操作，現有的MIM方法分為兩類。大多數方法使用Hadamard積進行掩蔽并保留掩蔽patch，即，其中M沿著通道維度廣播C次。與這些方法形成鮮明對比的是，掩蔽自動編碼器MAE (mask Autoencoders)提出在掩蔽階段拋去掩蔽patch，即：

MAE設計了一種非對稱的編碼器-解碼器結構：編碼器只對不帶mask token的可見patch進行操作；解碼器從可見patch和mask patch的表示中重建原始圖像。這種設計使MAE取得了具有競爭力的性能和極高的效率，例如，與所有patch操作相比，MAE達到了3倍的訓練速度。

然而，MAE只適用于各向同性的ViT（即非層次的ViT），目前尚不清楚如何將MAE的效率轉化為分層ViT，分層ViT在大多數視覺任務中幾乎一致優于各向同性的ViT。在本文中，作者試圖回答這個問題，并層次ViT提出了一個更綠色MIM的方法。

3.2 Green Hierarchical Vision Transformer for Masked Image Modeling

Base architecture

作者選擇swin Transformer作為本文的baseline，它主要由前饋網絡（FFN）和移位窗口注意組成。雖然FFN是逐點操作，并且只能對可見patch進行操作，但窗口注意力不能這樣做。

給定一個窗口大小p(例如，Swin為7)，窗口注意力首先將特征映射X劃分為個不重疊的局部窗口，其中每個包含個patch。由于每個窗口包含相同數量的patch，因此在每個窗口內并行地獨立執行MSA。然而，當局部窗口內patch數量不均勻時，如何高效并行計算注意力尚不清楚（如上圖所示，mask之后每個窗口內的patch可能不一致）。

為此，作者提出了一種高效的群體窗口注意力（Group Window Attention）方案，直接替換Swin中所有的移位的窗口注意力，使其只對可見patch進行操作。

Group Window Attention

針對上述問題，作者提出了一種組窗口注意方案，該方案大大提高了掩蔽特征上窗口注意的計算效率。給定mask特征，作者收集了一組不均勻的局部窗口，其中每個元素只包含可見token，相應的大小為。

如上圖所示，本文的組窗口注意力首先使用最優分組算法將不均勻的窗口劃分為幾個大小相等的組，然后在每個組內執行mask attention，以避免信息泄露。

3.3 Optimal Grouping with Dynamic Programming

General formulation

優化分組的第一步是找到一個索引分區與組大小相關：

其中為組的個數。條件約束分區包含所有沒有重復的局部窗口，并強制每個組的實際大小小于。基于分區，可以獲得一組分組token：

在這狙token上執行mask attention。最后，作者應用索引分區的逆操作來恢復輸出token的位置。

根據上面的公式，仍然有兩個問題沒有解決：1)如何選擇最優分組大小, 2)給定，如何獲得最優分區。為此，作者將目標表述為下面的min-min優化問題：

其中是一個cost函數，衡量的是分組token的注意力計算成本。直觀上，式(5)旨在找出最佳分組大小，使得的計算成本最小。方程(6)在方程(3)的約束下尋找最優分區。有了最優的分組大小，可以直接得到最優分區。

Group partition with Dynamic Programming

作者發現(6)和(3)式中的優化問題是具有相同容量的多個子集和問題(MSSP-I)的一種特殊情況，它是具有相同容量的0-1多個背包問題(MKP-I)的變種。在本文的例子中，群體的大小類似于背包的容量，可見token的數量類似于貨物的價值，貨物的重量與它們的價值相同，背包的數量是無界的。雖然一般的多背包問題是NP-complete的，但它的變體MSSP-I可以在偽多項式時間內使用動態規劃算法求解。具體來說，作者對單背包問題采用DP算法:

具體的偽代碼如下所示：

Cost function

因為本文主要關心效率，所以作者使用FLOPs來衡量分組token上的多頭注意力的計算復雜度，即：

其中C為通道數。雖然復雜度和group size的關系是二次的，但使用更小的可能產生更多的組，影響效率。因此，在訓練過程中可以自適應地確定最優 group size。

Putting everything together

模型從到，遍歷可能的組大小值，以找到最佳的組大小。對于每個選擇的分組大小，作者首先使用式(7)中的DP算法對窗口進行劃分，然后計算注意力在此劃分上的計算代價。選擇成本最小的群體作為最優群體規模。最優分組的偽代碼總結如下所示。

3.4 Masked Attention

由于不相鄰的局部窗口被劃分為相同的組，因此需要mask這些局部窗口之間的注意力權值，以避免這些局部窗口之間的信息交換。如上圖所示，在計算了attention map之后，作者只保留窗口內的注意權值(即對角線元素)，而丟棄窗口間的注意力權值。在檢索相對位置偏差時，作者也采用了類似的掩蔽方案，存儲每個token的原始絕對位置，并實時計算相對位置來檢索相應的偏差。

3.5 Batch-wise Random Masking

作者觀察到，每個樣本的隨機掩蔽策略會降低本文方法的效率：1)它可能會產生不同數量的局部窗口組，這對于被掩蔽注意力的并行計算是棘手的；2)當mask patch尺寸小于分層模型的最大patch尺寸時，有些patch可能同時包含可見和mask輸入。因此，作者提出將mask patch大小設置為與編碼器最大patch大小相同的值，并對相同GPU設備中的所有樣本使用相同的隨機mask。

實驗

上圖將本文的方法的效率與SimMIM進行了比較，可以看出本文方法的計算效率遠高于SimMIM。

由于層次模型有多個階段，具有不同的特征尺度，因此每個階段的最佳組大小也可能不同。對此，作者設計了一個實驗來分析不同階段的最優。在實驗中，作者隨機生成100個mask，計算不同選擇的成本，并在上圖中報告成本的平均值/標準偏差。一般來說，作者觀察到成本隨組大小的平方增加而增加，但在某些情況下，組大小正好等于局部窗口子集的和。

從上圖（a）可以看出，掩蔽比r在0.5到0.85之間變化時，本文的方法的性能非常穩定。在上圖（b）中，作者還研究了解碼器深度的影響。有趣的是，結果表明更少的解碼器塊可以產生更好的結果。這項研究支持SimMIM采用分層模型的簡單預測頭設計，并與MAE采用各向同性模型的觀察結果形成對比。此外，作者還研究了預訓練預算對本文的方法的影響。如上圖（c）所示，微調精度隨著訓練epoch的增加而穩步提高，并且似乎沒有停滯，這表明它有可能進一步提高性能。

如上表所示，窗口大小加倍的預訓練只會略微增加不到10%/20%的訓練時間/GPU內存，但會帶來適度的性能改善 。

上表展示了本文方法在ImageNet-1K上和其他SOTA方法的對比結果。

最后，作者評估了本文預訓練的模型到MS-COCO目標檢測和實例分割的遷移學習性能。可以看出，與有監督的預訓練Swin-B相比，本文的方法在所有指標方面都表現得更好。此外，作者還觀察到，本文的方法在密集預測任務上的性能仍與SimMIM相當。

總結

在本文中，作者提出了一種使用層次視覺Transformer（如Swin Transformer）的掩蔽圖像建模（MIM）的綠色方法，允許層次模型丟棄mask patch，只對可見patch進行操作。將高效的群體窗口注意力機制與基于DP算法的最優分組策略相結合，本文的方法可以訓練層次模型，并將速度提高約2.7倍，GPU內存消耗減少70%，同時在ImageNet分類方面仍具有競爭力，并具有下游MS-COCO目標檢測基準的優勢。

參考資料

[1]https:///abs/2205.13515
[2]https://github.com/LayneH/GreenMIM