在生成式 AI 時代，擴散模型已經成為圖像、視頻、3D、音頻和文本生成等生成式 AI 應用的流行工具。然而將擴散模型拓展到高分辨率領域仍然面臨巨大挑戰，這是因為模型必須在每個步驟重新編碼所有的高分辨率輸入。解決這些挑戰需要使用帶有注意力塊的深層架構，這使得優化更困難，消耗的算力和內存也更多。

怎么辦呢？最近的一些工作專注于研究用于高分辨率圖像的高效網絡架構。但是現有方法都沒有展示出超過 512×512 分辨率的效果，并且生成質量落后于主流的級聯或 latent 方法。

我們以 OpenAI DALL-E 2、谷歌 IMAGEN 和英偉達 eDiffI 為例，它們通過學習一個低分辨率模型和多個超分辨率擴散模型來節省算力，其中每個組件都單獨訓練。另一方面，latent 擴散模型（LDM）僅學習低分辨率擴散模型，并依賴單獨訓練的高分辨率自編碼器。對于這兩種方案，多階段式 pipeline 使訓練與推理復雜化，從而往往需要精心調整或進行超參。

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

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

該研究提出的主要觀點是將低分辨率擴散過程作為高分辨率生成的一部分，通過使用嵌套 UNet 架構在多個分辨率上執行聯合擴散過程。

該研究發現：MDM 與嵌套 UNet 架構一起實現了 1）多分辨率損失：大大提高了高分辨率輸入去噪的收斂速度；2）高效的漸進式訓練計劃，從訓練低分辨率擴散模型開始，按照計劃逐步添加高分辨率輸入和輸出。實驗結果表明，多分辨率損失與漸進式訓練相結合可以讓訓練成本和模型質量獲得更好的平衡。

該研究在類條件圖像生成以及文本條件圖像和視頻生成方面評估了 MDM。MDM 讓訓練高分辨率模型無需使用級聯或潛在擴散（latent diffusion）。消融研究表明，多分辨率損失和漸進訓練都極大地提高了訓練效率和質量。

我們來欣賞以下 MDM 生成的圖片和視頻。

方法概覽

研究者介紹稱，MDM 擴散模型在高分辨率中進行端到端訓練，同時利用層級結構的數據形成。MDM 首先在擴散空間中泛化了標準擴散模型，然后提出了專用的嵌套架構和訓練流程。

首先來看如何在擴展空間對標準擴散模型進行泛化。

與級聯或 latent 方法的不同之處在于，MDM 通過在一個擴展空間中引入多分辨率擴散過程，學得了具有層級結構的單個擴散過程。具體如下圖 2 所示。

具體來講，給定一個數據點 x ∈ R^N，研究者定義了與時間相關的隱變量 z_t =  z_t^1 , . . . , z_t^R  ∈ R^N_1+...NR。

研究者表示，在擴展空間中進行擴散建模有以下兩點優點。其一，我們在推理期間通常關心全分辨率輸出 z_t^R，那么所有其他中等分辨率被看作是額外的隱變量 z_t^r，增加了建模分布的復雜度。其二，多分辨率依賴性為跨 z_t^r 共享權重和計算提供了機會，從而以更高效的方式重新分配計算，并實現高效訓練和推理。

接下來看嵌套架構（NestedUNet）如何工作。

與典型的擴散模型類似，研究者使用 UNet 網絡結構來實現 MDM，其中并行使用殘差連接和計算塊以保留細粒度的輸入信息。這里的計算塊包含多層卷積和自注意力層。NestedUNet 與標準 UNet 的代碼分別如下。

除了相較于其他層級方法的簡單性，NestedUNet 允許以最高效的方式對計算進行分配。如下圖 3 所示，研究者早期探索發現，當以最低分辨率分配大部分參數和計算時，MDM 實現了明顯更好的擴展性。

最后是學習。

研究者使用常規去噪目標在多個分辨率下訓練 MDM，如下公式 (3) 所示。

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

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

實驗及結果

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

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

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

感興趣的讀者可以閱讀論文原文，了解更多研究內容。