BeiT v2 來襲 | BeiT升級，全面超越 MAE，實現(xiàn) Vision Transformer 微調(diào)自由！

Masked image modeling 通過恢復損壞的圖像塊在自監(jiān)督表示學習中展示了令人印象深刻的結(jié)果。然而，大多數(shù)方法仍然對 low-level 圖像像素進行操作，這阻礙了對表示模型的 high-level 語義的利用。

在這項研究中，作者建議使用語義豐富的 visual tokenizer 作為 Mask 預測的重建目標，為將 MIM 從像素級提升到語義級提供了一種系統(tǒng)的方法。

具體來說，引入向量量化知識蒸餾來訓練 tokenizer，它將連續(xù)的語義空間離散化為 compact codes。然后，通過預測 masked image patches 的原始 visual tokenizer 來預訓練 Vision Transformers 。此外，作者鼓勵模型將patches信息顯式聚合到全局圖像表示中，這有助于線性預測。

圖像分類和語義分割的實驗表明，本文的方法優(yōu)于所有比較的 MIM 方法。在 ImageNet-1K（224 大?。┥?，base-size BEIT V2 在微調(diào)時達到 85.5% 的 top-1 精度，在線性預測時達到 80.1% 的 top-1 精度。large-size BEIT V2 在 ImageNet-1K（224 大?。┪⒄{(diào)上獲得 87.3% 的 top-1 準確率，在 ADE20K 上獲得 56.7% 的 mIoU 用于語義分割。

Masked image modeling 在學習視覺表示方面顯示出令人印象深刻的結(jié)果，這極大地緩解了 Vision Transformer 的注釋饑餓問題。給定一張圖像，這些方法通常首先通過mask一些 patches 來破壞原始圖像，相關(guān)任務是恢復原始圖像。以開創(chuàng)性工作 BEiT 為例，每張圖像在預訓練期間都有2個視圖，即圖像 patches 和 visual tokenizer。原始圖像首先被tokenizer為離散 token。隨機采樣的圖像 patches 在被饋送到 Vision Transformer 之前被 Masked。預訓練的目標是根據(jù)損壞的圖像 patches 恢復原始 visual tokenizer。在預訓練視覺編碼器后，可以通過附加輕量級任務層直接在各種下游任務上微調(diào)模型。

在 mask-then-predict 框架下，與之前工作的主要區(qū)別在于重建目標，例如 visual tokenizer、原始像素和手工制作的 HOG 特征。然而，恢復low-level監(jiān)督往往會浪費建模能力來預訓練高頻細節(jié)和短程依賴關(guān)系。例如，當掩蓋戴在男人頭上的“帽子”時，更喜歡模型在給定整個上下文的情況下學習被掩蓋的“帽子”的高級概念，而不是在像素級細節(jié)上苦苦掙扎。相比之下，語言建模中的掩碼詞通常被認為具有比像素更多的語義。這促使通過在預訓練期間利用語義感知監(jiān)督來挖掘 MIM 的潛力。

在這項工作中，作者引入了一種自監(jiān)督的視覺表示模型 BEIT V2，旨在通過學習語義感知的 visual tokenizer 來改進 BEIT 預訓練。具體來說，提出了向量量化知識蒸餾（VQ-KD）算法來離散化語義空間。VQ-KD 編碼器首先根據(jù)可學習的 codebook 將輸入圖像轉(zhuǎn)換為離散token。然后****學習重建由教師模型編碼的語義特征，以離散token為條件。在訓練 VQ-KD 后，其編碼器用作 BEIT 預訓練的 visual tokenizer，其中離散 code 用作監(jiān)督信號。

此外，作者建議通過明確鼓勵 CLS token 聚合所有 patches 來預訓練全局圖像表示。該機制解決了 Masked image modeling 僅預訓練 patch-level 表示的問題。結(jié)果，在聚合全局表示的幫助下，線性預測的性能得到了提高。

在 ImageNet-1k 上對 base 和 big Vision Transformer 進行自監(jiān)督學習，并在多個下游任務上進行評估，例如圖像分類、線性預測和語義分割。如圖 1 所示，BEIT V2 在 ImageNet 微調(diào)上大大優(yōu)于以前的自監(jiān)督學習算法，例如，在 ViT-B/16 和 ViT-L/16 上都比 BEIT 提高了大約2個點。BEIT V2在 ImageNet 線性預測上優(yōu)于所有比較的 MIM 方法，同時在 ADE20k 上實現(xiàn)語義分割的巨大性能提升。

本研究的貢獻總結(jié)如下：

? 引入向量量化知識蒸餾，將Masked image modeling從像素級提升到語義級，以進行自監(jiān)督表示學習。

? 提出了一種Patch聚合策略，該策略在給定 Patch-Level Masked image modeling 的情況下強制執(zhí)行全局表示。

? 對下游任務進行了廣泛的實驗，例如 ImageNet 微調(diào)、線性預測和語義分割。實驗結(jié)果表明，BEIT V2顯著提高了模型大小、訓練步驟和下游任務的性能。

BEIT V2 繼承了用于Masked image modeling的 BEIT 框架。具體來說，給定輸入圖像，使用visual tokenizer將圖像標記為離散的visual token。然后屏蔽一部分圖像patch并將其輸入 Vision Transformer。預訓練任務是根據(jù)損壞的圖像恢復Mask的visual token。

使用Vision Transformer作為主干網(wǎng)絡來獲得圖像表示。給定輸入圖像 ，將圖像  reshape為  個patches ，其中  和  是patch-size。在實驗中，將每個 224×224 的圖像分割成一個 14×14 的圖像patch網(wǎng)格，其中每個patch為 16×16。然后將圖像patch  展平并線性投影到 Transformer 的輸入嵌入中。對于 N 個圖像塊，將編碼向量表示為 。

visual tokenizer 將圖像映射到一系列離散 tokens。具體來說，圖像 x 被標記為  ，其中vocabulary V（即visual codebook）包含 |V| 離散code。請注意，token的數(shù)量與圖像塊的數(shù)量相同。本文提出向量量化知識蒸餾（VQ-KD）來訓練visual tokenizer。如圖 2 所示，VQ-KD 在訓練過程中有兩個模塊，即visual tokenizer和****。visual tokenizer由Vision Transformer編碼器和量化器組成。分詞器首先將輸入圖像編碼為向量。接下來，向量量化器在codebook中查找每個patch表示  的最近鄰。讓  表示 codebook 嵌入。對于第 i 個圖像塊，其量化代碼由下式獲得：其中  歸一化用于codebook查找。上述距離相當于根據(jù)余弦相似度找code。在將圖像量化為visual token之后，將-normalized codebook embeddings 提供給****。****也是一個多層 Transformer 網(wǎng)絡。輸出向量  旨在重建教師模型的語義特征，例如 DINO 和 CLIP。讓  表示第 i 個圖像塊的教師模型的特征向量。最大化****輸出  和教師指導  之間的余弦相似度。因為量化過程（等式 1）是不可微分的。如圖 2 所示，為了將梯度反向傳播到編碼器，梯度直接從****輸入復制到編碼器輸出。直觀地說，量化器為每個編碼器輸出查找最近的code，因此codebook embeddings的梯度指示了編碼器的有用優(yōu)化方向。VQ-KD的訓練目標是：其中，sg[·]表示停止梯度算子，它是向前傳遞時的一個身份，而在反向傳遞過程中梯度為零，D表示用于tokenizer訓練的圖像數(shù)據(jù)。

矢量量化訓練的一個常見問題是 codebook 崩潰。換句話說，只使用了一小部分code。作者憑經(jīng)驗發(fā)現(xiàn)有幾種技術(shù)可用于提高 codebook 利用率。于等人建議對codebook查找應用降維和 l2 歸一化。等式 1 表明計算 l2 歸一化距離以找到最近的 code。此外，將查找空間的維度減少到 32 維。低維codebook嵌入在饋入****之前被映射回高維空間。此外，采用指數(shù)移動平均線來更新codebook嵌入，這在實驗中往往更加穩(wěn)定。

遵循BEIT中的MIM設(shè)置，以預訓練 Vision Transformers 的圖像表示。給定一個輸入圖像 x，分塊選擇大約40%的圖像 patch 進行 mask。如果 if ，其中δ（·）是指示函數(shù)，則使用共享的可學習嵌入來替換的原始圖像patch嵌入。隨后，在輸入前準備了一個可學習的CLS token，即，并將它們提供給vision Transformer。最終的編碼向量記為，其中為CLS token。接下來，使用一個mask圖像建模head，基于損壞的圖像來預測mask位置的 visual tokens。對于每個mask位置，一個softmax分類器預測 visual tokens ，其中是mask圖像，，是分類器的權(quán)重。visual tokens是由第2.2節(jié)中訓練的標記器獲得的，該節(jié)為mask圖像建模提供了監(jiān)督。最后，MIM的訓練損失可以表述為：其中表示原始圖像的visual tokens，D表示訓練前的圖像。

受(Gao和Callan，2021)的啟發(fā)，明確地預訓練了CLS token以進行全局表示。目標是減輕patch-level預訓練和 image-level 表示聚合之間的差異。如圖3所示，構(gòu)建了一個 representation bottleneck 來指導CLS token收集信息。對于L-layer Transformer，讓${h^l_i\}^N_{i=1}表示第層的輸出向量，其中l=1····l$。為了預先訓練最后一層的CLS token ，連接中間l層的patch向量${h^l_i\}^N_{i=1}，即S=[h^L_{CLS},h^l_1,···,h^l_N]$。然后，將S輸入一個淺層（例如，兩層）Transformer decoder，并進行mask預測。注意，這里也計算了第l層的MIM損失，如公式3所示。所以最終的訓練損失是兩項的總和，即，在第l層的原始損失，和淺層Transformer decoder的MIM損失。在實現(xiàn)中還共享2個頭的MIM Softmax權(quán)重。直觀地看，該模型傾向于將全局信息推到，因為該模型傾向于充分利用從第1層到第l+1層的參數(shù)，以減少額外的MIM損失。information-flow bottleneck鼓勵CLS token比未經(jīng)訓練的對應token獲得更可靠的全局表示。新添加的淺層****僅用于對CLS token進行預訓練，經(jīng)過預訓練后將被丟棄。

[1].BEIT V2: Masked Image Modeling with Vector-Quantized Visual Tokenizers

本文僅做學術(shù)分享，如有侵權(quán)，請聯(lián)系刪文。