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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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