基于生成模型的預(yù)訓(xùn)練方法

I實(shí)驗(yàn)

總結(jié)

參考

前言

請(qǐng)?zhí)砑訄D片描述

我們這次要介紹的文章被接收在 ICCV 2023 上，題為：DreamTeacher: Pretraining Image Backbones with Deep Generative Models，我認(rèn)為是個(gè)挺強(qiáng)挺有趣的自監(jiān)督方面的工作。DreamTeacher 用于從預(yù)訓(xùn)練的生成網(wǎng)絡(luò)向目標(biāo)圖像 Backbone 進(jìn)行知識(shí)蒸餾，作為一種通用的預(yù)訓(xùn)練機(jī)制，不需要標(biāo)簽。這篇文章中研究了特征蒸餾，并在可能有任務(wù)特定標(biāo)簽的情況下進(jìn)行標(biāo)簽蒸餾，我們會(huì)在后文詳細(xì)介紹這兩種類型的知識(shí)蒸餾。

事實(shí)上，之前已經(jīng)在 GiantPandaCV 上介紹過(guò)一種 diffusion 去噪自監(jiān)督預(yù)訓(xùn)練方法：DDeP，DDeP 的設(shè)計(jì)簡(jiǎn)單，但去噪預(yù)訓(xùn)練的方法很古老了。然而，DreamTeacher 開(kāi)創(chuàng)了如何有效使用優(yōu)質(zhì)的生成式模型蒸餾獲得相應(yīng)的知識(shí)。

補(bǔ)充：在 DDeP 這篇文章中，經(jīng)過(guò)讀者糾正，我們重新表述了加噪公式：

相關(guān)工作

Discriminative Representation Learning

最近比較流行的處理方法是對(duì)比表示學(xué)習(xí)方法，SimCLR 是第一個(gè)在線性探測(cè)和遷移學(xué)習(xí)方面表現(xiàn)出色的方法，而且沒(méi)有使用類標(biāo)簽，相較于監(jiān)督預(yù)訓(xùn)練方法。隨后的工作，如 MoCo，通過(guò)引入 memory bank 和梯度停止改進(jìn)了孿生網(wǎng)絡(luò)設(shè)計(jì)。然而，這些方法依賴于大量的數(shù)據(jù)增強(qiáng)和啟發(fā)式方法來(lái)選擇負(fù)例，可能不太適用于像 ImageNet 這樣規(guī)模的數(shù)據(jù)集。關(guān)于 memory bank 的概念，memory bank 是 MoCo 中的一個(gè)重要組件，用于存儲(chǔ)模型的特征向量。在 MoCo 的訓(xùn)練過(guò)程中，首先對(duì)一批未標(biāo)記的圖像進(jìn)行前向傳播，得到每個(gè)圖像的特征向量。然后，這些特征向量將被存儲(chǔ)到內(nèi)存庫(kù)中。內(nèi)存庫(kù)的大小通常會(huì)比較大，足夠存儲(chǔ)許多圖像的特征。訓(xùn)練過(guò)程的關(guān)鍵部分是建立正負(fù)樣本對(duì)。對(duì)于每個(gè)樣本，其特征向量將被視為查詢向量（Query），而來(lái)自內(nèi)存庫(kù)的其他特征向量將被視為候選向量（Candidate）。通常情況下，查詢向量和候選向量來(lái)自同一張圖片的不同視角或數(shù)據(jù)增強(qiáng)的版本。然后，通過(guò)比較查詢向量與候選向量之間的相似性來(lái)構(gòu)建正負(fù)樣本對(duì)。此外，還有一些其他方法和概念，我們就不在這篇解讀文章中介紹了。

Generative Representation Learning

DatasetGAN 是最早展示預(yù)訓(xùn)練 GAN 可以顯著改善感知任務(wù)表現(xiàn)的研究之一，特別是在數(shù)據(jù)標(biāo)記較少的情況下。SemanticGAN 提出了對(duì)圖像和標(biāo)簽的聯(lián)合建模。推理過(guò)程首先將測(cè)試圖像編碼為 StyleGAN 的潛在空間，然后使用任務(wù)頭部解碼標(biāo)簽。DDPM-seg 沿著這一研究方向，但使用了去噪擴(kuò)散概率模型（DDPMs）代替 StyleGAN。這篇文章繼續(xù)了這一研究方向，但重點(diǎn)放在從預(yù)訓(xùn)練的生成模型中，特別是擴(kuò)散模型，向下游圖像主干中提取知識(shí)，作為一種通用的預(yù)訓(xùn)練方式。

關(guān)于相關(guān)工作部分中涉及到的方法，如果有疑惑的推薦閱讀原文（鏈接在文末）。

DreamTeacher 框架介紹

