何愷明謝賽寧解剖擴(kuò)散模型，新作剛剛出爐

衡宇 2024-01-26 15:32:34 來源：量子位

編輯部 發(fā)自 凹非寺

量子位 | 公眾號 QbitAI

CV大神何愷明，也來搞擴(kuò)散模型（Diffusion Model）了！

大神最新論文剛剛掛上arXiv，還是熱乎的：解構(gòu)擴(kuò)散模型，提出一個高度簡化的新架構(gòu)l-DAE（小寫的L）。

并且通過與何愷明在視覺自監(jiān)督學(xué)習(xí)領(lǐng)域的代表作MAE（Masked Autoencoder）對比，更好地理解了擴(kuò)散模型內(nèi)部的工作原理。

不僅如此，這篇論文還發(fā)現(xiàn)在擴(kuò)散模型中，去噪過程比擴(kuò)散過程更重要。

這項工作陣容非常豪華，不僅有何愷明坐鎮(zhèn)，合著作者中還有紐約大學(xué)計算機(jī)科學(xué)助理教授、CV大牛謝賽寧。

以及曾和他共同發(fā)表ConvNeXT工作的劉壯——他同時是DenseNet的共同一作。

給擴(kuò)散模型開刀

團(tuán)隊認(rèn)為，盡管去噪擴(kuò)散模型在生成任務(wù)上表現(xiàn)出色，但它們在表示學(xué)習(xí)方面的能力尚未得到充分探索。

為此，他們找到一個新穎的研究方法：

希望通過解構(gòu)擴(kuò)散模型，將其逐步轉(zhuǎn)化為類似于MAE的架構(gòu)，以更深入地理解DDM（去噪擴(kuò)散模型）在自監(jiān)督學(xué)習(xí)中的表示學(xué)習(xí)能力。

先來一圖概括解構(gòu)過程：

（如果你也覺得這個圖很眼熟，沒錯，就是謝賽寧在代表作ConvNeXT中使用的同款。）

以使用了VQGAN tokenizer的DiT模型作為基線，接下來介紹團(tuán)隊如何一刀一刀把它改造成更簡化版本的。

• 移除類別條件

首先移除擴(kuò)散模型中基于類別標(biāo)簽的條件，很顯然，有條件標(biāo)簽與自監(jiān)督學(xué)習(xí)的目標(biāo)不符。

這一步驟顯著提高了線性探測（linear probing）的準(zhǔn)確率。

• 解構(gòu)VQGAN

接下來移除原本采用的VQGan感知損失和對抗損失，在自監(jiān)督學(xué)習(xí)中，這兩種損失函數(shù)都不是必須的。

感知損失涉及到對類別標(biāo)簽的依賴，這與自監(jiān)督學(xué)習(xí)的目標(biāo)不符；對抗損失涉及對生成過程的優(yōu)化，而不是直接學(xué)習(xí)數(shù)據(jù)的表示。

移除兩種損失函數(shù)后，模型表現(xiàn)受到影響下降，但下一步又一把拉了回來。

• 替換噪聲調(diào)度器

在原始擴(kuò)散模型中，噪聲調(diào)度通常遵循一個復(fù)雜的時間步序列。團(tuán)隊采用了一種更簡單的線性衰減策略，讓噪聲水平在訓(xùn)練過程中線性地從最大值衰減到零，而不是遵循復(fù)雜的非線性衰減路徑。

這種簡化的噪聲調(diào)度策略使得模型在訓(xùn)練過程中更多地關(guān)注于較干凈的數(shù)據(jù)，而不是在高度噪聲的數(shù)據(jù)上花費(fèi)太多計算資源。

• 簡化tokenizer

作者進(jìn)一步改造了擴(kuò)散模型中的tokenizer，這是將輸入圖像映射到潛在空間的關(guān)鍵組件。

他們比較了幾種不同的方法，包括卷積變分自編碼器（conv. VAE）、基于塊的變分自編碼器（patch-wise VAE）、基于塊的自編碼器（patch-wise AE）和基于塊的主成分分析（patch-wis PCA）。

最終發(fā)現(xiàn)，即使是簡單的PCA也能有效地工作。

通過逆PCA（inverse PCA）將輸入圖像投影到潛在空間，添加噪聲，然后再將噪聲圖像投影回圖像空間。這種方法允許模型直接在圖像上進(jìn)行操作，而不需要tokenizer。

• 改變預(yù)測目標(biāo)

與現(xiàn)代擴(kuò)散模型通常預(yù)測噪聲不同，團(tuán)隊讓模型預(yù)測干凈的數(shù)據(jù)，通過調(diào)整損失函數(shù)實(shí)現(xiàn)。

