在生成式 AI 時(shí)代，擴(kuò)散模型已經(jīng)成為圖像、視頻、3D、音頻和文本生成等生成式 AI 應(yīng)用的流行工具。然而將擴(kuò)散模型拓展到高分辨率領(lǐng)域仍然面臨巨大挑戰(zhàn)，這是因?yàn)槟Ｐ捅仨氃诿總€(gè)步驟重新編碼所有的高分辨率輸入。解決這些挑戰(zhàn)需要使用帶有注意力塊的深層架構(gòu)，這使得優(yōu)化更困難，消耗的算力和內(nèi)存也更多。

怎么辦呢？最近的一些工作專注于研究用于高分辨率圖像的高效網(wǎng)絡(luò)架構(gòu)。但是現(xiàn)有方法都沒(méi)有展示出超過(guò) 512×512 分辨率的效果，并且生成質(zhì)量落后于主流的級(jí)聯(lián)或 latent 方法。

我們以 OpenAI DALL-E 2、谷歌 IMAGEN 和英偉達(dá) eDiffI 為例，它們通過(guò)學(xué)習(xí)一個(gè)低分辨率模型和多個(gè)超分辨率擴(kuò)散模型來(lái)節(jié)省算力，其中每個(gè)組件都單獨(dú)訓(xùn)練。另一方面，latent 擴(kuò)散模型（LDM）僅學(xué)習(xí)低分辨率擴(kuò)散模型，并依賴單獨(dú)訓(xùn)練的高分辨率自編碼器。對(duì)于這兩種方案，多階段式 pipeline 使訓(xùn)練與推理復(fù)雜化，從而往往需要精心調(diào)整或進(jìn)行超參。

本文中，研究者提出了俄羅斯套娃式擴(kuò)散模型（Matryoshka Diffusion Models，MDM）它是用于端到端高分辨率圖像生成的全新擴(kuò)散模型。代碼很快將釋出。

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

該研究提出的主要觀點(diǎn)是將低分辨率擴(kuò)散過(guò)程作為高分辨率生成的一部分，通過(guò)使用嵌套 UNet 架構(gòu)在多個(gè)分辨率上執(zhí)行聯(lián)合擴(kuò)散過(guò)程。

該研究發(fā)現(xiàn)：MDM 與嵌套 UNet 架構(gòu)一起實(shí)現(xiàn)了 1）多分辨率損失：大大提高了高分辨率輸入去噪的收斂速度；2）高效的漸進(jìn)式訓(xùn)練計(jì)劃，從訓(xùn)練低分辨率擴(kuò)散模型開(kāi)始，按照計(jì)劃逐步添加高分辨率輸入和輸出。實(shí)驗(yàn)結(jié)果表明，多分辨率損失與漸進(jìn)式訓(xùn)練相結(jié)合可以讓訓(xùn)練成本和模型質(zhì)量獲得更好的平衡。

該研究在類條件圖像生成以及文本條件圖像和視頻生成方面評(píng)估了 MDM。MDM 讓訓(xùn)練高分辨率模型無(wú)需使用級(jí)聯(lián)或潛在擴(kuò)散（latent diffusion）。消融研究表明，多分辨率損失和漸進(jìn)訓(xùn)練都極大地提高了訓(xùn)練效率和質(zhì)量。

我們來(lái)欣賞以下 MDM 生成的圖片和視頻。

方法概覽

研究者介紹稱，MDM 擴(kuò)散模型在高分辨率中進(jìn)行端到端訓(xùn)練，同時(shí)利用層級(jí)結(jié)構(gòu)的數(shù)據(jù)形成。MDM 首先在擴(kuò)散空間中泛化了標(biāo)準(zhǔn)擴(kuò)散模型，然后提出了專用的嵌套架構(gòu)和訓(xùn)練流程。

首先來(lái)看如何在擴(kuò)展空間對(duì)標(biāo)準(zhǔn)擴(kuò)散模型進(jìn)行泛化。