DreamTeacher 框架能在兩種場(chǎng)景下的工作：無(wú)監(jiān)督表示學(xué)習(xí)和半監(jiān)督學(xué)習(xí)。在無(wú)監(jiān)督表示學(xué)習(xí)中，預(yù)訓(xùn)練階段沒(méi)有可用的標(biāo)簽信息，而在半監(jiān)督學(xué)習(xí)中，只有部分?jǐn)?shù)據(jù)擁有標(biāo)簽?？蚣苁褂糜?xùn)練好的生成模型 G 來(lái)傳遞其學(xué)到的表示知識(shí)到目標(biāo)圖像主干 f。無(wú)論在哪種場(chǎng)景下，框架的訓(xùn)練方法和所選的生成模型 G 與圖像主干 f 的選擇都是一樣的。首先，它創(chuàng)建一個(gè)包含圖像和相應(yīng)特征的特征數(shù)據(jù)集 。然后，通過(guò)將生成模型的特征傳遞到圖像主干 f 的中間特征中來(lái)訓(xùn)練圖像主干 f。作者特別關(guān)注使用卷積主干 f 的情況，而對(duì) Transformer 的探索留給未來(lái)的研究。

Unsupervised Representation Learning

對(duì)于無(wú)監(jiān)督表示學(xué)習(xí)，給定一個(gè)特征數(shù)據(jù)集 D，在圖像主干 f 的不同層次上附加特征回歸器，以回歸出對(duì)應(yīng)的生成特征 從圖像 中。我們首先如何討論創(chuàng)建特征數(shù)據(jù)集，然后設(shè)計(jì)特征回歸器，最后介紹蒸餾目標(biāo)。創(chuàng)建特征數(shù)據(jù)集 D 的方法有兩種。一種是通過(guò)從生成模型 G 中采樣圖像，并記錄生成過(guò)程中提取的中間特征來(lái)創(chuàng)建合成數(shù)據(jù)集。這種方法可以合成無(wú)限大小的數(shù)據(jù)集，但可能會(huì)出現(xiàn) mode dropping（生成模型可能沒(méi)有學(xué)習(xí)到分布的某些部分）的問(wèn)題。另一種方法是將實(shí)際圖像通過(guò)編碼過(guò)程編碼到生成模型 G 的潛在空間中，然后記錄生成過(guò)程中提取的中間特征，創(chuàng)建編碼數(shù)據(jù)集。合成數(shù)據(jù)集適用于采樣速度快、無(wú)法編碼真實(shí)圖像的生成模型（如 GAN），而編碼數(shù)據(jù)集適用于具有編碼器網(wǎng)絡(luò)的生成模型（如 VAE）和擴(kuò)散模型。這兩種方法的特征數(shù)據(jù)集可以在離線預(yù)先計(jì)算，也可以在訓(xùn)練過(guò)程中在線創(chuàng)建，以實(shí)現(xiàn)快速的內(nèi)存訪問(wèn)和高效的樣本生成和刪除，從而適用于任何大小的數(shù)據(jù)集和特征預(yù)訓(xùn)練，同時(shí)增加下游Backbone 網(wǎng)絡(luò)的魯棒性。DreamTeacher 框架的整體流程如下圖所示，圖里表示創(chuàng)建特征數(shù)據(jù)集 D 使用的是第二種方法。為了將生成式表示 ? 蒸餾到通用主干 f 中，設(shè)計(jì)了一個(gè)特征回歸器模塊，將圖像主干的多層特征映射并對(duì)齊到生成式特征上。受到 Feature Pyramid Network（FPN）的設(shè)計(jì)啟發(fā)，特征回歸器采用自頂向下的架構(gòu)，并使用側(cè)向跳線連接來(lái)融合主干特征，并輸出多尺度特征。在圖像主干的最后一層之前應(yīng)用了類似于 PSPNet 中的金字塔池化模塊（PPM），上圖（底部）直觀地描述了這個(gè)架構(gòu)。接下來(lái)，我們關(guān)注如何做特征蒸餾的。將編碼器 f 的不同級(jí)別的中間特征表示為 ，對(duì)應(yīng)的特征回歸器輸出為 。使用一個(gè) 1×1 的卷積來(lái)匹配 和 的通道數(shù)，如果它們不同的話。特征回歸損失非常簡(jiǎn)單，受到 FitNet 的啟發(fā)，它提出了通過(guò)模擬中間特征激活將教師網(wǎng)絡(luò)上的知識(shí)蒸餾到學(xué)生網(wǎng)絡(luò)上：

在這里，W 是一個(gè)不可學(xué)習(xí)的白化算子，使用 LayerNorm 實(shí)現(xiàn)，用于對(duì)不同層次上的特征幅值進(jìn)行歸一化。層數(shù) l = {2, 3, 4, 5}，對(duì)應(yīng)于相對(duì)于輸入分辨率的 步長(zhǎng)處的特征。

此外，這篇文章還探索了基于激活的注意力轉(zhuǎn)移（AT）目標(biāo)。AT 使用一個(gè)運(yùn)算符 ，對(duì)空間特征的每個(gè)維度生成一個(gè)一維的“注意力圖”，其中 |Ai| 表示特征激活 A 在通道維度 C 上的絕對(duì)值和。這種方法相比直接回歸高維特征可以提高收斂速度。具體來(lái)說(shuō)，AT 損失函數(shù)如下：

