0. 基本信息

DropPos: Pre-Training Vision Transformers by Reconstructing Dropped Positions. InNeurIPS 2023.

論文：arxiv.org/pdf/2309.03576

代碼：github.com/Haochen-Wang409/DropPos

今天介紹我們在自監督視覺預訓練領域的一篇原創工作DropPos:Pre-Training Vision Transformers by Reconstructing Dropped Positions.

目前 DropPos 已被 NeurIPS 2023 接收，相關代碼已開源，有任何問題歡迎在 GitHub 提出。

1. TL;DR

我們提出了一種全新的自監督代理任務 DropPos，首先在 ViT 前向過程中屏蔽掉大量的 position embeddings (PE)，然后利用簡單的 cross-entropy loss 訓練模型，讓模型重建那些無 PE token 的位置信息。這個及其簡單的代理任務就能在多種下游任務上取得有競爭力的性能。

2. Motivation

在 MoCo v3 的論文中有一個很有趣的現象：ViT 帶與不帶 position embedding，在 ImageNet 上的分類精度相差無幾。

表 1. MoCo v3 中的實驗現象 (原文第 6 頁)

這一實驗結果背后，隱含著「ViT 的建模主要關注于不同 patch 的 visual appearence，對于 position 的 awareness 較差」這一信息。即，如果把圖片切 patch 然后再隨機打亂之后，ViT 能夠在亂序的情況下準確識別該圖片的類別。這一點和人類直覺有很大出入。同時，有可能是因為 ViT 過擬合到了 ImageNet 這個特定數據集導致的。

基于此，我們首先做了一些 tiny experiments，探究 position awareness 與模型的識別準確率到底是否有正相關的關系。具體來說，我們凍結了 MoCo v3 和 MAE 的 pre-train/fine-tune 權重，在其后接一個全連接層，并用 position classification 這個任務做 linear probing。即，在 forward 過程中隨機丟棄 75% 的 PE，并把 ViT 的 feature 映射到 196 維 (一張圖有 14x14 個 patch)，期望讓最終的線性層正確分類該 patch 的位置。

表 2. Position awareness 對于下游任務的影響

表中結果表明，fine-tune 后的模型權重，更適合預測位置這一任務。說明「強大的對位置的建模能力，對于圖像分類任務是有益的」。基于此，我們想探究一種能夠提升 ViT 對于位置建模能力的全新自監督代理任務。

一種可行的方案是「簡單地把 ViT 的 PE 隨機丟棄一部分，然后讓模型預測這些不帶 PE 的 token 的精確位置」，即 reconstructDroppedPositions (DropPos).

圖 1. DropPos 與 CL 和 MIM 的對比

DropPos 有如下的優勢：

對比 CL，DropPos 不需要精心設計的數據增強 (例如 multi-crop)。

對比 MIM，DropPos 不需要精心設計的掩碼策略和重建目標。

下面我們介紹 DropPos 的具體運行流程

3. Method

圖 2. DropPos 的流程圖

即使 DropPos 的想法很直觀也很簡單，但這類方法一直沒有成為預訓練的主流，主要是由于在設計上有以下三個難點：

如果簡單地把所有 PE 丟棄，讓模型直接重建每個 patch 的位置，會導致上下游的 discrepency。因為下游任務需要 PE，而上游預訓練的模型又完全沒見過 PE。

ViT 對于 long-range 的建模能力很強，這個簡單的位置重建任務可能沒辦法讓模型學到非常 high-level 的語義特征。

看上去相似的不同 patch (例如純色的背景) 的位置無需被精準重建，因此決定哪些 patch 的位置需要被重建非常關鍵。

針對上述難點，我們提出了三個解決手段：

針對問題一，我們采用了一個簡單的隨機丟棄策略。每次訓練過程中丟棄 75% 的 PE，保留 25% 的 PE。

針對問題二，我們采取了高比例的 patch mask，既能提高代理任務的難度，又能加快訓練的速度。

針對問題三，我們提出了 position smoothing 和 attentive reconstruction 的策略。

3.1 DropPos 前向過程

算法 1. DropPos 的前向過程

DropPos 的前向過程包括兩段 mask，分別是第一步 patch mask (類似 MAE)，和第二步的 position mask (用可學習的 position mask 代替 dropped positions)。具體可以參見上方的偽代碼。

3.2 Objective

我們使用了一個最簡單的 cross-entropy loss 作為預訓練的目標函數：

其中，o 是模型的輸出，即第 i 個 patch 的預測位置是 j 的 logit，y 是真實的位置信息。

gamma 是第一步的 patch mask ratio，N 為總 patch 數量。

Mpose是 0-1 的 position mask，1 表示該 patch 帶有 PE，不應當被重建，而 0 表示該 patch 不帶 PE，需要被重建。

我們接下來引入 position smoothing 和 attentive reconstruction 技術來松弛這個問題，以解決相似但不同 patch 的位置重建問題。

3.2.1 Position Smoothing

我們采用一個高斯核來平滑原本的 position targets

此處，w(i, j) 表示當真實位置為 i，而預測位置為 j 時，平滑后的 position target。

此外，我們還讓 sigma 自大變小，讓模型一開始不要過分關注精確的位置重建，而訓練后期則越來越關注于精準的位置重建。

3.2.2 Attentive Reconstruction

我們采用 [CLS] token 和其他 patch 的相似度作為親和力矩陣，作為目標函數的額外權重。

其中 f 為不同 token 的特征，tau 為超參數，控制了 affinity 的平滑程度。

4. Experiments

4.1 與其他方法的對比

4.2 消融實驗

本文主要有四個超參：patch mask ratio (gamma)，position mask ratio (gamma_pos)，sigma，和 tau。

由表，我們可以得出一些比較有趣的結論：

一般來說，更高的 position 重建精度會帶來更高的下游任務性能。

上述結論存在例外：當 sigma = 0 時，即不做位置平滑時，位置預測精度高，而下游任務表現反而低；當 tau = inf 時，即不做 attentive reconstruction 時，位置預測精度高，而下游表現反而低。

因此，過分關注于預測每一個 patch 的精確的位置，會導致局部最優，對于下游任務不利。

上圖是 DropPos 位置重建的可視化結果，黑色 patch 代表的是前向過程中被 mask 掉的 patch；白色 patch 的位置被錯誤重建，而剩余 patch 的位置被精準重建。

DropPos 在極端情況 (例如 gamma=0.75) 時，依然可以做到大部分 patch 的精準重建。