與級(jí)聯(lián)或 latent 方法的不同之處在于，MDM 通過(guò)在一個(gè)擴(kuò)展空間中引入多分辨率擴(kuò)散過(guò)程，學(xué)得了具有層級(jí)結(jié)構(gòu)的單個(gè)擴(kuò)散過(guò)程。具體如下圖 2 所示。

具體來(lái)講，給定一個(gè)數(shù)據(jù)點(diǎn) x ∈ R^N，研究者定義了與時(shí)間相關(guān)的隱變量 z_t =  z_t^1 , . . . , z_t^R  ∈ R^N_1+...NR。

研究者表示，在擴(kuò)展空間中進(jìn)行擴(kuò)散建模有以下兩點(diǎn)優(yōu)點(diǎn)。其一，我們?cè)谕评砥陂g通常關(guān)心全分辨率輸出 z_t^R，那么所有其他中等分辨率被看作是額外的隱變量 z_t^r，增加了建模分布的復(fù)雜度。其二，多分辨率依賴性為跨 z_t^r 共享權(quán)重和計(jì)算提供了機(jī)會(huì)，從而以更高效的方式重新分配計(jì)算，并實(shí)現(xiàn)高效訓(xùn)練和推理。

接下來(lái)看嵌套架構(gòu)（NestedUNet）如何工作。

與典型的擴(kuò)散模型類似，研究者使用 UNet 網(wǎng)絡(luò)結(jié)構(gòu)來(lái)實(shí)現(xiàn) MDM，其中并行使用殘差連接和計(jì)算塊以保留細(xì)粒度的輸入信息。這里的計(jì)算塊包含多層卷積和自注意力層。NestedUNet 與標(biāo)準(zhǔn) UNet 的代碼分別如下。

除了相較于其他層級(jí)方法的簡(jiǎn)單性，NestedUNet 允許以最高效的方式對(duì)計(jì)算進(jìn)行分配。如下圖 3 所示，研究者早期探索發(fā)現(xiàn)，當(dāng)以最低分辨率分配大部分參數(shù)和計(jì)算時(shí)，MDM 實(shí)現(xiàn)了明顯更好的擴(kuò)展性。

最后是學(xué)習(xí)。

研究者使用常規(guī)去噪目標(biāo)在多個(gè)分辨率下訓(xùn)練 MDM，如下公式 (3) 所示。

這里用到了漸進(jìn)式訓(xùn)練。研究者按照上述公式 (3) 直接對(duì) MDM 進(jìn)行端到端訓(xùn)練，并展示出了比原始基線方法更好的收斂性。他們發(fā)現(xiàn)，使用類似于 GAN 論文中提出的簡(jiǎn)單漸進(jìn)式訓(xùn)練方法，極大地加速了高分辨率模型的訓(xùn)練。

這一訓(xùn)練方法從一開(kāi)始就避免了高成本的高分辨率訓(xùn)練，加速了整體收斂。不僅如此，他們還合并了混合分辨率訓(xùn)練，該訓(xùn)練方法在單個(gè) batch 中同時(shí)訓(xùn)練具有不同最終分辨率的樣本。

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

MDM 是一種通用技術(shù)，適用于可以逐步壓縮輸入維度的任何問(wèn)題。MDM 與基線方法的比較如下圖 4 所示。

表 1 給出了在 ImageNet（FID-50K）和 COCO（FID-30K）上的比較結(jié)果。

下圖 5、6、7 展示了 MDM 在圖像生成（圖 5）、文本到圖像（圖 6）和文本到視頻（圖 7）方面的結(jié)果。盡管是在相對(duì)較小的數(shù)據(jù)集上進(jìn)行訓(xùn)練的，但 MDM 仍顯示出生成高分辨率圖像和視頻的強(qiáng)大零樣本（zero-shot）能力。

感興趣的讀者可以閱讀論文原文，了解更多研究?jī)?nèi)容。