文章鏈接：https://arxiv.org/pdf/2310.15111
項(xiàng)目鏈接：https://github.com/apple/ml-mdm

亮點(diǎn)直擊

• 提出了Matryoshka Diffusion Models (MDM)，通過聯(lián)合處理多個(gè)分辨率的輸入，避免了級聯(lián)或潛在擴(kuò)散方法的復(fù)雜性，并引入了Nested UNet架構(gòu)，嵌套不同尺度的特征和參數(shù)，提升了高分辨率生成的效果。
• 采用了多分辨率損失，顯著加速了高分辨率去噪的收斂速度，同時(shí)使用漸進(jìn)式訓(xùn)練策略，從低分辨率開始逐步引入高分辨率內(nèi)容，實(shí)現(xiàn)了訓(xùn)練成本和生成質(zhì)量的良好平衡。
• 在多個(gè)生成任務(wù)中表現(xiàn)優(yōu)異，包括類條件圖像生成、文本到圖像生成和視頻生成，尤其是在1024×1024像素分辨率下，使用較小的數(shù)據(jù)集（CC12M）依然取得了高質(zhì)量生成結(jié)果。
• MDM具備強(qiáng)大的泛化能力，不僅在圖像生成中表現(xiàn)突出，還自然擴(kuò)展到視頻生成，展現(xiàn)了廣泛的應(yīng)用潛力。

總結(jié)速覽

解決的問題：

• 生成高質(zhì)量圖像和視頻的擴(kuò)散模型面臨高維度學(xué)習(xí)的計(jì)算和優(yōu)化挑戰(zhàn)，尤其是在處理高分辨率圖像時(shí)。

提出的方案：

• 引入Matryoshka Diffusion Model (MDM)，一種新穎的高分辨率圖像和視頻合成框架，采用聯(lián)合降噪的擴(kuò)散過程，在多個(gè)分辨率下處理輸入。
• 使用NestedUNet架構(gòu)，在大尺度輸入的特征和參數(shù)中嵌套小尺度的特征和參數(shù)。
• 設(shè)計(jì)了一種從低分辨率到高分辨率的漸進(jìn)式訓(xùn)練方式，優(yōu)化高分辨率生成的效果。

應(yīng)用的技術(shù)：

• MDM通過多分辨率聯(lián)合降噪的擴(kuò)散過程，使用NestedUNet架構(gòu)，在像素空間直接進(jìn)行高分辨率生成。
• 使用漸進(jìn)式訓(xùn)練策略，逐步從低分辨率到高分辨率，解決了高分辨率生成的優(yōu)化問題。

達(dá)到的效果：

• MDM在多個(gè)基準(zhǔn)測試中表現(xiàn)出色，包括類條件圖像生成、高分辨率文本到圖像生成和文本到視頻生成任務(wù)。
• 成功訓(xùn)練了單一的像素空間模型，最高分辨率可達(dá)1024×1024像素。
• 在使用僅包含1200萬張圖像的CC12M數(shù)據(jù)集上展示了強(qiáng)大的零樣本泛化能力。

Matryoshka擴(kuò)散模型

本節(jié)介紹了Matryoshka擴(kuò)散模型（MDM），一個(gè)新型的擴(kuò)散模型類別，在高分辨率空間中進(jìn)行訓(xùn)練，同時(shí)利用數(shù)據(jù)生成的層次結(jié)構(gòu)。MDM首先在擴(kuò)展空間中推廣了標(biāo)準(zhǔn)擴(kuò)散模型，為此提出了專門的嵌套架構(gòu)和訓(xùn)練過程。

擴(kuò)展空間中的擴(kuò)散模型

與級聯(lián)或潛在方法不同，MDM通過在擴(kuò)展空間中引入多分辨率擴(kuò)散過程，學(xué)習(xí)一個(gè)具有層次結(jié)構(gòu)的單一擴(kuò)散過程。下圖2展示了該過程的示意圖。

NestedUNet架構(gòu)

除了相對于其他層次化方法的簡潔性，NestedUNet 還允許以最有效的方式分配計(jì)算。如下圖 3 所示，早期探索發(fā)現(xiàn)，當(dāng)將大部分參數(shù)和計(jì)算分配到最低分辨率時(shí)，MDM 在可擴(kuò)展性方面表現(xiàn)得更好。

訓(xùn)練

研究者們在多個(gè)分辨率上聯(lián)合使用正常的去噪目標(biāo)來訓(xùn)練 MDM，具體如下：

漸進(jìn)式訓(xùn)練

實(shí)驗(yàn)

MDM 是一種通用技術(shù)，適用于任何輸入維度可以逐步壓縮的問題?？紤]了兩個(gè)超出類條件圖像生成的應(yīng)用，展示了本文方法的有效性——文本到圖像和文本到視頻生成。

