Survey: Efficient MAE towards Large-Scale Vision Transformers

2024-12-24
| WACV 2024

Chen et al.

📑 投影片連結:20241224 Efficient MAE.pptx

當前挑戰

遮蔽自編碼器(Masked Autoencoder, MAE)在自監督視覺預訓練(Self-Supervised Pre-training)中取得了巨大的成功。標準 MAE 透過隨機遮蔽 75% 的圖像 Patch,僅將剩下 25% 的可見 Token 送入 Vision Transformer (ViT) 編碼器,大幅降低了計算複雜度。

若想進一步提升大模型預訓練的效率,最直接的想法是:能不能把遮蔽率進一步拉高到 85% 甚至 90%?

然而,直接提高遮蔽率會面臨以下瓶頸:

超高遮蔽率下重建難度不均導致表徵崩潰

當遮蔽率達到 85% 以上時,許多被遮蔽的 Patch 周遭完全沒有任何可見的鄰近像素。這些「孤立 Patch」的重建在數學上幾乎是純粹的隨機猜測,導致解碼器難以收斂,進而拖垮整個編碼器的特徵表徵品質。

均等損失權重讓模型過度卡在極端困難的不可預測 Patch 上

傳統 MAE 對所有被遮蔽的 Patch 給予相同的 MSE 損失權重。在極限遮蔽下,模型會耗費大量的梯度與容量去硬背那些毫無上下文線索的極端困難 Patch,忽略了真正具有語義關聯的有效結構。

小筆記:MAE 的非對稱設計與算力瓶頸
MAE 的核心優勢在於「非對稱架構」——龐大的 Encoder 只吃未被遮蔽的少數 Token,輕量的 Decoder 才吃完整 Token。因此,遮蔽率越高,Encoder 處理的 Token 越少,訓練速度越快。但遮蔽率若突破臨界點,特徵學習品質就會斷崖式下滑。


主要貢獻

這篇論文提出了 Efficient MAE,透過量化 Patch 重建難度並平衡梯度分佈,讓大尺度 Vision Transformer 在極高遮蔽率(85%~90%)下依然能維持優異的特徵學習能力。

解決: 超高遮蔽率下重建難度不均導致表徵崩潰

作者提出了量化指標與自適應遮蔽策略:
1. 鄰近像素平均距離(P2Dist 指標):計算每個被遮蔽 Patch 到最近可見 Patch 之間的空間與特徵距離,以此作為量化重建難度的客觀依據。
2. 基於難度的遮蔽策略(Difficulty-based Masking, DM):設定難度閾值 \beta,主動識別並過濾那些四周完全沒有可見資訊的極端困難 Patch,避免解碼器在毫無資訊量的區域空轉。

解決: 均等損失權重讓模型過度卡在極端困難的不可預測 Patch 上

作者提出了難度平坦化損失(Difficulty-Flatten Loss, DF Loss)
– 根據 P2Dist 動態調整不同難度 Patch 在重建損失中的權重。
– 降低極端容易(已有豐富周邊資訊)與極端困難(幾乎無資訊)Patch 的梯度貢獻,將模型的學習重心集中在「中等難度、具備足夠上下文且富含結構資訊」的 Patch 上。


延伸探討

極限遮蔽率下的加速效益與下游表現

實驗結果顯示,在 85%~90% 的極高遮蔽率下,Efficient MAE 不僅將預訓練速度提升了數倍、大幅節省 GPU 顯存開銷,而且在 ImageNet-1K 線性微調(Fine-tuning)以及下游物件偵測、語義分割任務上的性能,依然能與標準 75% 遮蔽率的 MAE 持平甚至微幅超越。