論文筆記 — Green Hierarchical Vision Transformer for Masked Image Modeling
Code: https://github.com/LayneH/GreenMIM
Paper: https://arxiv.org/abs/2205.13515
Introduction
MAE(Masked Auto-Encoder) 為 MIM(Masked Image Model) 提出了一種非對稱編碼器-解碼器架構，其中編碼器（例如，標準 ViT 模型）僅在可見補丁上運行，輕量級解碼器恢復所有補丁以進行掩碼預測。
一方面，非對稱編碼器-解碼器架構顯著降低了預訓練的計算負擔。 另一方面，MAE 僅支持各向同性 ViT架構作為編碼器，而大多數現代視覺模型採用分層結構，部分原因是需要處理尺度變化 視覺元素。
如何將分層視覺轉換器（例如 Swin Transformer）集成到 MAE 框架中仍然不是很簡單。 儘管 SimMIM 的工作已經實現了 MIM 的 Swin Transformer，我們發現即使是 SimMIM 的基本尺寸模型也無法在具有八個 32GB GPU 的單台機器上進行訓練。 計算負擔使得更廣泛的研究人員難以深入這一研究領域。
我們提出了一種用於 Masked Image Modeling 的 Green Hierarchical Vision Transformer，它提倡一種更實用的方法，可以顯著提高效率。
我們的方法不會對主幹模型的體系結構進行任何修改，因此我們可以與在可見和屏蔽補丁上運行的基線進行同類比較。 我們的方法需要更少的訓練時間並且消耗更少的 GPU 內存，同時在 ImageNet1K和 MS-COCO對象檢測上的性能與基線相當。
Approach
Green Hierarchical Vision Transformer for Masked Image Modeling
Adapting existing hierarchical ViTs for MIM
典型的分層 ViT 主要由Feed-Forward Network (FFN) 和 Efficient Attention 組成，它們不能僅在可見的補丁上運行。 我們發現主要障礙是分層 ViT 的歸納偏差，即
- 具有非重疊窗口分區的窗口注意力
- 卷積與池化引起的局部性
在 FFN 或具有重疊窗口分區的注意力中。 這些運算符與隨機掩碼不兼容，因為它會創建各種大小的局部窗口，這些窗口對於並行計算是不可行的。 為此，我們提出了兩個關鍵見解，使大多數分層 ViT 能夠以綠色方式僅在可見補丁上運行。
- Group Window Attention: 如下圖，
2. 結合稀疏卷積： 我們的直覺是掩碼輸入的特徵可以被視為稀疏張量，其中僅保留（少量）可見塊的特徵，而其他特徵被省略。 從這個角度來看，我們建議用最初為稀疏 3D 點雲數據設計的高度優化的稀疏卷積 直接替換分層 ViT 中的所有捲積/池化操作。
Optimal Grouping with Dynamic Programming
(1) General formulation
最優分組的第一步是找到一個關於組大小 gs 的索引分區 Π：
其中 ng 是結果組的數量。 等式 (3) 中的條件將分區約束為包含所有沒有重複的本地窗口，並強制每個組的大小需要小於 gs 。
在(4)結果上執行 masked attention 。最後，我們應用分區 Π的逆運算來恢復輸出標記的位置。 有了上面的公式，還有兩個問題沒有解決:
- 如何選擇最優組大小 gs∗
- 如何在給定 gs∗ 的情況下獲得最優分區 Π∗。
為此，我們將目標制定為以下 min-min 優化問題
其中 C(·) 是一個成本函數，用分組的標記衡量注意力的計算成本。等式（5）旨在找到最佳組大小 gs*，即最佳分區的計算成本。 gs* 是最小值。 等式（6）在等式（3）的約束下搜索最優分區。 有了最優組大小，我們可以直接得到最優分區 Π∗ = Partition(gs∗, {wi} )。 接下來，我們將詳細介紹如何解決上述優化問題。
(2) Group partition with Dynamic Programming
我們發現等式（6）和（3）中的優化問題是具有相同容量的多子集，是眾所周知的 0–1 多背包問題的變體。 在我們的例子中，組大小類似於背包的容量，可見令牌的數量 {wi}類似於貨物的價值，貨物的重量與其價值相同，背包的數量是無限的。
(3) Cost function
因為我們主要關心效率，所以我們使用 FLOPs 來衡量多頭注意力對分組標記的計算複雜度，即：
其中 C 是通道數，複雜度是二次方的 gs，使用較小的 gs 可能會產生更多的組（和更多的填充）並且效率不佳。 因此，最佳組大小是在訓練期間自適應確定的。
Masked Attention
因為不相鄰的本地窗口被劃分到相同的組中，所以需要屏蔽注意力權重以避免這些本地窗口之間的信息交換。 如圖所示，計算出注意力圖後，我們僅保留窗口內註意力權重（下圖藍色部分）並丟棄窗口間注意力權重(下圖灰色部分)。 ，我們存儲每個token的原始絕對位置並即時計算相對位置以檢索相應的偏差。
Batch-wise Random Masking
我們觀察到每個樣本的隨機掩蔽策略會降低我們方法的效率， 因此，我們建議將掩碼補丁大小設置為與編碼器的最大補丁大小相同的值（例如：32）並且對所有使用相同的隨機掩碼。
Experiments
Conclusion
結合高效的 Group Window Attention 方案、基於 DP 算法的最優分組策略和稀疏卷積，我們的方法可以將分層模型的訓練速度提高 ∼2.7 倍，並將 GPU 內存消耗減少多達 70%，同時仍然享受 ImageNet 分類的競爭性能和下游 MS-COCO 對象檢測基準的優越性。 我們希望這項工作能夠促進未來以有效性和效率為目標的自我監督學習方法。
Limitations
我們算法的局限性之一是它需要分批屏蔽(Batch Mask)方案以實現最佳效率。 儘管這種限制對 MIM 預訓練的影響很小，但它限制了我們的方法在更廣泛的設置上的應用。
Broader Impact
這項工作為具有分層 ViT 的 MIM 提出了一種綠色方法，顯著減輕了 MIM 的繁重計算負擔。