實(shí)驗(yàn)設(shè)置

數(shù)據(jù)集

在論文中廣泛依賴 CC12M 作為文本到圖像生成模型的數(shù)據(jù)集，顯著不同于以往依賴極大且有時(shí)無法獲取的數(shù)據(jù)集的研究。CC12M 足以構(gòu)建高質(zhì)量的文本到圖像模型，并具備強(qiáng)大的zero-shot 能力，訓(xùn)練時(shí)間相對較短。這使得社區(qū)能夠?qū)Ψ椒ㄟM(jìn)行更一致的比較，因?yàn)樵摂?shù)據(jù)集是免費(fèi)提供的，且訓(xùn)練時(shí)間是可行的。因此 CC12M 更適合作為該問題研究的共同訓(xùn)練和評估基線。

評估
根據(jù)以往的研究，使用 Fréchet Inception Distance（ImageNet, CC12M）和 CLIP 分?jǐn)?shù)（CC12M）來評估本文的圖像生成模型。為了檢查它們的zero-shot 能力，我們還報(bào)告了使用 COCO驗(yàn)證集生成圖像的 FID/CLIP 分?jǐn)?shù)，使用 CC12M 訓(xùn)練的模型。我們還在補(bǔ)充材料中提供了圖像和視頻合成的其他定性樣本。

實(shí)現(xiàn)細(xì)節(jié)

基線模型除了與現(xiàn)有最先進(jìn)方法的比較外，還在控制設(shè)置下對 MDM 與三個(gè)基線模型進(jìn)行了詳細(xì)分析：

1. 簡單 DM：將標(biāo)準(zhǔn) UNet 架構(gòu)直接應(yīng)用于高分辨率輸入；我們還考慮了 Nested UNet 架構(gòu)，但忽略低分辨率損失；這兩種情況本質(zhì)上與最近的端到端擴(kuò)散模型如 Hoogeboom 等（2023）相同。
2. 級聯(lián) DM：遵循 Saharia 等（2022）的實(shí)現(xiàn)細(xì)節(jié)，訓(xùn)練一個(gè)直接與 MDM 可比較的 CDM，其中上采樣器的配置與我們的 NestedUNet 相同。我們還對低分辨率條件圖像應(yīng)用噪聲增強(qiáng)，并在推理過程中遍歷最佳噪聲水平。
3. 潛在 DM：利用來自 Rombach 等（2022）自動編碼器的潛在編碼，隨后訓(xùn)練與 MDM UNet 維度匹配的擴(kuò)散模型。

主要結(jié)果

比較簡單 DM 和 MDM，可以清晰地看到 MDM 收斂速度更快，最終性能更好。這表明，多分辨率擴(kuò)散過程結(jié)合多分辨率損失有效地改善了模型的收斂性，同時(shí)帶來的復(fù)雜性微乎其微。當(dāng)遵循漸進(jìn)式訓(xùn)練計(jì)劃時(shí)，看到 MDM 的性能和收斂速度進(jìn)一步提高。作為直接比較，發(fā)現(xiàn)級聯(lián) DM 基線顯著低于 MDM 的表現(xiàn)，盡管兩者都從相同的64X64模型開始。需要注意的是，這一點(diǎn)非常顯著，因?yàn)榧壜?lián) DM 的參數(shù)總數(shù)大于 MDM（因?yàn)?MDM 在不同分辨率之間有廣泛的參數(shù)共享），且推理步驟是其兩倍。假設(shè)級聯(lián) DM 性能較差的主要原因在于64X64模型沒有經(jīng)過嚴(yán)格訓(xùn)練，這導(dǎo)致訓(xùn)練與推理之間在條件輸入方面存在較大差距。最后，與 LDM 相比，MDM 的性能也更好。盡管這不是一個(gè)直接的對比，因?yàn)?LDM 確實(shí)由于其小輸入尺寸而更有效，但 MDM 的訓(xùn)練和推理管道更為簡單。

與文獻(xiàn)的比較在下表 1 中，MDM 與現(xiàn)有文獻(xiàn)中的方法進(jìn)行了比較，報(bào)告了 ImageNet256X256的 FID-50K 和 MSCOCO 的zero-shot FID-30K。對于 ImageNet，我們的架構(gòu)和超參數(shù)并沒有經(jīng)過優(yōu)化，MDM 能夠達(dá)到 3.51 的競爭性 FID 值，與 CFG 相比。我們的 FID 結(jié)果與文獻(xiàn)相當(dāng)，而 MDM 在訓(xùn)練時(shí)使用的數(shù)據(jù)量明顯少于基線模型，如 Imagen 和 DALL·E 2。

