擴(kuò)散模型如何突破瓶頸？成本高又難訓(xùn)練的DiT/SiT模型如何提升效率？

對(duì)于這個(gè)問(wèn)題，紐約大學(xué)謝賽寧團(tuán)隊(duì)最近發(fā)表的一篇論文找到了一個(gè)全新的切入點(diǎn)：提升表征（representation）的質(zhì)量。

論文的核心或許就可以用一句話(huà)概括：「表征很重要！」

用謝賽寧的話(huà)來(lái)說(shuō)，即使只是想讓生成模型重建出好看的圖像，仍然需要先學(xué)習(xí)強(qiáng)大的表征，然后再去渲染高頻的、使圖像看起來(lái)更美觀的細(xì)節(jié)。

這個(gè)觀點(diǎn)，Yann LeCun之前也多次強(qiáng)調(diào)過(guò)。

有網(wǎng)友還在線(xiàn)幫謝賽寧想標(biāo)題：你這篇論文不如就叫「Representation is all you need」（手動(dòng)狗頭）

由于觀點(diǎn)一致，這篇研究也獲得了同在紐約大學(xué)的Yann LeCun的轉(zhuǎn)發(fā)。

當(dāng)使用自監(jiān)督學(xué)習(xí)訓(xùn)練視覺(jué)編碼器時(shí)，我們知道一個(gè)事實(shí)，使用具有重建損失（reconstruction loss）的解碼器的效果遠(yuǎn)遠(yuǎn)不如具有特征預(yù)測(cè)損失（feature prediction loss）和崩潰預(yù)防機(jī)制的聯(lián)合嵌入架構(gòu)。 

這篇來(lái)自紐約大學(xué)謝賽寧團(tuán)隊(duì)的論文表明，即使只對(duì)生成像素感興趣（例如，使用擴(kuò)散Transformer生成漂亮的圖片），包含特征預(yù)測(cè)損失也是值得的，以便解碼器的內(nèi)部表示可以基于預(yù)訓(xùn)練的視覺(jué)編碼器（例如 DINOv2）進(jìn)行特征預(yù)測(cè)。

REPA的核心思想非常簡(jiǎn)單，就是讓擴(kuò)散模型中的表征與外部更強(qiáng)大的視覺(jué)表征進(jìn)行對(duì)齊，但提升效果非常顯著，頗有「他山之石，可以攻玉」的意味。

僅僅是在損失函數(shù)添加一項(xiàng)相似度最大化，就能將SiT/DiT的訓(xùn)練速度提升將近18倍，還刷新了模型的SOTA性能，在ImageNet 256x256上實(shí)現(xiàn)了最先進(jìn)的FID=1.42。

謝賽寧表示，剛看到實(shí)驗(yàn)結(jié)果時(shí)，他自己也被震驚到了，因?yàn)楦杏X(jué)并沒(méi)有發(fā)明什么全新的東西，而只是意識(shí)到了，我們幾乎完全不理解擴(kuò)散模型和SSL方法學(xué)習(xí)到的表示。

論文簡(jiǎn)介

論文地址：https://arxiv.org/abs/2410.06940

項(xiàng)目地址：https://sihyun.me/REPA/

在生成高維的視覺(jué)數(shù)據(jù)方面，基于去噪方法（如擴(kuò)散模型）或基于流的生成模型，已經(jīng)成為了一種可擴(kuò)展的途徑，并在有挑戰(zhàn)性的的零樣本文生圖/文生視頻任務(wù)上取得了非常成功的結(jié)果。

最近的研究表明，生成擴(kuò)散模型中的去噪過(guò)程可以在模型內(nèi)部的隱藏狀態(tài)中引入有意義的表示，但這些表示的質(zhì)量目前仍落后于自監(jiān)督學(xué)習(xí)方法，例如DINOv2。

作者認(rèn)為，訓(xùn)練大規(guī)模擴(kuò)散模型的一個(gè)主要瓶頸，就在于無(wú)法有效學(xué)習(xí)到高質(zhì)量的內(nèi)部表示。

如果能夠結(jié)合高質(zhì)量的外部視覺(jué)表示，而不是僅僅依靠擴(kuò)散模型來(lái)獨(dú)立學(xué)習(xí)，就可以使訓(xùn)練過(guò)程變得更容易。

為了實(shí)現(xiàn)這一點(diǎn)，論文基于經(jīng)典的擴(kuò)散Transformer架構(gòu)，引入了一種簡(jiǎn)單的正則化方法REPA（REPresentation Alignment）。

簡(jiǎn)單來(lái)說(shuō)，就是將去噪網(wǎng)絡(luò)中從噪聲輸入 得到的隱藏狀態(tài)??的投影，與外部自監(jiān)督預(yù)訓(xùn)練的視覺(jué)編碼器從干凈圖像??獲得的視覺(jué)表示??_*進(jìn)行對(duì)齊。