最后一步，作者又讓模型直接預(yù)測原始圖像，而不是經(jīng)過PCA編碼的潛空間。這涉及到在PCA空間中計算殘差，并在損失函數(shù)中對PCA重建誤差進(jìn)行加權(quán)。

最終，整個模型的工作流程就非常簡單了：

輸入是一張有噪聲的圖片，噪聲添加在PCA潛空間里。輸出是原始的干凈圖片。

經(jīng)過這一系列改造，團(tuán)隊發(fā)現(xiàn)：

• 模型的表示能力主要來源于去噪過程，而非擴(kuò)散過程。
• 低維潛在空間非常關(guān)鍵，允許模型學(xué)習(xí)到數(shù)據(jù)的壓縮表示。

并且提出只有很少的現(xiàn)代組件對學(xué)習(xí)良好的表示至關(guān)重要，而許多其他組件是非必要的。

經(jīng)過改造簡化后的模型稱為l-DAE，在自監(jiān)督學(xué)習(xí)任務(wù)中表現(xiàn)出競爭力，同時在結(jié)構(gòu)上更接近于傳統(tǒng)的DAE（去噪自編碼器）。

在實(shí)驗(yàn)中，MAE和l-DAE兩種自編碼器方法表現(xiàn)都超過了MoCo v3，在ViT-B（86m參數(shù)）上，l-DAE表現(xiàn)與MAE相當(dāng)，但在更大規(guī)模的ViT-L（304M）上還所有落后。

最后，團(tuán)隊在簡短的結(jié)論中提出：

希望我們的發(fā)現(xiàn)能夠重新激發(fā)對基于去噪方法在當(dāng)今自監(jiān)督學(xué)習(xí)研究領(lǐng)域的興趣。

兩位作者已離開Meta

論文一經(jīng)公布，幾位作者就激情當(dāng)起了自個兒的首批自來水（doge）。

謝賽寧在推特上轉(zhuǎn)發(fā)了別人對I-DAE的分享和討論，并表示：

越來越多研究表明，擴(kuò)散模型也可以作為效果優(yōu)秀的特征提取器來使用。

作者劉壯也在一旁開麥，表示擴(kuò)散模型這東西，不僅僅能做生成相關(guān)的工作：

還有開發(fā)者驚訝地發(fā)現(xiàn)，居然還有人使用主成分分析，可有一陣子不見了。

不過，大神們的工作剛放出來嘛，還有待大家多多傳播和具體感知。

除這項最新研究涉及了MAE外，何愷明在提出代表作MAE后，還在此基礎(chǔ)上還發(fā)表了一系列研究。

例如，提出了一個用來訓(xùn)練CLIP的快速、簡單且有效的方法FLIP（Fast Language-Image Pre-training），對CLIP架構(gòu)引入類似于MAE的思路。

只在模型的極簡結(jié)構(gòu)上，施加了一個簡單的mask，就讓新模型的速度快了3.7倍，同時，性能還可以做到不降反升。

此外，還提出了不對ViT引入分層設(shè)計，只用普通ViT就能搞定目標(biāo)檢測的ViTDet。

他們使用普通ViT作為骨干網(wǎng)絡(luò)，基于MAE方法進(jìn)行預(yù)訓(xùn)練，由此得到的ViTDet能與之前所有基于分層骨干網(wǎng)絡(luò)的先進(jìn)方法競爭。

其他學(xué)者同樣將MAE方法用在了眾多項目之中，MAE開山論文目前谷歌學(xué)術(shù)引用量達(dá)4174次。

而何愷明近期的一次公開露面，是在香港中文大學(xué)參加了一個講座，講道“自己做科研也會emo”。

同時，他在講座上回答了校友們提出的關(guān)于大模型、AI for Science等相關(guān)的諸多問題。

其中干貨滿滿，在現(xiàn)場的多半排了很久的長隊才一睹大神真容，不在現(xiàn)場的則像追劇似的搜羅網(wǎng)友發(fā)在網(wǎng)上的演講片段視頻。

一作陳鑫磊， 目前是Meta FAIR實(shí)驗(yàn)室的研究科學(xué)家，也是浙大校友，研究興趣集中于預(yù)訓(xùn)練，特別是有自監(jiān)督或是多模態(tài)視覺表示的預(yù)訓(xùn)練。

作者劉壯，同樣是Meta FAIR實(shí)驗(yàn)室的研究科學(xué)家，同時也是清華姚班校友，研究重點(diǎn)為神經(jīng)網(wǎng)絡(luò)。

至于另一位合作者謝賽寧這邊，他最近從Meta AI離職加入了紐約大學(xué)，不過依然還是LeCun的同事。

論文地址：https://arxiv.org/pdf/2401.14404.pdf

參考鏈接：
[1]https://twitter.com/sainingxie/status/1750741794080407893