其中分別是回歸器和生成模型在第 l 層中的特征的矢量形式中的第 j 對(duì)。

最后，綜合特征回歸損失為：

Label-Guided Representation Learning

在這里插入圖片描述

在半監(jiān)督設(shè)置中，在預(yù)訓(xùn)練階段在凍結(jié)的生成網(wǎng)絡(luò) G 之上訓(xùn)練了一個(gè)任務(wù)相關(guān)的分支，稱為特征解釋器，采用 DatasetGAN 的方法進(jìn)行監(jiān)督訓(xùn)練。與 DatasetGAN 合成用于訓(xùn)練下游任務(wù)網(wǎng)絡(luò)的帶標(biāo)簽數(shù)據(jù)集不同，DreamTeacher 改用軟標(biāo)簽蒸餾，即在編碼和合成的數(shù)據(jù)集中都包含了預(yù)測(cè)的軟標(biāo)簽，也就是特征數(shù)據(jù)集 D 中包含了軟標(biāo)簽。這在上圖（d）中進(jìn)行了可視化。

這篇文章探索了使用分割標(biāo)簽對(duì)解釋器分支進(jìn)行訓(xùn)練（半監(jiān)督情景下），并使用交叉熵和 Dice 目標(biāo)的組合來(lái)訓(xùn)練：

其中是特征解釋器的權(quán)重，y 是任務(wù)標(biāo)簽。H(·, ·) 表示像素級(jí)的交叉熵?fù)p失，D(·, ·) 表示 Dice Loss。

對(duì)于標(biāo)簽蒸餾，使用以下?lián)p失函數(shù)：

其中 和 分別是特征解釋器和目標(biāo)圖像主干 f 的 logits。H 是交叉熵?fù)p失，而 τ 是溫度參數(shù)。

將標(biāo)簽蒸餾目標(biāo)與特征蒸餾目標(biāo)相結(jié)合，得到混合損失函數(shù)：

使用混合蒸餾損失對(duì)預(yù)訓(xùn)練數(shù)據(jù)集中的所有圖像進(jìn)行預(yù)訓(xùn)練，無(wú)論是帶標(biāo)簽還是無(wú)標(biāo)簽的。帶標(biāo)簽的標(biāo)簽僅用于訓(xùn)練特征解釋器，而 DreamTeacher 只使用特征解釋器生成的軟標(biāo)簽對(duì)圖像主干 f 進(jìn)行蒸餾預(yù)訓(xùn)練。

實(shí)驗(yàn)

在實(shí)驗(yàn)中，使用的生成模型包含：unconditional BigGAN、ICGAN、StyleGAN2；對(duì)于基于擴(kuò)散的模型，使用了 ADM 和 stable diffusion 模型。使用的數(shù)據(jù)集包含：bdd100k、ImageNet-1k(IN1k-1M)、LSUN 和 ffhq。下表將 DreamTeacher 與 ImageNet 和 COCO 上的自監(jiān)督學(xué)習(xí)的 SOTA 方法進(jìn)行比較：

在這里插入圖片描述

對(duì)于 Label-efficient 的語(yǔ)義分割 benchmark。下表將 DreamTeacher與各種表示學(xué)習(xí)基線進(jìn)行比較。

下圖是使用 DreamTeacher 特征蒸餾預(yù)訓(xùn)練的 ConvNX-B 模型在 LSUN-cat 無(wú)標(biāo)簽圖像上的定性結(jié)果。

在這里插入圖片描述

總結(jié)

這篇文章的研究聚焦于提出一種名為 DreamTeacher 的框架，旨在從生成模型向目標(biāo)圖像 Backbone 傳遞知識(shí)（知識(shí)蒸餾）。在這個(gè)框架下，進(jìn)行了多個(gè)實(shí)驗(yàn)，涵蓋了不同的 settings ，包括生成模型、目標(biāo)圖像 Backbone 和評(píng)估 benchmark。其目標(biāo)是探究生成式模型在大規(guī)模無(wú)標(biāo)簽數(shù)據(jù)集上學(xué)習(xí)語(yǔ)義上有意義特征的能力，并將這些特征成功地傳遞到目標(biāo)圖像 Backbone 上。

通過(guò)實(shí)驗(yàn)，這篇文章發(fā)現(xiàn)使用生成目標(biāo)的生成網(wǎng)絡(luò)能夠?qū)W習(xí)到具有意義的特征，這些特征可以有效地應(yīng)用于目標(biāo)圖像主干。與現(xiàn)有自監(jiān)督學(xué)習(xí)方法相比，這篇文章基于生成模型的預(yù)訓(xùn)練方法表現(xiàn)更為優(yōu)異，這些 benchmark 測(cè)試包括 COCO、ADE20K 和 BDD100K 等。

這篇文章的工作為生成式預(yù)訓(xùn)練提供了新的視角和方法，并在視覺(jué)任務(wù)中充分利用了生成模型。在近兩年的論文中，生成式預(yù)訓(xùn)練技術(shù)是一個(gè)比較有趣的方向。

責(zé)任編輯：彭菁