這樣一個(gè)非常直給的策略，卻獲得了驚人的結(jié)果：應(yīng)用于流行的SiT或DiT時(shí)，模型的訓(xùn)練效率和生成質(zhì)量都得到了顯著提高。

具體來(lái)說(shuō)，REPA可以將SiT的訓(xùn)練速度加快17.5×以上，以不到40萬(wàn)步的訓(xùn)練量匹配有700萬(wàn)步訓(xùn)練的SiT-XL模型的性能，同時(shí)實(shí)現(xiàn)了FID=1.42的SOTA結(jié)果。

REPA：使用表征對(duì)齊的正則化

統(tǒng)一視角的擴(kuò)散模型+流模型

由于論文希望同時(shí)優(yōu)化基于流的模型SiT和基于去噪的擴(kuò)散模型DiT，因此首先從統(tǒng)一的隨機(jī)插值視角，對(duì)這兩種模型進(jìn)行簡(jiǎn)要的回顧。

考慮在t∈[0,T]的連續(xù)時(shí)間步中，對(duì)數(shù)據(jù)??_*~p(??)使用高斯分布ε～??(0,??)添加隨機(jī)噪音：

其中，α_t和σ_t分別表示t的遞減和遞增函數(shù)。在公式(1)給定的過(guò)程中，存在一個(gè)帶有速度場(chǎng)（velocity field）的概率流常微分方程：

其中t步時(shí)的分布就等于邊際概率p_t(??)。

速度??(??,t)可以表示為如下兩個(gè)條件期望之和：

這個(gè)值可以通過(guò)最小化如下訓(xùn)練目標(biāo)得到近似值??_θ(??,t)：

同時(shí)，還存在一個(gè)反向的隨機(jī)微分方程（SDE），帶有擴(kuò)散系數(shù)w_t，其中的邊際概率p_t(??)與公式(2)相符：

其中，??(??_t,t)是一個(gè)條件期望值，定義為：

對(duì)任意t>0，都可以通過(guò)速度??(??,t)計(jì)算出??(??,t)的值：

這表明，數(shù)據(jù)??_t也可以通過(guò)求解公式(5)的SDE來(lái)以另一種方式生成。

以上定義對(duì)類(lèi)似的擴(kuò)散模型變體，例如DDPM，同樣適用，只是需要將連續(xù)的時(shí)間步離散化。

方法概述

令p(??)為數(shù)據(jù)??∈??的未知目標(biāo)分布，我們的訓(xùn)練目標(biāo)就是通過(guò)模型對(duì)數(shù)據(jù)的學(xué)習(xí)得到p(??)的近似。

為了降低計(jì)算成本，最近流行的「潛在擴(kuò)散」方法（latent diffusion）提出學(xué)習(xí)潛在變量??=E(??)的分布p(??)，其中E表示來(lái)自預(yù)訓(xùn)練自編碼器（例如KL-VAE）中的編碼部分。

要學(xué)習(xí)到分布p(??)，就需要訓(xùn)練擴(kuò)散模型??_θ(??_t,t)，訓(xùn)練目標(biāo)是進(jìn)行速度預(yù)測(cè)，具體方法如上一節(jié)所述。

放在自監(jiān)督表示學(xué)習(xí)的背景中，可以將擴(kuò)散模型看成編碼器f_θ：?????和解碼器g_θ：?????的組合，其中編碼器負(fù)責(zé)隱式地學(xué)習(xí)到表示??_t以重建目標(biāo)??_t。

然而，作者提出，用于生成的大型擴(kuò)散模型并不擅長(zhǎng)表征學(xué)習(xí)，因此REPA引入了外部的語(yǔ)義豐富的表示，從而顯著提升生成性能。

REPA方法概述

模型觀察

擴(kuò)散模型是否真的不擅長(zhǎng)表征學(xué)習(xí)？這需要更進(jìn)一步地觀察模型才能確定，為此，研究人員測(cè)量并比對(duì)了diffusion transformer和當(dāng)前的SOTA自監(jiān)督模型DINOv2之間的表征差距，包括語(yǔ)義差距和特征對(duì)齊兩種角度。

語(yǔ)義差距

從圖2a可知，預(yù)訓(xùn)練SiT的隱藏層表示在第20層達(dá)到最佳狀態(tài)，這與之前的研究結(jié)果相符，但仍遠(yuǎn)遠(yuǎn)落后于DINOv2。

特征對(duì)齊

如圖2b和2c所示，使用CKNNA值測(cè)量SiT和DINOv2之間的表征對(duì)齊程度后發(fā)現(xiàn)，SiT的對(duì)齊效果會(huì)隨著模型增大和訓(xùn)練迭代步數(shù)增加而逐漸改善，但即使增加到7M次迭代，和DINOv2之間的對(duì)齊程度仍然不足。

事實(shí)上，這種差距不僅在SiT中存在，根據(jù)附錄C.2的實(shí)驗(yàn)結(jié)果，DiT等其他基于去噪的生成式Transformer模型也存在類(lèi)似的問(wèn)題。

縮小表征差距