定性結(jié)果下面展示訓(xùn)練后的 MDM 隨機(jī)樣本，用于圖像生成（ImageNet 256X256 ，下圖 5）、文本到圖像生成（CC12M 1024X1024， 下圖 6）和文本到視頻生成（WebVid-10M，下圖 7）。盡管在相對較小的數(shù)據(jù)集上進(jìn)行訓(xùn)練，MDM 展現(xiàn)出強(qiáng)大的zero-shot 能力，能夠生成高分辨率的圖像和視頻。值得注意的是，對所有三個(gè)任務(wù)使用相同的訓(xùn)練流程，表明其處理各種數(shù)據(jù)類型的多樣化能力。

消融研究

漸進(jìn)式訓(xùn)練的效果實(shí)驗(yàn)了漸進(jìn)式訓(xùn)練計(jì)劃，在該計(jì)劃中，改變了低分辨率模型在繼續(xù)訓(xùn)練目標(biāo)分辨率之前的訓(xùn)練迭代次數(shù)（下圖 8a）?？吹礁嗟牡头直媛视?xùn)練明顯有利于高分辨率的 FID 曲線。需要注意的是，在低分辨率輸入上進(jìn)行訓(xùn)練在內(nèi)存和時(shí)間復(fù)雜度方面更為高效，因此漸進(jìn)式訓(xùn)練為在訓(xùn)練過程中尋找最佳計(jì)算權(quán)衡提供了一種直接的選擇。

嵌套層數(shù)的效果接下來，比較了在 CC12M 上使用不同數(shù)量嵌套分辨率的性能。結(jié)果如上圖 8b 所示。從兩個(gè)分辨率層增加到三個(gè)分辨率層始終改善了模型的收斂性。值得注意的是，增加嵌套層數(shù)僅帶來了微不足道的成本。

CLIP-FID 權(quán)衡最后，在上圖 8c 中展示了 COCO 的zero-shot 評估中 CLIP-FID 的帕累托曲線，這是通過改變無分類器引導(dǎo)（CFG）權(quán)重實(shí)現(xiàn)的。MDM 對 CFG 的響應(yīng)與其他擴(kuò)散模型變體類似。作為比較，疊加了 Imagen 報(bào)告的相同圖。Imagen 通常展示出更小的 FID，這歸因于其在大數(shù)據(jù)集上訓(xùn)練所導(dǎo)致的更高多樣性。然而，MDM 展示出強(qiáng)大的 CLIP 得分，而在實(shí)踐中發(fā)現(xiàn)，這種高 CLIP 得分與生成圖像的視覺質(zhì)量有很好的相關(guān)性。

討論與未來方向

本文展示了跨不同分辨率共享表示可以加快訓(xùn)練速度并獲得高質(zhì)量結(jié)果，尤其是當(dāng)?shù)头直媛适紫缺挥?xùn)練時(shí)。相信這是因?yàn)槟Ｐ湍軌蚋行У乩貌煌直媛手g的相關(guān)性，既在空間上也在時(shí)間上。盡管在這里僅探討了一小部分架構(gòu)，預(yù)計(jì)通過對權(quán)重共享架構(gòu)的更詳細(xì)探索，以及在當(dāng)前架構(gòu)中不同分辨率間參數(shù)分配的新方法，可以實(shí)現(xiàn)更多改進(jìn)。

本文工作的另一個(gè)獨(dú)特方面是使用了擴(kuò)展空間，其中在多個(gè)分辨率上同時(shí)進(jìn)行去噪。在這種形式下，時(shí)間和空間上的分辨率以相同的方式處理，時(shí)間和空間中的相關(guān)結(jié)構(gòu)差異由權(quán)重共享模型的不同參數(shù)學(xué)習(xí)。對多分辨率的聯(lián)合優(yōu)化進(jìn)行更一般的概念化的方法是將不同分辨率的損失解耦，并給予它們不同的權(quán)重?？梢栽O(shè)想，在訓(xùn)練低分辨率到高分辨率的過程中可以實(shí)現(xiàn)平滑的過渡。

還注意到，盡管本文將方法與 LDM 進(jìn)行了比較，這些方法是互補(bǔ)的。可以在自編碼器代碼的基礎(chǔ)上構(gòu)建 MDM。盡管并沒有聲稱基于 MDM 的模型達(dá)到了最先進(jìn)的水平，但將對 MDM 在大規(guī)模數(shù)據(jù)集和模型規(guī)模上的評估留作未來工作。

本文轉(zhuǎn)自 AI生成未來，作者：AI生成未來

原文鏈接:??https://mp.weixin.qq.com/s/BLkteFzrZJq1-tCf678DOA??