0. 基本信息

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

論文：arxiv.org/pdf/2309.03576

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

今天介紹我們在自監(jiān)督視覺預(yù)訓(xùn)練領(lǐng)域的一篇原創(chuàng)工作DropPos:Pre-Training Vision Transformers by Reconstructing Dropped Positions.

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

1. TL;DR

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

2. Motivation

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

表 1. MoCo v3 中的實(shí)驗(yàn)現(xiàn)象 (原文第 6 頁)

這一實(shí)驗(yàn)結(jié)果背后，隱含著「ViT 的建模主要關(guān)注于不同 patch 的 visual appearence，對于 position 的 awareness 較差」這一信息。即，如果把圖片切 patch 然后再隨機(jī)打亂之后，ViT 能夠在亂序的情況下準(zhǔn)確識別該圖片的類別。這一點(diǎn)和人類直覺有很大出入。同時(shí)，有可能是因?yàn)?ViT 過擬合到了 ImageNet 這個(gè)特定數(shù)據(jù)集導(dǎo)致的。

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

表 2. Position awareness 對于下游任務(wù)的影響

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

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

圖 1. DropPos 與 CL 和 MIM 的對比

DropPos 有如下的優(yōu)勢：

對比 CL，DropPos 不需要精心設(shè)計(jì)的數(shù)據(jù)增強(qiáng) (例如 multi-crop)。

對比 MIM，DropPos 不需要精心設(shè)計(jì)的掩碼策略和重建目標(biāo)。

下面我們介紹 DropPos 的具體運(yùn)行流程

3. Method

圖 2. DropPos 的流程圖

即使 DropPos 的想法很直觀也很簡單，但這類方法一直沒有成為預(yù)訓(xùn)練的主流，主要是由于在設(shè)計(jì)上有以下三個(gè)難點(diǎn)：

如果簡單地把所有 PE 丟棄，讓模型直接重建每個(gè) patch 的位置，會導(dǎo)致上下游的 discrepency。因?yàn)橄掠稳蝿?wù)需要 PE，而上游預(yù)訓(xùn)練的模型又完全沒見過 PE。

ViT 對于 long-range 的建模能力很強(qiáng)，這個(gè)簡單的位置重建任務(wù)可能沒辦法讓模型學(xué)到非常 high-level 的語義特征。

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

針對上述難點(diǎn)，我們提出了三個(gè)解決手段：

針對問題一，我們采用了一個(gè)簡單的隨機(jī)丟棄策略。每次訓(xùn)練過程中丟棄 75% 的 PE，保留 25% 的 PE。

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

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

3.1 DropPos 前向過程

算法 1. DropPos 的前向過程

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

3.2 Objective

我們使用了一個(gè)最簡單的 cross-entropy loss 作為預(yù)訓(xùn)練的目標(biāo)函數(shù)：

其中，o 是模型的輸出，即第 i 個(gè) patch 的預(yù)測位置是 j 的 logit，y 是真實(shí)的位置信息。

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

Mpose是 0-1 的 position mask，1 表示該 patch 帶有 PE，不應(yīng)當(dāng)被重建，而 0 表示該 patch 不帶 PE，需要被重建。

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

3.2.1 Position Smoothing

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

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

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

3.2.2 Attentive Reconstruction

我們采用 [CLS] token 和其他 patch 的相似度作為親和力矩陣，作為目標(biāo)函數(shù)的額外權(quán)重。

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

4. Experiments

4.1 與其他方法的對比

4.2 消融實(shí)驗(yàn)

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

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

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

上述結(jié)論存在例外：當(dāng) sigma = 0 時(shí)，即不做位置平滑時(shí)，位置預(yù)測精度高，而下游任務(wù)表現(xiàn)反而低；當(dāng) tau = inf 時(shí)，即不做 attentive reconstruction 時(shí)，位置預(yù)測精度高，而下游表現(xiàn)反而低。

因此，過分關(guān)注于預(yù)測每一個(gè) patch 的精確的位置，會導(dǎo)致局部最優(yōu)，對于下游任務(wù)不利。

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

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