那么，REPA方法究竟如何縮小這種表征差距，讓diffusion transformer在噪聲輸入中也能學(xué)到有用的語(yǔ)義特征？

定義N,D分別表示patch數(shù)量預(yù)訓(xùn)練編碼器f的嵌入維度，編碼器輸入為無(wú)噪聲的圖像??_*，輸出為??_*=f(??_*)∈?^N×D。

Diffusion transformer將編碼器輸出??_t=f_θ(??_t)通過(guò)一個(gè)可訓(xùn)練的投影頭h_φ（MLP）投影為h_φ(??_t)∈?^N×D。

之后，REPA負(fù)責(zé)將h_φ(??_t)與??_*進(jìn)行對(duì)齊，通過(guò)最大化兩者間的patch間相似度：

在實(shí)際實(shí)現(xiàn)中，將這一項(xiàng)添加到公式(4)定義的基于擴(kuò)散的訓(xùn)練目標(biāo)中，就得到總體的訓(xùn)練目標(biāo)：

其中超參數(shù)λ>0用于控制模型在去噪目標(biāo)和表征對(duì)齊間的權(quán)衡。

從圖3結(jié)果可知，REPA減少了表示中的語(yǔ)義差距。

有趣的是，使用REPA后，僅對(duì)齊前幾個(gè)Transformer塊就能實(shí)現(xiàn)足夠程度的表示對(duì)齊，從而讓diffusion transformer的靠后層專(zhuān)注于捕獲高頻細(xì)節(jié)，從而進(jìn)一步提高生成性能。

實(shí)驗(yàn)結(jié)果

為了驗(yàn)證REPA方法的有效性，實(shí)驗(yàn)在兩種流行的擴(kuò)散模型訓(xùn)練目標(biāo)（即??_velocity）上進(jìn)行了實(shí)驗(yàn)，包括DiT中改進(jìn)后的DDPM和SiT中的線(xiàn)性隨機(jī)插值，但實(shí)際中也同樣可以考慮其他的訓(xùn)練目標(biāo)。

所用模型默認(rèn)嚴(yán)格遵循SiT和DiT的原始結(jié)構(gòu)（除非有特別說(shuō)明），包括B/2、L/2、XL/2三種參數(shù)設(shè)置，如表1所示。

以下實(shí)驗(yàn)旨在回答3個(gè)問(wèn)題：

- REPA能否顯著提升diffusion transformer的訓(xùn)練？

- REPA在模型規(guī)模和表征質(zhì)量方面是否具有可擴(kuò)展性？

- 擴(kuò)散模型的表征能否和多種視覺(jué)表征進(jìn)行對(duì)齊？

REPA提升視覺(jué)縮放

首先比較兩個(gè)SiT-XL/2模型在前400K次迭代期間生成的圖像，它們共享相同的噪聲、采樣器和采樣步數(shù)，但其中使用REPA訓(xùn)練的模型顯示出更好的進(jìn)展。

REPA在各個(gè)方面都展現(xiàn)出了強(qiáng)大的可擴(kuò)展性

研究人員還改變了預(yù)訓(xùn)練編碼器和Diffusion Transformer的模型大小來(lái)檢驗(yàn)REPA的可擴(kuò)展性。

圖5a結(jié)果表明，與更好的視覺(jué)表示相結(jié)合可以改善生成效果和線(xiàn)性探測(cè)的結(jié)果。

此外，如圖5b和c所示，增加模型大小可以在生成和線(xiàn)性評(píng)估方面帶來(lái)更快的收益，也就是說(shuō)，模型規(guī)模越大，REPA的加速效果越明顯，表現(xiàn)出了強(qiáng)大的可擴(kuò)展性。

REPA顯著提高訓(xùn)練效率和生成質(zhì)量

最后，論文比較了普通DiT或SiT模型在訓(xùn)練中使用REPA前后的FID值。

在沒(méi)有指導(dǎo)的情況下，REPA在400K次迭代時(shí)實(shí)現(xiàn)了FID=7.9，優(yōu)于普通模型在7M次迭代后的性能。

此外，使用無(wú)分類(lèi)器引導(dǎo)時(shí)，帶有REPA的SiT-XL/2的性能優(yōu)于SOTA性能（FID=1.42），同時(shí)迭代次數(shù)減少了7倍。

作者介紹

Sihyun Yu

本文一作Sihyun Yu是KAIST（韓國(guó)科學(xué)技術(shù)院）人工智能專(zhuān)業(yè)最后一年的博士生，此前他同樣在KAIST獲得了數(shù)學(xué)和計(jì)算機(jī)科學(xué)的雙專(zhuān)業(yè)學(xué)士學(xué)位。

他的研究主要集中在減少大型生成模型訓(xùn)練（和采樣）的內(nèi)存和計(jì)算負(fù)擔(dān)，其中，對(duì)大規(guī)模且高效的視頻生成特別感興趣；博士期間，他還曾在英偉達(dá)和谷歌研究院擔(dān)任實(shí)習(xí)生。